#!/usr/bin/env python3
"""
expand_skeleton.py
Rasterize Manhattan polylines (JSON) back to a binary layout image.
"""
import argparse
from pathlib import Path
import json
import numpy as np
import cv2


def draw_polylines(polylines, shape, line_width=3):
    h, w = shape
    img = np.zeros((h, w), dtype='uint8')
    for pl in polylines:
        if len(pl) < 2:
            continue
        pts = [(int(x), int(y)) for x,y in pl]
        for i in range(len(pts)-1):
            p0 = pts[i]
            p1 = pts[i+1]
            cv2.line(img, p0, p1, color=255, thickness=line_width)
    return (img > 127).astype('uint8')


def main():
    p = argparse.ArgumentParser()
    p.add_argument('polyjson')
    p.add_argument('outpng')
    p.add_argument('--line-width', type=int, default=3)
    args = p.parse_args()

    data = json.load(open(args.polyjson,'r'))
    polylines = data.get('polylines', [])
    shape = tuple(data.get('shape', [512,512]))
    img = draw_polylines(polylines, shape, line_width=args.line_width)
    cv2.imwrite(args.outpng, img*255)
    print('Wrote', args.outpng)


if __name__ == '__main__':
    main()