import os

import matplotlib
import pandas as pd

matplotlib.use("Agg")
import matplotlib.pyplot as plt

# Constants
PROJECT_ROOT = os.path.dirname(os.path.dirname(__file__))
RESULTS_DIR = os.path.join(PROJECT_ROOT, "results")
TABLE_DATA_FILE = os.path.join(RESULTS_DIR, "table_data.csv")

# Ensure the results directory exists
os.makedirs(RESULTS_DIR, exist_ok=True)


def load_table_data(file_path):
    """Loads table data from a CSV file into a pandas DataFrame."""
    return pd.read_csv(file_path)


def format_number(x):
    """Formats a number to display in thousands (K) or millions (M)."""
    if x >= 1e9:
        return f"{x / 1e9:.2f}B"
    elif x >= 1e6:
        return f"{x / 1e6:.2f}M"
    elif x >= 1e3:
        return f"{x / 1e3:.2f}K"
    else:
        return f"{x:.0f}"


def format_large_numbers(data):
    """Applies number formatting to all numeric columns except 'Input Size'."""
    for column in data.columns[1:]:  # Skip 'Input Size'
        data[column] = data[column].apply(format_number)
    return data


def generate_table_markdown(data, output_path):
    """Generates a markdown file containing the formatted data table."""
    with open(output_path, "w") as file:
        file.write("# Sorting Algorithm Comparison Tables\n\n")
        file.write(data.to_markdown(index=False))
    print(f"Markdown table saved to {output_path}")


def generate_table_image(data, image_output_path):
    """Generates and saves an image of the data table."""
    # Create figure and axis
    fig, ax = plt.subplots(figsize=(18, 9))
    ax.axis("tight")
    ax.axis("off")

    # Set font locally for this plot
    plt.rcParams["font.family"] = "Arial"

    # Adjust column labels: wrap words and handle "Non-decreasing"/"Non-increasing" separately
    def wrap_label(label):
        # Replace specific patterns with wrapped versions
        label = (
            label.replace("(Random) ", "(Random)\n")
            .replace("(Deterministic) ", "(Deterministic)\n")
            .replace("Non-decreasing", "Non-\ndecreasing")
            .replace("Non-increasing", "Non-\nincreasing")
        )

        # Handle special case when "Sort" is in the label
        if "Sort" in label:
            algorithm, *rest = label.split("Sort", 1)
            return f"{algorithm}Sort\n{rest[0].strip()}" if rest else f"{algorithm}Sort"

        # General word wrapping
        return "\n".join(label.split())

    # Apply wrapping to all column labels
    col_labels = [wrap_label(label) for label in data.columns]

    # Generate the table
    table = ax.table(
        cellText=data.values,
        colLabels=col_labels,
        cellLoc="center",
        loc="center",
    )
    table.set_fontsize(20)

    num_rows, num_cols = data.shape

    # Customize cell properties
    for (row, col), cell in table.get_celld().items():
        # Set cell border color
        cell.set_edgecolor("#CCD3DB")

        if row == 0:  # Header row formatting
            cell.set_text_props(
                weight="bold", ha="right" if col == 0 else "left", linespacing=1.5
            )
            cell.set_facecolor("#ffffff")
            cell.set_height(0.20)

        else:  # Data row formatting
            if col == 0:  # Make 'Input Size' column bold
                cell.set_text_props(weight="bold", ha="right")
            else:
                cell.set_text_props(ha="left")
            # Alternate row colors
            cell.set_facecolor("#ffffff" if (row - 1) % 2 == 0 else "#f6f8fa")
            cell.set_height(0.06)

    # Adjust column widths
    for col_index in range(num_cols):
        table.auto_set_column_width(col=[col_index])

    # Add a title to the figure
    title_text = "Number of Comparisons Made by Sorting Algorithms for Various Input Sizes on Sorted Datasets"
    fig.text(0.5, 0.88, title_text, ha="center", fontsize=20, fontweight="bold")

    # Add a legend for "K" and "M" notation at the bottom
    legend_text = "* K: thousands\n* M: millions\n* B: billions"
    fig.text(0.05, 0.05, legend_text, ha="left", fontsize=14, fontweight="bold")

    # Apply tight layout to minimize white space
    plt.tight_layout()

    # Save the image with a tight bounding box
    plt.savefig(image_output_path, bbox_inches="tight", pad_inches=0.1, dpi=150)
    plt.close()
    print(f"Table image saved to {image_output_path}")


def main():
    # Load and format data
    table_data = load_table_data(TABLE_DATA_FILE)
    formatted_table_data = format_large_numbers(table_data.copy())

    # Output file paths
    markdown_output_file = os.path.join(RESULTS_DIR, "comparison_table.md")
    image_output_file = os.path.join(RESULTS_DIR, "comparison_table.png")

    # Generate markdown and image outputs
    generate_table_markdown(formatted_table_data, markdown_output_file)
    generate_table_image(formatted_table_data, image_output_file)


if __name__ == "__main__":
    main()