#!/usr/bin/env catnip
# Vision par ordinateur avec OpenCV
# OpenCV : espaces colorimétriques, pipelines, batch, détourage, détection
#
# DEPS: opencv-python

cv2 = import('cv2')
np = import('numpy')
sys = import('sys')
import('pathlib', 'Path')

# Résolution relative au script (fonctionne depuis n'importe quel CWD)
script_dir = Path(META.file).parent
output_dir = script_dir / "output"
output_dir.mkdir(exist_ok=True)

# Sauvegarde sans pertes d'une matrice (OpenCV encode les images couleur depuis BGR)
save = (name: str, mat) => {
    filename = f"catnip_cv_{name}.png"
    path = output_dir / filename
    if not cv2.imwrite(str(path), mat) {
        print(f"✗ Impossible d'écrire {path}")
        sys.exit(1)
    }
    print(f"  {name} → {filename}")
    mat
}

# Struct Op : une opération nommée appliquée à une image (pipeline déclaratif)
struct Op {
    name: str; apply

    run(self, img) => { save(self.name, self.apply(img)) }
}

# Chargement (OpenCV lit en BGR, pas RGB)
print("⇒ Chargement")
input_path = script_dir / "data/catnip_test.jpg"
img = cv2.imread(str(input_path))
if img is None {
    print(f"✗ Impossible de lire {input_path}")
    sys.exit(1)
}
h = img.shape[0]
w = img.shape[1]
print(f"  taille: {w}x{h}, canaux: {img.shape[2]}, dtype: {img.dtype}")

# Espaces colorimétriques
print()
print("⇒ Espaces colorimétriques")
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
save("gray", gray)

# Une matrice HSV n'est pas directement affichable comme une image BGR.
# La teinte est colorisée à saturation et luminosité maximales ; les deux
# autres composantes sont affichées en niveaux de gris.
hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
hsv_channels = cv2.split(hsv)
hue_preview = cv2.merge(tuple(
        hsv_channels[0],
        np.full_like(hsv_channels[0], 255),
        np.full_like(hsv_channels[0], 255),
    ))
save("hsv_hue", cv2.cvtColor(hue_preview, cv2.COLOR_HSV2BGR))
save("hsv_saturation", hsv_channels[1])
save("hsv_value", hsv_channels[2])

# Transformations géométriques via broadcasting
print()
print("⇒ Transformations géométriques")
rot_mat = cv2.getRotationMatrix2D(tuple(w / 2, h / 2), 30, 1.0)

geo = list(
    Op("resized", (i) => { cv2.resize(i, tuple(200, 150)) }),
    Op("rotated", (i) => { cv2.warpAffine(i, rot_mat, tuple(w, h)) }),
    Op("flip_h", (i) => { cv2.flip(i, 1) }),
    Op("flip_v", (i) => { cv2.flip(i, 0) }),
)
geo.[(o: Op) => { o.run(img) }]

# Filtres et lissage via broadcasting
print()
print("⇒ Filtres et lissage")

filters = list(
    Op("gaussian", (i) => { cv2.GaussianBlur(i, tuple(9, 9), 0) }),
    Op("median", (i) => { cv2.medianBlur(i, 7) }),
    Op("bilateral", (i) => { cv2.bilateralFilter(i, 9, 75, 75) }),
)
filters.[(o: Op) => { o.run(img) }]

# Pipeline séquentiel : fold fait circuler l'image dans chaque opération
print()
print("⇒ Pipeline Catnip : gris → flou → contours")

edge_pipeline = list(
    Op("to_gray", (i) => { cv2.cvtColor(i, cv2.COLOR_BGR2GRAY) }),
    Op("denoise", (i) => { cv2.GaussianBlur(i, tuple(5, 5), 0) }),
    Op("edges", (i) => { cv2.Canny(i, 60, 140) }),
)

run_pipeline = (source) => {
    fold(edge_pipeline, source, (current, operation) => { operation.apply(current) })
}

save("pipeline_edges", run_pipeline(img))

# Batch : le même pipeline est diffusé sur plusieurs variantes de l'image
print()
print("⇒ Batch Catnip : un pipeline, plusieurs images")

struct Variant { name: str; image; }

variants = list(
    Variant("original", img),
    Variant("flip_h", cv2.flip(img, 1)),
    Variant("brighter", cv2.convertScaleAbs(img, alpha=1.2, beta=20)),
    Variant("rotated", cv2.warpAffine(img, rot_mat, tuple(w, h))),
)

batch_edges = variants.[(variant: Variant) => {
    Variant(variant.name, run_pipeline(variant.image))
}]
batch_edges.[(variant: Variant) => { save(f"batch_{variant.name}", variant.image) }]

# Une planche-contact étiquetée rend le batch visible en un coup d'œil
label_result = (variant: Variant) => {
    canvas = cv2.cvtColor(variant.image, cv2.COLOR_GRAY2BGR)
    cv2.putText(canvas, variant.name, tuple(20, 45), cv2.FONT_HERSHEY_SIMPLEX, 1.2, tuple(0, 255, 0), 2)
    canvas
}
contact_sheet = cv2.hconcat(batch_edges.[(variant: Variant) => { label_result(variant) }])
save("batch_contact_sheet", contact_sheet)

# Détourage classique sans modèle externe : NumPy construit un masque guidé,
# puis GrabCut affine la séparation entre le sujet et l'arrière-plan.
print()
print("⇒ Détourage automatique (GrabCut)")

grabcut_mask = np.where(np.less(gray, 100), cv2.GC_PR_FGD, cv2.GC_PR_BGD).astype(np.uint8)
# Le centre du sujet est certain ; la bordure et la plomberie à droite sont du fond.
cv2.ellipse(grabcut_mask, tuple(330, 440), tuple(75, 180), 0, 0, 360, cv2.GC_FGD, -1)
cv2.rectangle(grabcut_mask, tuple(0, 0), tuple(w - 1, h - 1), cv2.GC_BGD, 8)
cv2.rectangle(grabcut_mask, tuple(610, 0), tuple(w - 1, h - 1), cv2.GC_BGD, -1)
mat_hint = np.array(list(tuple(0, 775), tuple(170, h - 1), tuple(0, h - 1)), dtype=np.int32)
cv2.fillPoly(grabcut_mask, list(mat_hint), cv2.GC_BGD)

background_model = np.zeros(tuple(1, 65), np.float64)
foreground_model = np.zeros(tuple(1, 65), np.float64)

cv2.grabCut(
    img,
    grabcut_mask,
    None,
    background_model,
    foreground_model,
    5,
    cv2.GC_INIT_WITH_MASK,
)

is_background = np.logical_or(
    np.equal(grabcut_mask, cv2.GC_BGD),
    np.equal(grabcut_mask, cv2.GC_PR_BGD),
)
alpha = np.where(is_background, 0, 255).astype(np.uint8)
alpha = cv2.morphologyEx(alpha, cv2.MORPH_OPEN, np.ones(tuple(3, 3), np.uint8))
alpha = cv2.morphologyEx(alpha, cv2.MORPH_CLOSE, np.ones(tuple(11, 11), np.uint8))
alpha = cv2.GaussianBlur(alpha, tuple(5, 5), 0)

bgr_channels = cv2.split(img)
cutout = cv2.merge(tuple(bgr_channels[0], bgr_channels[1], bgr_channels[2], alpha))
cutout_preview = np.full_like(img, 32)
cv2.copyTo(img, alpha, cutout_preview)
save("foreground_mask", alpha)
save("foreground_rgba", cutout)
save("foreground_preview", cutout_preview)

# Détection de contours
print()
print("⇒ Détection de contours")
edges = cv2.Canny(gray, 50, 150)
save("canny", edges)

# threshold renvoie un tuple (seuil, image) -- on garde l'image
save("threshold", cv2.threshold(gray, 127, 255, cv2.THRESH_BINARY)[1])
save("adaptive", cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 11, 2))

# findContours renvoie (contours, hiérarchie) ; on dessine les contours externes
contours = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[0]
outlined = img.copy()
cv2.drawContours(outlined, contours, -1, tuple(0, 255, 0), 1)
save("contours", outlined)
print(f"  {len(contours)} contours externes détectés")

# Morphologie via broadcasting (sur la carte de contours binaire)
print()
print("⇒ Morphologie")
kernel = np.ones(tuple(5, 5), np.uint8)

morph = list(
    Op("erode", (i) => { cv2.erode(i, kernel) }),
    Op("dilate", (i) => { cv2.dilate(i, kernel) }),
    Op("opening", (i) => { cv2.morphologyEx(i, cv2.MORPH_OPEN, kernel) }),
    Op("closing", (i) => { cv2.morphologyEx(i, cv2.MORPH_CLOSE, kernel) }),
)
morph.[(o: Op) => { o.run(edges) }]

# Détection de coins (Shi-Tomasi) : chaque coin est une ligne (1, 2) de l'array
print()
print("⇒ Détection de coins")
corners = cv2.goodFeaturesToTrack(gray, 100, 0.01, 10)
marked = img.copy()
if corners is None {
    print("  aucun coin détecté")
} else {
    for c in corners {
        pt = c[0]
        cv2.circle(marked, tuple(int(pt[0]), int(pt[1])), 4, tuple(0, 0, 255), -1)
    }
    print(f"  {corners.shape[0]} coins marqués")
}
save("corners", marked)

# Interop numpy : ajustements pixel par pixel (les images sont des arrays numpy)
print()
print("⇒ Ajustements numpy")
save("brighter", np.clip(img.astype(np.float32) * 1.3 + 25, 0, 255).astype(np.uint8))
save("inverted", 255 - img)

# Statistiques par canal -- cv2.split sépare BGR en trois plans 2D
print()
print("⇒ Statistiques par canal (BGR)")
names = list("bleu", "vert", "rouge")
channels = cv2.split(img)
for idx in range(3) {
    ch = channels[idx]
    print(f"  {names[idx]} : moyenne={round(float(ch.mean()), 1)}, écart-type={round(float(ch.std()), 1)}")
}

# Résumé
print()
print("⇒ Résumé")
print(f"Images générées dans {output_dir} :")

summary = list(
    "  - Couleur : grayscale, canaux HSV",
    "  - Géométrie : resize, rotate, flip",
    "  - Filtres : gaussian, median, bilateral",
    "  - Catnip : pipeline avec fold, batch avec broadcasting",
    "  - Détourage : masque GrabCut et PNG transparent",
    "  - Détection : canny, threshold, adaptive, contours",
    "  - Morphologie : erode, dilate, opening, closing",
    "  - Features : coins Shi-Tomasi",
    "  - numpy : luminosité, inversion",
)
for line in summary { print(line) }