main.py

import os
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.metrics import confusion_matrix, f1_score
import argparse
import json
import numpy as np


'''Helper function for saving confusion matrices in a standardized way.'''
def save_confusion_matrix(cm, labels, xlabel, ylabel, title, output_path):
    cm_percent = cm
    annot = np.array([[f"{val:.2f}%" for val in row] for row in cm_percent]).reshape((3,3))

    fig, ax = plt.subplots()
    sns.heatmap(
        cm, annot=annot, fmt='', cmap='Blues', cbar=False,
        xticklabels=labels, yticklabels=labels, ax=ax
    )
    ax.set_xlabel(xlabel.title().replace('4O', '4o'))
    ax.set_ylabel(ylabel.title().replace('4O', '4o'))
    ax.set_title(title)
    plt.savefig(output_path, bbox_inches='tight', dpi=300)
    plt.close(fig)


'''Computes the F1 score from a full confusion matrix (neither row nor column normalized.)'''
def f1_from_confusion_matrix(cm):

    cm_arr = np.array(cm)
    n = cm_arr.shape[0]
    f1_scores = []
    for i in range(n):
        tp = cm_arr[i, i]
        fn = cm_arr[i, :].sum() - tp
        fp = cm_arr[:, i].sum() - tp
        precision = tp / (tp + fp) if (tp + fp) > 0 else 0.0
        recall = tp / (tp + fn) if (tp + fn) > 0 else 0.0

        if precision + recall > 0:
            f1 = 2 * precision * recall / (precision + recall)
        else:
            f1 = 0.0
        f1_scores.append(f1)
    return float(np.mean(f1_scores))


'''Saves confusion matrices as images from JSON'''
def generateJsonConfusionMatrices():
    input_file = 'data/matrices/matrices.json'
    output_dir = 'ground_truth_matrices'
    os.makedirs(output_dir, exist_ok=True)
    with open(input_file) as f:
        matrices = json.load(f)
    labels = ['context', 'uses', 'extends']
    for name, matrix in matrices.items():
        cm = matrix
        total = sum(sum(row) for row in cm)
        true = []
        pred = []
        for i, row in enumerate(cm):
            for j, count in enumerate(row):
                true.extend([labels[i]] * count)
                pred.extend([labels[j]] * count)
        f1 = f1_score(true, pred, average='macro')
        xlabel = name.replace('_', ' ')
        ylabel = 'Ground Truth'
        title = f"Total Samples = {total}, F1-Score = {f1:.2f}"
        output_path = os.path.join(output_dir, f"{name}.png")
        save_confusion_matrix(cm, labels, xlabel, ylabel, title, output_path)
        print(f"Saved JSON confusion matrix for {name} to {output_path}")

'''Generates and saves confusion matrices from labeled datasets.'''
def generateLabeledConfusionMatrices():

    input_dir = 'data/labels'
    output_dir = 'ground_truth_matrices'
    os.makedirs(output_dir, exist_ok=True)

    weights_map = generateCorpusDistributions()
    filename_to_classification = {'gpt-4.1-mini-labeled.csv': 'gpt-4.1-mini_classification',
                                  'llama-gpt-4o-labeled.csv': 'llama_gpt-4o-mini_classification',
                                  'llama-labeled.csv': 'llama_classification'}


    for filename in os.listdir(input_dir):
        if filename.endswith('.csv'):
            weights = weights_map.get(filename_to_classification.get(filename))
            file_path = os.path.join(input_dir, filename)
            df = pd.read_csv(file_path)

            true = df['manual_classification']
            predicted_cols = [col for col in df.columns if col not in ['multisentence', 'manual_classification']]
            if not predicted_cols:
                continue
            pred_col = predicted_cols[0]
            pred = df[pred_col]

            labels = ['context', 'uses', 'extends']
            cm = confusion_matrix(true, pred, labels=labels)
            if weights is not None:
                print(f"Using weights when calcualting matrix for {filename}")
                weights_arr = np.array(weights)
                if weights_arr.shape[0] != len(labels):
                    raise ValueError(f"Length of weights ({weights_arr.shape[0]}) does not match number of labels ({len(labels)})")
                cm = cm/cm.sum(axis = 0)
                cm = cm * weights_arr
                cm = (cm/cm.sum() * 100)  

            total = len(df.dropna(subset=['manual_classification', pred_col]))

            f1 = f1_from_confusion_matrix(cm)

            xlabel = pred_col.replace('_', ' ')
            ylabel = 'Ground Truth'
            title = f"N = {total}, F1-Score = {f1:.2f}"
            output_path = os.path.join(output_dir, f"{os.path.splitext(filename)[0]}.png")

            save_confusion_matrix(cm, labels, xlabel, ylabel, title, output_path)
            print(f"Saved labeled confusion matrix for {filename} to {output_path}")

'''Generates and saves confusion matrices from sentence corpus between each method pair.'''
def generateInterConfusionMatrices():
    df = pd.read_csv('data/corpus/sentence_classification_corpus.csv')

    os.makedirs('interconfusion_matrices', exist_ok=True)

    cols = [
        'gpt-4.1-mini_classification',
        'llama_classification',
        'llama_gpt-4o-mini_classification'
    ]

    for i in range(len(cols)):
        for j in range(i + 1, len(cols)):
            col2 = cols[i]
            col1 = cols[j]

            df_cm = df[[col1, col2]].dropna()

            y_true = df_cm[col1]
            y_pred = df_cm[col2]

            labels = ['context', 'uses', 'extends']
            cm = confusion_matrix(y_true, y_pred, labels=labels)
            cm = cm/cm.sum() * 100

            total_samples = len(df_cm)
            f1 = f1_from_confusion_matrix(cm)
            xlabel = col2.replace('_', ' ')
            ylabel = col1.replace('_', ' ').title()
            title = f"N = {total_samples//1000}k, F1-Score = {f1:.2f}"
            filename = f"{col1.replace('_classification', '')}_vs_{col2.replace('_classification', '')}.png"
            output_path = os.path.join('interconfusion_matrices', filename)

            save_confusion_matrix(cm, labels, xlabel, ylabel, title, output_path)
            print(f"Saved confusion matrix for {col1} vs {col2} to {output_path}")

'''Generates the predicted distributions for each classifier, or uses a cached value for efficiency.'''
def generateCorpusDistributions(use_cache=False):
    json_path = 'data/corpus/corpus_distributions.json'
    csv_path = 'data/corpus/sentence_classification_corpus.csv'

    if use_cache and os.path.exists(json_path):
        with open(json_path) as f:
            distributions = json.load(f)
        return distributions


    df = pd.read_csv(csv_path)
    cols = [
        'gpt-4.1-mini_classification',
        'llama_classification',
        'llama_gpt-4o-mini_classification'
    ]
    labels = ['context', 'uses', 'extends']
    distributions = {}
    for col in cols:
        counts = df[col].value_counts().reindex(labels, fill_value=0)
        total = counts.sum()
        fractions = (counts / total).tolist()
        distributions[col] = fractions


    os.makedirs(os.path.dirname(json_path), exist_ok=True)
    with open(json_path, 'w') as f:
        json.dump(distributions, f, indent=4)

    return distributions

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Generate corpus data or confusion matrices')
    parser.add_argument('--interconfusion-matrices', dest='interconfusion_matrices', action='store_true',
                        help='Generate confusion matrices')
    parser.add_argument('--ground-truth-matrices', dest='ground_truth_matrices', action='store_true',
                        help='Generate labeled confusion matrices')
    parser.add_argument('--all', dest='all', action='store_true',
                        help='Generate all confusion matrices.')

    args = parser.parse_args()


    if args.interconfusion_matrices or args.all:
        generateInterConfusionMatrices()
    if args.ground_truth_matrices or args.all:
        generateLabeledConfusionMatrices()
        generateJsonConfusionMatrices()

    if not (args.interconfusion_matrices or args.ground_truth_matrices or args.all):
        parser.print_help()