#!/usr/bin/env python3
"""
Script to add split information to combined_annotations_trajectories.csv
"""

import os

import pandas as pd


def add_splits_to_trajectories():
    # Define file paths
    base_dir = "/home/mashalimay/webarena/modular_agent_merge/utils_agrb"
    splits_file = os.path.join(base_dir, "splits.csv")
    trajectories_file = os.path.join(base_dir, "combined_annotations_trajectories.csv")
    output_file = os.path.join(base_dir, "combined_annotations_trajectories_with_splits.csv")

    print("Reading splits.csv...")
    # Read the splits file
    splits_df = pd.read_csv(splits_file)
    print(f"Loaded {len(splits_df)} rows from splits.csv")

    print("Reading combined_annotations_trajectories.csv...")
    # Read the trajectories file
    trajectories_df = pd.read_csv(trajectories_file)
    print(f"Loaded {len(trajectories_df)} rows from combined_annotations_trajectories.csv")

    print("Merging data...")
    # Create a normalized task_id for merging (remove .resized from visualwebarena tasks)
    trajectories_df["normalized_task_id"] = trajectories_df["task_id"].str.replace(".resized", "", regex=False)

    # Merge the dataframes on normalized_task_id and benchmark
    # Left join to keep all trajectories, even if no split info is found
    merged_df = trajectories_df.merge(
        splits_df[["task_id", "benchmark", "split"]],
        left_on=["normalized_task_id", "benchmark"],
        right_on=["task_id", "benchmark"],
        how="left",
        suffixes=("", "_splits"),
    )

    # Drop the helper columns and splits task_id (keep original task_id from trajectories)
    merged_df = merged_df.drop(["normalized_task_id", "task_id_splits"], axis=1)

    print(f"Merged dataframe has {len(merged_df)} rows")

    # Check for any missing split information
    missing_splits = merged_df["split"].isna().sum()
    if missing_splits > 0:
        print(f"Warning: {missing_splits} rows have missing split information")
        print("Tasks with missing splits:")
        missing_tasks = merged_df[merged_df["split"].isna()][["task_id", "benchmark"]].drop_duplicates()
        for _, row in missing_tasks.iterrows():
            print(f"  - {row['task_id']} ({row['benchmark']})")

    # Show split distribution
    print("\nSplit distribution:")
    split_counts = merged_df["split"].value_counts(dropna=False)
    for split, count in split_counts.items():
        print(f"  {split}: {count}")

    print(f"Saving merged data to {output_file}...")
    # Save the merged dataframe
    merged_df.to_csv(output_file, index=False)

    print("Done! The file with split information has been saved.")
    return output_file


if __name__ == "__main__":
    output_file = add_splits_to_trajectories()
    print(f"\nOutput saved to: {output_file}")
