"""
Assemble panorama strips into a single comparison grid image.
"""

import argparse
import os
import glob
from PIL import Image, ImageDraw, ImageFont
from config import MODEL_ORDER


LABEL_W = 100
STEP_LABEL_H = 40
PAD = 2


def find_panorama(base_dir, tag, seed):
    prefix_map = {"Q2_K": "q2k", "Q3_K_S": "q3ks", "Q4_0": "q4_0",
                  "Q4_1": "q4_1", "Q5_0": "q5_0", "Q5_1": "q5_1",
                  "Q6_K": "q6k", "Q8_0": "q8_0"}
    prefix = prefix_map.get(tag, tag.lower())
    pattern = os.path.join(base_dir, f"panorama_{prefix}_s{seed}_*.png")
    files = sorted(glob.glob(pattern))
    return files[-1] if files else None


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--panorama-dir", required=True)
    parser.add_argument("--output-dir", required=True)
    parser.add_argument("--seed", type=int, default=42)
    args = parser.parse_args()

    os.makedirs(args.output_dir, exist_ok=True)

    rows = []
    for tag in MODEL_ORDER:
        path = find_panorama(args.panorama_dir, tag, args.seed)
        if path:
            rows.append((tag, Image.open(path)))
            print(f"  [OK] {tag}")
        else:
            print(f"  [--] {tag}: not found")

    if not rows:
        print("No panoramas found.")
        return

    pw, ph = rows[0][1].size
    n_steps = 20
    step_w = pw // n_steps

    grid_w = LABEL_W + pw + PAD
    grid_h = STEP_LABEL_H + (ph + PAD) * len(rows)
    grid = Image.new("RGB", (grid_w, grid_h), (20, 20, 20))
    draw = ImageDraw.Draw(grid)

    try:
        font = ImageFont.truetype("arial.ttf", 20)
        font_sm = ImageFont.truetype("arial.ttf", 14)
    except:
        font = ImageFont.load_default()
        font_sm = font

    for j in range(n_steps):
        x = LABEL_W + j * step_w + step_w // 2 - 8
        draw.text((x, 10), str(j+1), fill="white", font=font_sm)

    for i, (tag, pano) in enumerate(rows):
        y = STEP_LABEL_H + (ph + PAD) * i
        draw.text((8, y + ph // 2 - 10), tag, fill=(255, 220, 80), font=font)
        grid.paste(pano, (LABEL_W, y))

    out_path = os.path.join(args.output_dir, "comparison_grid.png")
    grid.save(out_path, quality=95)
    print(f"\nSaved: {out_path} ({grid_w}×{grid_h})")

    # Half-size version
    small = grid.resize((grid_w // 2, grid_h // 2), Image.LANCZOS)
    small.save(out_path.replace(".png", "_half.png"), quality=95)


if __name__ == "__main__":
    main()