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

plt.rcParams.update({
    'font.size': 16,
    'font.family': 'serif',
    'axes.labelsize': 18,
    'axes.titlesize': 20,
    'xtick.labelsize': 14,
    'ytick.labelsize': 14,
    'legend.fontsize': 14,
    'figure.dpi': 300,
    'axes.linewidth': 1.5,
    'xtick.direction': 'in',
    'ytick.direction': 'in',
    'xtick.major.size': 6,
    'ytick.major.size': 6,
    'xtick.minor.size': 3,
    'ytick.minor.size': 3,
})


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 calculate_distance(row):
    origin = np.array([row['XR_Origin_X'], row['XR_Origin_Y'], row['XR_Origin_Z']])
    tracked = np.array([row['TrackedObject_X'], row['TrackedObject_Y'], row['TrackedObject_Z']])
    return np.linalg.norm(tracked - origin)

def process_position_data(data_dir):
    data_by_scene = {}
    all_distances = []
    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))
        df['Distance'] = df.apply(calculate_distance, axis=1)
        all_distances.extend(df['Distance'])
        if scene_name not in data_by_scene:
            data_by_scene[scene_name] = []
        data_by_scene[scene_name].extend(df['Distance'])
    return all_distances, data_by_scene

def plot_boxplots(all_distances, data_by_scene):
    plt.figure(figsize=(10, 7))
    data = [distances for _, distances in sorted(data_by_scene.items())]
    labels = [scene for scene in sorted(data_by_scene.keys())]
    data = [all_distances] + data
    labels = ['All Scenes'] + labels
    box = plt.boxplot(
        data, labels=labels, showmeans=True, meanline=True, showfliers=False,
        boxprops=dict(linewidth=2), whiskerprops=dict(linewidth=2), capprops=dict(linewidth=2),
        medianprops=dict(linewidth=2, color='black'), meanprops=dict(linewidth=2, color='C1')
    )
    plt.ylabel('Distance to Tracked Object (units)')
    plt.xlabel('Scene')
    plt.title('Distribution of Distances to Tracked Objects by Scene')
    plt.xticks(rotation=20)
    plt.grid(axis='y', linestyle='--', linewidth=1, alpha=0.7)
    plt.tight_layout()
    plt.savefig('boxplot_distances.png', dpi=600)
    plt.close()

def main():
    all_distances, data_by_scene = process_position_data('Recordings')
    plot_boxplots(all_distances, data_by_scene)

if __name__ == "__main__":
    main()