import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
import argparse
import sys

def plot_servo_data(csv_file):
    # Read CSV file
    try:
        df = pd.read_csv(csv_file)
    except FileNotFoundError:
        print(f"Error: File '{csv_file}' not found")
        sys.exit(1)
    except Exception as e:
        print(f"Error reading CSV file: {e}")
        sys.exit(1)

    # Preprocess data: subtract 500 and multiply by 0.5
    for servo in ['Servo 1', 'Servo 2', 'Servo 3']:
        df[servo] = (df[servo] - 500) * 0.25

    # Set plot style
    sns.set_style("whitegrid")
    plt.figure(figsize=(12, 6))

    # Plot servo angles
    sns.lineplot(data=df, x='Step', y='Servo 1', label='Servo 1')
    sns.lineplot(data=df, x='Step', y='Servo 2', label='Servo 2')
    sns.lineplot(data=df, x='Step', y='Servo 3', label='Servo 3')

    # Set title and labels
    plt.title('Servo Angles vs Time (Normalized)')
    plt.xlabel('Time (seconds)')
    plt.ylabel('Angle (degrees)')

    # Add legend
    plt.legend(title='Servo ID')

    # Optimize display
    plt.tight_layout()

    # Generate output filename based on input
    output_file = csv_file.rsplit('.', 1)[0] + '_normalized_plot.png'
    
    # Save plot
    plt.savefig(output_file, dpi=300, bbox_inches='tight')
    print(f"Plot saved as: {output_file}")

def main():
    # Set up argument parser
    parser = argparse.ArgumentParser(description='Plot servo angles from CSV file')
    parser.add_argument('csv_file', help='Path to the CSV file containing servo data')
    
    # Parse arguments
    args = parser.parse_args()
    
    # Generate plot
    plot_servo_data(args.csv_file)

if __name__ == "__main__":
    main()