import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from pathlib import Path

def load_csv_with_standard_decimal(file_path):
    with open(file_path, 'r') as f:
        scene_line = f.readline().strip()
        scene_name = scene_line.split(': ')[1].split('_')[0]
    df = pd.read_csv(file_path, skiprows=1)
    return df, scene_name

def process_position_data(data_dir):
    # Group by scene
    scene_data = {}
    csv_files = list(Path(data_dir).glob('**/P*_PositionData*.csv'))
    for file_path in csv_files:
        df, scene_name = load_csv_with_standard_decimal(str(file_path))
        if scene_name not in scene_data:
            scene_data[scene_name] = []
        scene_data[scene_name].append(df)
    return scene_data

def plot_relative_positions(scene_data):
    scenes = list(scene_data.keys())
    # Create one figure with 4 subplots
    fig, axes = plt.subplots(2, 2, figsize=(16, 14))
    axes = axes.flatten()
    for idx, (scene, recordings) in enumerate(scene_data.items()):
        all_rel_x = []
        all_rel_z = []
        avg_obj_xs = []
        avg_obj_zs = []
        for df in recordings:
            obj_x = df['TrackedObject_X'].values
            obj_z = df['TrackedObject_Z'].values
            avg_obj_x = np.mean(obj_x)
            avg_obj_z = np.mean(obj_z)
            avg_obj_xs.append(avg_obj_x)
            avg_obj_zs.append(avg_obj_z)
            rel_x = df['XR_Origin_X'].values - avg_obj_x
            rel_z = df['XR_Origin_Z'].values - avg_obj_z
            all_rel_x.extend(rel_x)
            all_rel_z.extend(rel_z)
        ax = axes[idx]
        ax.scatter(all_rel_x, all_rel_z, alpha=0.2, c='blue', label='Player Positions (relative)')
        ax.scatter(0, 0, c='red', s=800, marker='*', label='Agent', edgecolor='black', linewidths=2, zorder=10)
        ax.set_title(f'Scene: {scene}')
        ax.set_xlabel('X Position (relative to Agent)')
        ax.set_ylabel('Z Position (relative to Agent)')
        ax.legend()
        ax.axis('equal')
        ax.grid(True, linestyle='--', alpha=0.7)
        # Also save individual images
        plt.figure(figsize=(8, 8))
        plt.scatter(all_rel_x, all_rel_z, alpha=0.2, c='blue', label='Player Positions (relative)')
        plt.scatter(0, 0, c='red', s=800, marker='*', label='Agent', edgecolor='black', linewidths=2, zorder=10)
        plt.title(f'Player Positions Relative to Agent\nScene: {scene}')
        plt.xlabel('X Position (relative to Agent)')
        plt.ylabel('Z Position (relative to Agent)')
        plt.legend()
        plt.axis('equal')
        plt.grid(True, linestyle='--', alpha=0.7)
        plt.tight_layout()
        plt.savefig(f'room_layout_{scene}.png', dpi=300)
        plt.close()
    plt.tight_layout()
    plt.savefig('room_layout_all_scenes.png', dpi=300)
    plt.close()

def main():
    scene_data = process_position_data('Recordings')
    plot_relative_positions(scene_data)

if __name__ == "__main__":
    main()