#!/usr/bin/env python3
from __future__ import annotations

import argparse
import csv
import json
import math
from dataclasses import asdict, dataclass
from pathlib import Path

import numpy as np
from PIL import Image, ImageDraw
from skimage import color, feature, measure, morphology, transform


IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".webp"}


@dataclass
class ImageFeatures:
    file: str
    stem: str
    width: int
    height: int
    foreground_fraction: float
    component_count: int
    largest_component_fraction: float
    hole_count: int
    weighted_compactness: float
    edge_density: float
    color_bins: int
    circular_signal: float
    rectilinear_signal: float
    thin_profile_signal: float
    symmetry_signal: float
    gear_signal: float
    standard_signal: float
    complexity_penalty: float
    restorability_score: float
    recommended_backend: str
    confidence: str
    reasons: str


def clamp(value: float, lo: float = 0.0, hi: float = 1.0) -> float:
    return max(lo, min(hi, value))


def percentile_score(value: float, p10: float, p90: float, invert: bool = False) -> float:
    if p90 <= p10:
        score = 0.5
    else:
        score = clamp((value - p10) / (p90 - p10))
    return 1.0 - score if invert else score


def list_images(image_dir: Path) -> list[Path]:
    return sorted(
        p for p in image_dir.iterdir()
        if p.is_file() and p.suffix.lower() in IMAGE_EXTS
    )


def load_small_image(path: Path, max_width: int) -> tuple[np.ndarray, tuple[int, int]]:
    image = Image.open(path).convert("RGB")
    original_size = image.size
    if image.width > max_width:
        height = round(image.height * max_width / image.width)
        image = image.resize((max_width, height), Image.Resampling.LANCZOS)
    return np.asarray(image, dtype=np.uint8), original_size


def foreground_mask(rgb: np.ndarray) -> np.ndarray:
    arr = rgb.astype(np.float32) / 255.0
    hsv = color.rgb2hsv(arr)
    value = hsv[..., 2]
    saturation = hsv[..., 1]

    # SOLIDWORKS screenshots have pale gradient backgrounds and many light-gray
    # parts, so include medium-light gray entity pixels as well as colored parts.
    mask = (value < 0.88) | ((saturation > 0.09) & (value < 0.99))
    mask = morphology.closing(mask, morphology.disk(2))
    mask = morphology.remove_small_objects(mask, max_size=max(32, mask.size // 9000))
    mask = morphology.remove_small_holes(mask, max_size=max(32, mask.size // 14000))
    return mask


def crop_component(mask: np.ndarray, label_image: np.ndarray, label_id: int) -> np.ndarray:
    ys, xs = np.where(label_image == label_id)
    if len(xs) == 0:
        return np.zeros((1, 1), dtype=bool)
    pad = 4
    y0 = max(int(ys.min()) - pad, 0)
    y1 = min(int(ys.max()) + pad + 1, mask.shape[0])
    x0 = max(int(xs.min()) - pad, 0)
    x1 = min(int(xs.max()) + pad + 1, mask.shape[1])
    return label_image[y0:y1, x0:x1] == label_id


def symmetry_for(component: np.ndarray) -> float:
    if component.size <= 1:
        return 0.0
    target = transform.resize(
        component.astype(float),
        (64, 64),
        order=0,
        preserve_range=True,
        anti_aliasing=False,
    ) > 0.5
    area = target.sum()
    if area == 0:
        return 0.0
    horiz = np.logical_and(target, np.fliplr(target)).sum() / np.logical_or(target, np.fliplr(target)).sum()
    vert = np.logical_and(target, np.flipud(target)).sum() / np.logical_or(target, np.flipud(target)).sum()
    return float(max(horiz, vert))


def raw_features(path: Path, max_width: int) -> dict[str, float | int | str]:
    rgb, (width, height) = load_small_image(path, max_width=max_width)
    gray = color.rgb2gray(rgb)
    mask = foreground_mask(rgb)
    label_image = measure.label(mask)
    props = [
        prop for prop in measure.regionprops(label_image)
        if prop.area >= max(48, mask.size // 4500)
    ]
    props.sort(key=lambda item: item.area, reverse=True)

    foreground_fraction = float(mask.mean())
    component_count = len(props)
    largest = props[0] if props else None
    largest_component_fraction = float((largest.area / mask.size) if largest else 0.0)
    total_area = sum(prop.area for prop in props) or 1

    compactness_values = []
    holes = 0
    circular_signal = 0.0
    rectilinear_signal = 0.0
    thin_profile_signal = 0.0
    gear_signal = 0.0
    symmetry_signal = 0.0

    for index, prop in enumerate(props[:16]):
        area = max(float(prop.area), 1.0)
        perimeter = max(float(prop.perimeter), 1.0)
        compactness = perimeter * perimeter / (4.0 * math.pi * area)
        compactness_values.append(compactness * (area / total_area))
        component_holes = max(0, 1 - int(prop.euler_number))
        holes += component_holes

        minr, minc, maxr, maxc = prop.bbox
        box_h = maxr - minr
        box_w = maxc - minc
        box_area = max(float(box_h * box_w), 1.0)
        aspect = max(box_w / max(box_h, 1), box_h / max(box_w, 1))
        extent = prop.area / box_area
        solidity = float(getattr(prop, "solidity", 0.0))
        eccentricity = float(getattr(prop, "eccentricity", 1.0))

        circular = clamp((1.0 - eccentricity) * 1.7) * clamp(solidity * 1.25) * clamp(2.2 - compactness / 2.1)
        rect = clamp((extent - 0.35) / 0.45) * clamp(solidity * 1.15)
        thin = clamp((aspect - 3.0) / 8.0) * clamp(solidity * 1.2)

        weight = area / total_area
        circular_signal = max(circular_signal, circular)
        rectilinear_signal = max(rectilinear_signal, rect)
        thin_profile_signal = max(thin_profile_signal, thin)

        if index < 4:
            comp = crop_component(mask, label_image, prop.label)
            symmetry_signal = max(symmetry_signal, symmetry_for(comp) * clamp(weight * 4.0))

        teeth_like = clamp((compactness - 2.2) / 7.5)
        gear_signal = max(
            gear_signal,
            clamp((1.0 - eccentricity) * 2.0) * teeth_like * clamp((component_holes + 1) / 3.0),
        )

    weighted_compactness = float(sum(compactness_values))
    edges = feature.canny(gray, sigma=1.1)
    edge_density = float(edges.mean())

    object_pixels = rgb[mask]
    if len(object_pixels) == 0:
        color_bins = 0
    else:
        quantized = (object_pixels // 32).astype(np.uint8)
        color_bins = int(len(np.unique(quantized.reshape(-1, 3), axis=0)))

    return {
        "file": str(path),
        "stem": path.stem,
        "width": width,
        "height": height,
        "foreground_fraction": foreground_fraction,
        "component_count": component_count,
        "largest_component_fraction": largest_component_fraction,
        "hole_count": holes,
        "weighted_compactness": weighted_compactness,
        "edge_density": edge_density,
        "color_bins": color_bins,
        "circular_signal": float(circular_signal),
        "rectilinear_signal": float(rectilinear_signal),
        "thin_profile_signal": float(thin_profile_signal),
        "symmetry_signal": float(symmetry_signal),
        "gear_signal": float(gear_signal),
    }


def score_rows(rows: list[dict[str, float | int | str]]) -> list[ImageFeatures]:
    compactness = np.asarray([float(row["weighted_compactness"]) for row in rows])
    edge_density = np.asarray([float(row["edge_density"]) for row in rows])
    components = np.asarray([float(row["component_count"]) for row in rows])
    color_bins = np.asarray([float(row["color_bins"]) for row in rows])
    holes = np.asarray([float(row["hole_count"]) for row in rows])

    p = {
        "compactness": np.percentile(compactness, [10, 90]),
        "edge": np.percentile(edge_density, [10, 90]),
        "components": np.percentile(components, [10, 90]),
        "colors": np.percentile(color_bins, [10, 90]),
        "holes": np.percentile(holes, [50, 95]),
    }

    scored: list[ImageFeatures] = []
    for row in rows:
        comp_score = percentile_score(float(row["weighted_compactness"]), *p["compactness"], invert=True)
        edge_score = percentile_score(float(row["edge_density"]), *p["edge"], invert=True)
        component_score = percentile_score(float(row["component_count"]), *p["components"], invert=True)
        color_score = percentile_score(float(row["color_bins"]), *p["colors"], invert=True)
        many_hole_penalty = percentile_score(float(row["hole_count"]), *p["holes"])

        circular = float(row["circular_signal"])
        rect = float(row["rectilinear_signal"])
        thin = float(row["thin_profile_signal"])
        symmetry = float(row["symmetry_signal"])
        gear = float(row["gear_signal"])
        holes_good = 1.0 if 1 <= int(row["hole_count"]) <= 12 else 0.35 if int(row["hole_count"]) == 0 else 0.0
        standard = max(circular, rect, thin, gear)

        complexity_penalty = (
            0.34 * (1.0 - comp_score)
            + 0.28 * (1.0 - edge_score)
            + 0.25 * (1.0 - component_score)
            + 0.13 * (1.0 - color_score)
        )

        special_family = (
            gear >= 0.80
            and circular >= 0.40
            and float(row["edge_density"]) >= 0.048
        )

        score = (
            25.0
            + 16.0 * comp_score
            + 12.0 * edge_score
            + 8.0 * component_score
            + 4.0 * color_score
            + 12.0 * standard
            + 5.0 * symmetry
            + 4.0 * holes_good
            + (7.0 * gear if special_family else 0.0)
            - 10.0 * many_hole_penalty
        )
        score = max(0.0, min(100.0, score))

        if special_family:
            backend = "simplecadapi"
        elif score >= 72.0:
            backend = "build123d"
        elif score >= 58.0:
            backend = "build123d_or_manual_review"
        else:
            backend = "manual_review_before_modeling"

        if score >= 82:
            confidence = "high"
        elif score >= 68:
            confidence = "medium"
        else:
            confidence = "low"

        reasons = []
        if special_family:
            reasons.append("gear/ring/rack-like standard mechanical family")
        if circular >= 0.45:
            reasons.append("circular or coaxial feature signal")
        if rect >= 0.55:
            reasons.append("simple plate/block extrusion signal")
        if thin >= 0.45:
            reasons.append("shaft/rail/thin extrusion signal")
        if symmetry >= 0.55:
            reasons.append("strong silhouette symmetry")
        if holes_good >= 1.0:
            reasons.append("moderate visible hole count")
        if complexity_penalty >= 0.55:
            reasons.append("penalized for dense edges/components/colors")
        if not reasons:
            reasons.append("generic low-complexity B-Rep signal")

        scored.append(
            ImageFeatures(
                restorability_score=round(score, 3),
                recommended_backend=backend,
                confidence=confidence,
                standard_signal=round(standard, 4),
                complexity_penalty=round(complexity_penalty, 4),
                reasons="; ".join(reasons),
                **{
                    key: row[key] for key in (
                        "file",
                        "stem",
                        "width",
                        "height",
                        "foreground_fraction",
                        "component_count",
                        "largest_component_fraction",
                        "hole_count",
                        "weighted_compactness",
                        "edge_density",
                        "color_bins",
                        "circular_signal",
                        "rectilinear_signal",
                        "thin_profile_signal",
                        "symmetry_signal",
                        "gear_signal",
                    )
                },
            )
        )
    return sorted(scored, key=lambda item: item.restorability_score, reverse=True)


def write_csv(rows: list[ImageFeatures], path: Path) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    fields = list(asdict(rows[0]).keys()) if rows else []
    with path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(handle, fieldnames=fields)
        writer.writeheader()
        for row in rows:
            writer.writerow(asdict(row))


def write_json(rows: list[ImageFeatures], path: Path) -> None:
    path.write_text(
        json.dumps([asdict(row) for row in rows], ensure_ascii=False, indent=2),
        encoding="utf-8",
    )


def make_contact_sheet(rows: list[ImageFeatures], output: Path, limit: int, title: str) -> None:
    selected = rows[:limit]
    if not selected:
        return
    thumb_w, thumb_h = 220, 316
    label_h = 42
    cols = 8
    rows_count = math.ceil(len(selected) / cols)
    sheet = Image.new("RGB", (cols * thumb_w, rows_count * (thumb_h + label_h) + 28), "white")
    draw = ImageDraw.Draw(sheet)
    draw.text((8, 8), title, fill=(0, 0, 0))
    y_offset = 28
    for index, item in enumerate(selected):
        path = Path(item.file)
        image = Image.open(path).convert("RGB")
        image.thumbnail((thumb_w, thumb_h), Image.Resampling.LANCZOS)
        col = index % cols
        row = index // cols
        x = col * thumb_w + (thumb_w - image.width) // 2
        y = y_offset + row * (thumb_h + label_h)
        sheet.paste(image, (x, y))
        label = f"{index + 1:02d} {item.stem}  {item.restorability_score:.1f}"
        draw.text((col * thumb_w + 4, y + thumb_h + 3), label, fill=(0, 0, 0))
        draw.text((col * thumb_w + 4, y + thumb_h + 20), item.recommended_backend[:28], fill=(70, 70, 70))
    output.parent.mkdir(parents=True, exist_ok=True)
    sheet.save(output, quality=90)


def main() -> int:
    parser = argparse.ArgumentParser(description="Rank SOLIDWORKS screenshot restorability for cad-router backends.")
    parser.add_argument("image_dir", type=Path)
    parser.add_argument("--out-dir", type=Path, default=Path("models/sldprt_screen_rank"))
    parser.add_argument("--max-width", type=int, default=420)
    parser.add_argument("--contact-limit", type=int, default=80)
    parser.add_argument("--skip-contact", action="store_true", help="Only write CSV/JSON rankings.")
    args = parser.parse_args()

    paths = list_images(args.image_dir)
    if not paths:
        raise SystemExit(f"No images found in {args.image_dir}")

    raw = []
    for index, path in enumerate(paths, start=1):
        raw.append(raw_features(path, max_width=args.max_width))
        if index % 100 == 0:
            print(f"processed={index}/{len(paths)}", flush=True)
    ranked = score_rows(raw)

    args.out_dir.mkdir(parents=True, exist_ok=True)
    write_csv(ranked, args.out_dir / "ranking.csv")
    write_json(ranked, args.out_dir / "ranking.json")
    if not args.skip_contact:
        make_contact_sheet(ranked, args.out_dir / "top_80_contact_sheet.jpg", args.contact_limit, "Top restorability candidates")
        make_contact_sheet(list(reversed(ranked)), args.out_dir / "bottom_40_contact_sheet.jpg", 40, "Lowest restorability candidates")

    print(f"images={len(ranked)}")
    print(f"csv={args.out_dir / 'ranking.csv'}")
    print(f"json={args.out_dir / 'ranking.json'}")
    print(f"top_sheet={args.out_dir / 'top_80_contact_sheet.jpg'}")
    print("top10=")
    for item in ranked[:10]:
        print(f"{item.stem},{item.restorability_score:.1f},{item.recommended_backend},{item.reasons}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())