"""
Generate paper table images using matplotlib.
"""

import argparse
import os
import numpy as np
import matplotlib.pyplot as plt


def make_table1(output_dir):
    data = {
        "Config":   ["Q8_0","Q6_K","Q5_1","Q5_0","Q4_1","Q4_0","Q3_K_S","Q2_K"],
        "LPIPS":    ["—","0.032","0.055","0.069","0.101","0.117","0.196","0.304"],
        "CLIP":     ["0.361","0.361","0.360","0.361","0.361","0.360","0.356","0.358"],
        "Quality":  ["0.0306","0.0308","0.0305","0.0305","0.0300","0.0300","0.0271","0.0247"],
        "IR":       ["1.252","1.261","1.241","—","—","1.229","1.068","1.022"],
        "Cat A":    ["—","0.038","0.062","0.076","0.131","0.131","0.228","0.291"],
        "Cat B":    ["—","0.027","0.044","0.051","0.075","0.090","0.189","0.254"],
        "Cat C":    ["—","0.042","0.080","0.091","0.123","0.152","0.215","0.357"],
        "Cat D":    ["—","0.020","0.046","0.055","0.078","0.102","0.162","0.276"],
        "Cat E":    ["—","0.031","0.044","0.071","0.098","0.109","0.186","0.340"],
    }

    cols = list(data.keys())
    cell_text = list(zip(*[data[c] for c in cols]))

    fig, ax = plt.subplots(figsize=(14, 5))
    ax.axis("off")
    ax.set_title("Table 1: Quantization Ladder — Uniform Models vs Q8_0 Baseline",
                 fontsize=13, fontweight="bold", pad=20)

    table = ax.table(cellText=cell_text, colLabels=cols,
                     cellLoc="center", loc="center")
    table.auto_set_font_size(False)
    table.set_fontsize(10)
    table.scale(1.2, 1.6)

    for j in range(len(cols)):
        table[0, j].set_facecolor("#2c2c2c")
        table[0, j].set_text_props(color="white", fontweight="bold")
    for i in range(1, len(cell_text) + 1):
        color = "#f5f5f5" if i % 2 == 0 else "white"
        for j in range(len(cols)):
            table[i, j].set_facecolor(color)

    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, "table1_quant_ladder.png"),
                dpi=200, bbox_inches="tight")
    plt.close()
    print("Generated: table1_quant_ladder.png")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--output-dir", required=True)
    args = parser.parse_args()
    os.makedirs(args.output_dir, exist_ok=True)
    make_table1(args.output_dir)


if __name__ == "__main__":
    main()