import os

import matplotlib.pyplot as plt
import pandas as pd


def plot_coordinates(df, output_folder):
    os.makedirs(output_folder, exist_ok=True)

    for _, row in df.iterrows():
        fig, ax = plt.subplots(figsize=(10, 10))

        # Plot points
        ax.scatter(
            row["left_shoulder_x"],
            row["left_shoulder_y"],
            color="red",
            label="Left Shoulder",
        )
        ax.scatter(
            row["right_shoulder_x"],
            row["right_shoulder_y"],
            color="blue",
            label="Right Shoulder",
        )
        ax.scatter(
            row["left_hip_x"], row["left_hip_y"], color="green", label="Left Hip"
        )
        ax.scatter(
            row["right_hip_x"], row["right_hip_y"], color="purple", label="Right Hip"
        )
        ax.scatter(
            row["stethoscope_x"],
            row["stethoscope_y"],
            color="orange",
            label="Stethoscope",
        )

        # Connect shoulders and hips
        ax.plot(
            [row["left_shoulder_x"], row["right_shoulder_x"]],
            [row["left_shoulder_y"], row["right_shoulder_y"]],
            "k-",
        )
        ax.plot(
            [row["left_hip_x"], row["right_hip_x"]],
            [row["left_hip_y"], row["right_hip_y"]],
            "k-",
        )
        ax.plot(
            [row["left_shoulder_x"], row["left_hip_x"]],
            [row["left_shoulder_y"], row["left_hip_y"]],
            "k-",
        )
        ax.plot(
            [row["right_shoulder_x"], row["right_hip_x"]],
            [row["right_shoulder_y"], row["right_hip_y"]],
            "k-",
        )
        ax.set_aspect("equal")
        # Invert y-axis to match image coordinates (origin at top-left)
        ax.invert_yaxis()
        ax.legend()
        ax.set_title(f"Coordinates for {row['image_file_name']}")

        plt.savefig(
            os.path.join(output_folder, f"{row['image_file_name'].split('.')[0]}.png")
        )
        plt.close()


# Read CSV files
result_df = pd.read_csv("output/results/results.csv")
result_convert_df = pd.read_csv("output/results/results-convert.csv")

# Plot and save results
plot_coordinates(result_df, "result_plot")
plot_coordinates(result_convert_df, "result-convert-plot")

""" result_df = pd.read_csv("data(卓球玉正規化座標)/01/results-convert.csv")
plot_coordinates(result_df, "data-confirm") """

print(
    "Plotting completed. Results saved in 'result_plot' and 'result-convert-plot' folders."
)