import pandas as pd
import re

# Input and output files
input_csv = "results.csv"
output_csv = "submission.csv"

# Load frame-level results
df = pd.read_csv(input_csv)

# Extract base video ID (remove _f###.jpg)
def extract_video_id(filename):
    return re.sub(r"_f\d+\.\w+$", "", filename)

df["video_id"] = df["file_name"].apply(extract_video_id)

# Aggregate probabilities per video (mean of frame scores)
video_scores = (
    df.groupby("video_id")["predicted_prob"]
    .mean()
    .reset_index()
    .rename(columns={"video_id": "id", "predicted_prob": "score"})
)

# Assign label: generated if >0.5 else real
video_scores["pred"] = video_scores["score"].apply(
    lambda x: "generated" if x > 0.5 else "real"
)

# Reorder columns
video_scores = video_scores[["id", "pred", "score"]]

# Save submission file
video_scores.to_csv(output_csv, index=False)
print(f"Saved submission file to {output_csv}")