render_floorplan.py

"""
Render a floorplan JSON as an SVG or PNG image.
Supports: plot boundary, buildable boundary, rooms (color-coded by type),
          doors, windows, dimension labels, north arrow.
Usage:
    python render_floorplan.py --input floorplan.json --output floorplan.svg
    python render_floorplan.py --input floorplan.json --output floorplan.png
"""
import json
import argparse
from typing import List, Dict, Any, Tuple

ROOM_COLORS = {
    "living": "#E8F4FD", "bedroom": "#FFF3E0", "master_bedroom": "#FFE0B2",
    "kitchen": "#F3E5F5", "dining": "#E0F2F1", "toilet": "#FFEBEE",
    "bathroom": "#FFEBEE", "pooja": "#FFFDE7", "study": "#E8F5E9",
    "balcony": "#F5F5F5", "parking": "#ECEFF1", "staircase": "#FBE9E7",
    "corridor": "#ECEFF1", "utility": "#ECEFF1", "store": "#ECEFF1",
}
ROOM_STROKES = {
    "living": "#1976D2", "bedroom": "#F57C00", "master_bedroom": "#E65100",
    "kitchen": "#7B1FA2", "dining": "#00796B", "toilet": "#C62828",
    "bathroom": "#C62828", "pooja": "#F9A825", "study": "#388E3C",
    "balcony": "#616161", "parking": "#455A64", "staircase": "#D84315",
    "corridor": "#78909C", "utility": "#78909C", "store": "#78909C",
}

def polygon_bbox(poly):
    xs = [p[0] for p in poly]; ys = [p[1] for p in poly]
    return min(xs), min(ys), max(xs), max(ys)

def polygons_bbox(polys):
    all_x, all_y = [], []
    for poly in polys:
        for p in poly: all_x.append(p[0]); all_y.append(p[1])
    return min(all_x), min(all_y), max(all_x), max(all_y)

def render_floorplan_svg(floorplan, width=1200, padding=80, show_dimensions=True, show_labels=True):
    plot = floorplan.get("plot", {})
    rooms = floorplan.get("rooms", [])
    doors = floorplan.get("doors", [])
    windows = floorplan.get("windows", [])

    all_polys = []
    if plot.get("outer_boundary"): all_polys.append(plot["outer_boundary"])
    if plot.get("buildable_boundary"): all_polys.append(plot["buildable_boundary"])
    for room in rooms:
        if room.get("polygon"): all_polys.append(room["polygon"])

    if not all_polys:
        return '<svg xmlns="http://www.w3.org/2000/svg" width="100" height="100"></svg>'

    minx, miny, maxx, maxy = polygons_bbox(all_polys)
    plot_w, plot_h = maxx - minx, maxy - miny
    scale = (width - 2 * padding) / max(plot_w, plot_h)
    svg_h = int(plot_h * scale + 2 * padding)
    svg_w = int(plot_w * scale + 2 * padding)

    def tx(x): return (x - minx) * scale + padding
    def ty(y): return svg_h - ((y - miny) * scale + padding)
    def tpts(poly): return " ".join(f"{tx(p[0])},{ty(p[1])}" for p in poly)

    lines = []
    lines.append(f'<svg xmlns="http://www.w3.org/2000/svg" width="{svg_w}" height="{svg_h}" viewBox="0 0 {svg_w} {svg_h}">')
    lines.append(f'<rect width="{svg_w}" height="{svg_h}" fill="white"/>')

    project_name = floorplan.get("project_name", "Floor Plan")
    lines.append(f'<text x="{svg_w//2}" y="30" text-anchor="middle" font-family="Arial,sans-serif" font-size="18" font-weight="bold">{project_name}</text>')

    if plot.get("outer_boundary"):
        lines.append(f'<polygon points="{tpts(plot["outer_boundary"])}" fill="none" stroke="#333" stroke-width="3"/>')
    if plot.get("buildable_boundary"):
        lines.append(f'<polygon points="{tpts(plot["buildable_boundary"])}" fill="none" stroke="#666" stroke-width="2" stroke-dasharray="8,4"/>')

    for room in rooms:
        poly = room.get("polygon", [])
        if not poly: continue
        rtype = room.get("type", "unknown")
        fill = ROOM_COLORS.get(rtype, "#F5F5F5")
        stroke = ROOM_STROKES.get(rtype, "#333")
        lines.append(f'<polygon points="{tpts(poly)}" fill="{fill}" stroke="{stroke}" stroke-width="1.5"/>')

        if show_labels:
            rminx, rminy, rmaxx, rmaxy = polygon_bbox(poly)
            cx, cy = tx((rminx + rmaxx) / 2), ty((rminy + rmaxy) / 2)
            name = room.get("name", rtype)
            area = room.get("area_sqm", "")
            label = f"{name}"
            if area: label += f" ({area}m\u00b2)"
            lines.append(f'<text x="{cx}" y="{cy+4}" text-anchor="middle" font-family="Arial,sans-serif" font-size="10" fill="#333">{label}</text>')

    for door in doors:
        pos = door.get("position", [0, 0])
        dw = door.get("width", 0.9) * scale
        orient = door.get("orientation", "horizontal")
        x, y = tx(pos[0]), ty(pos[1])
        if orient == "horizontal":
            lines.append(f'<line x1="{x-dw/2}" y1="{y}" x2="{x+dw/2}" y2="{y}" stroke="#2E7D32" stroke-width="4"/>')
        else:
            lines.append(f'<line x1="{x}" y1="{y-dw/2}" x2="{x}" y2="{y+dw/2}" stroke="#2E7D32" stroke-width="4"/>')

    for win in windows:
        pos = win.get("position", [0, 0])
        ww = win.get("width", 1.2) * scale
        orient = win.get("orientation", "horizontal")
        x, y = tx(pos[0]), ty(pos[1])
        if orient == "horizontal":
            lines.append(f'<line x1="{x-ww/2}" y1="{y}" x2="{x+ww/2}" y2="{y}" stroke="#1565C0" stroke-width="3" stroke-dasharray="4,2"/>')
        else:
            lines.append(f'<line x1="{x}" y1="{y-ww/2}" x2="{x}" y2="{y+ww/2}" stroke="#1565C0" stroke-width="3" stroke-dasharray="4,2"/>')

    # North arrow
    arrow_x, arrow_y = svg_w - 60, 60
    lines.append(f'<g transform="translate({arrow_x}, {arrow_y})">')
    lines.append('<polygon points="0,-20 -6,10 6,10" fill="#D32F2F"/>')
    lines.append(f'<text x="0" y="25" text-anchor="middle" font-family="Arial,sans-serif" font-size="12" font-weight="bold">N</text>')
    lines.append('</g>')

    # Legend
    legend_x, legend_y = 20, svg_h - 20 - len(ROOM_COLORS) * 18
    lines.append(f'<rect x="{legend_x-5}" y="{legend_y-15}" width="150" height="{len(ROOM_COLORS)*18+25}" fill="white" stroke="#ccc" stroke-width="1"/>')
    lines.append(f'<text x="{legend_x}" y="{legend_y}" font-family="Arial,sans-serif" font-size="11" font-weight="bold">Legend</text>')
    for i, (rtype, color) in enumerate(ROOM_COLORS.items()):
        ly = legend_y + 15 + i * 16
        lines.append(f'<rect x="{legend_x}" y="{ly-8}" width="12" height="12" fill="{color}" stroke="#666" stroke-width="0.5"/>')
        lines.append(f'<text x="{legend_x+18}" y="{ly+2}" font-family="Arial,sans-serif" font-size="10">{rtype.replace("_", " ").title()}</text>')

    dims = floorplan.get("dimensions", {})
    meta_text = f"Built-up: {dims.get('total_built_up_area_sqm', 'N/A')}m\u00b2 | Carpet: {dims.get('total_carpet_area_sqm', 'N/A')}m\u00b2"
    lines.append(f'<text x="{svg_w//2}" y="{svg_h-10}" text-anchor="middle" font-family="Arial,sans-serif" font-size="11" fill="#666">{meta_text}</text>')
    lines.append('</svg>')
    return "\n".join(lines)

def render_floorplan_png(floorplan, output_path, width=1200):
    try:
        import cairosvg
    except ImportError:
        raise ImportError("cairosvg required for PNG. Install: pip install cairosvg")
    svg_content = render_floorplan_svg(floorplan, width=width)
    cairosvg.svg2png(bytestring=svg_content.encode('utf-8'), write_to=output_path,
                     output_width=width, output_height=int(width * 0.75))
    print(f"PNG saved to {output_path}")

def main():
    parser = argparse.ArgumentParser(description="Render a floorplan JSON to SVG/PNG")
    parser.add_argument("--input", type=str, required=True, help="Path to floorplan JSON")
    parser.add_argument("--output", type=str, required=True, help="Output path (.svg or .png)")
    parser.add_argument("--width", type=int, default=1200)
    parser.add_argument("--no-labels", action="store_true")
    parser.add_argument("--no-dimensions", action="store_true")
    args = parser.parse_args()

    with open(args.input) as f:
        floorplan = json.load(f)

    show_labels = not args.no_labels
    show_dimensions = not args.no_dimensions

    if args.output.lower().endswith(".svg"):
        svg = render_floorplan_svg(floorplan, width=args.width, show_labels=show_labels, show_dimensions=show_dimensions)
        with open(args.output, "w") as f:
            f.write(svg)
        print(f"SVG saved to {args.output}")
    elif args.output.lower().endswith(".png"):
        render_floorplan_png(floorplan, args.output, width=args.width)
    else:
        print("Error: output must be .svg or .png")

if __name__ == "__main__":
    main()