#!/usr/bin/env python3
"""
skeleton_extract.py
Extract Manhattan-style skeleton from binary layout images.
Supports optional color inversion and basic denoising.
Outputs skeleton PNG and a JSON summary with connected-component counts.
"""
import argparse
from pathlib import Path
import numpy as np
from PIL import Image
import cv2
from skimage.morphology import skeletonize
from skimage import img_as_ubyte
import json


def load_image(path, invert=False):
    im = Image.open(path).convert('L')
    a = np.array(im)
    # auto threshold by Otsu
    _, th = cv2.threshold(a, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    if invert:
        th = 255 - th
    return th // 255


def denoise(bin_img, kernel=3):
    # simple opening/closing
    k = cv2.getStructuringElement(cv2.MORPH_RECT, (kernel, kernel))
    img = cv2.morphologyEx((bin_img * 255).astype('uint8'), cv2.MORPH_OPEN, k)
    img = cv2.morphologyEx(img, cv2.MORPH_CLOSE, k)
    return img // 255


def extract_skeleton(bin_img):
    # skimage expects bool image
    sk = skeletonize(bin_img > 0)
    return sk.astype('uint8')


def save_png(arr, path):
    im = Image.fromarray((arr * 255).astype('uint8'))
    im.save(path)


def main():
    p = argparse.ArgumentParser()
    p.add_argument('input')
    p.add_argument('outdir')
    p.add_argument('--invert', action='store_true', help='Invert black/white before processing')
    p.add_argument('--denoise', type=int, default=3, help='Denoise kernel size')
    args = p.parse_args()

    inp = Path(args.input)
    out = Path(args.outdir)
    out.mkdir(parents=True, exist_ok=True)

    bin_img = load_image(inp, invert=args.invert)
    if args.denoise and args.denoise > 0:
        bin_img = denoise(bin_img, kernel=args.denoise)

    sk = extract_skeleton(bin_img)

    save_png(bin_img, out / (inp.stem + '_bin.png'))
    save_png(sk, out / (inp.stem + '_sk.png'))

    # summary
    num_pixels = int(sk.sum())
    components = cv2.connectedComponents((sk * 255).astype('uint8'))[0] - 1
    info = {'input': str(inp), 'pixels_in_skeleton': int(num_pixels), 'components': int(components)}
    with open(out / (inp.stem + '_sk.json'), 'w') as f:
        json.dump(info, f, indent=2)

    print('Saved:', out / (inp.stem + '_sk.png'))


if __name__ == '__main__':
    main()