from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd


def plot_descriptors():
    input_path = "visualizations/raw_data/descriptors.csv"

    output_dir = "visualizations/plots/descriptors"
    Path(output_dir).mkdir(parents=True, exist_ok=True)

    df = pd.read_csv(input_path)
    properties = df["property"].unique()

    colors = plt.cm.viridis(np.linspace(0, 1, len(df["dataset"].unique())))

    for prop in properties:
        prop_df = df[df["property"] == prop]
        datasets = prop_df["dataset"].tolist()
        stats = [
            {
                "label": row["dataset"],
                "q1": row["q1"],
                "med": row["median"],
                "q3": row["q3"],
                "whislo": row["p5"],
                "whishi": row["p95"],
                "fliers": [],
            }
            for _, row in prop_df.iterrows()
        ]

        fig, ax = plt.subplots(figsize=(10, 6))

        # prepare bar plot positions
        # we want extra space for "SuperNatural3", which has index 5
        base_spacing = 0.4
        extra_gap = 0.1
        positions = []
        for i in range(len(datasets)):
            if i < 5:
                positions.append(i * base_spacing)
            elif i == 5:
                positions.append(i * base_spacing + extra_gap)
            else:
                positions.append(i * base_spacing + 2 * extra_gap)

        ax.bxp(
            stats,
            patch_artist=True,
            positions=positions,
            widths=0.3,
            boxprops=dict(edgecolor="black", linewidth=1.5),
            medianprops=dict(color="black", linewidth=2),
            whiskerprops=dict(color="black", linewidth=1.5),
            capprops=dict(color="black", linewidth=1.5),
        )

        # Set box colors
        for patch, color in zip(ax.patches, colors, strict=False):
            patch.set_facecolor(color)

        ax.set_xticks(positions)
        ax.set_xticklabels(datasets, rotation=30, ha="center", fontsize=11)
        ax.tick_params(axis="y", labelsize=11)
        ax.grid(True, axis="y", linestyle="--", alpha=0.7)

        # tighter before first and after last bar
        ax.set_xlim(min(positions) - 0.25, max(positions) + 0.25)

        plt.tight_layout()

        output_path = f"{output_dir}/{prop}_boxplot.png"
        plt.savefig(
            output_path,
            dpi=300,
            bbox_inches="tight",
            facecolor="white",
            edgecolor="none",
        )
        plt.close(fig)


if __name__ == "__main__":
    plot_descriptors()
