scripts/import_layout.py

#!/usr/bin/env python3
"""
import_layout.py — Spatial asset registration for SolarWine.

Reads site geometry from config/settings.py and the ThingsBoard device
registry, assigns 3D coordinates to every asset (panels, vine rows,
sensors), and writes ``Data/layout.json`` for the ShadowModel and
dashboard map view.

Usage:
    python -m scripts.import_layout            # write Data/layout.json
    python -m scripts.import_layout --print    # dump to stdout instead

The output schema:
{
  "site": { lat, lon, altitude, timezone },
  "panel_geometry": { width, height, row_spacing, row_azimuth_deg },
  "canopy_geometry": { height, width, n_vertical_zones, lai_distribution },
  "rows": [
    { "row_id": 501, "type": "treatment", "devices": [...],
      "panel_center_x_m": ..., "panel_center_y_m": ... },
    ...
  ],
  "devices": {
    "Air1": { uuid, area, row, label, position_m: [x, y, z] },
    ...
  }
}
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

# Add project root to path so we can import config/src
_PROJECT_ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(_PROJECT_ROOT))

from config.settings import (
    CANOPY_HEIGHT,
    CANOPY_WIDTH,
    PANEL_HEIGHT,
    PANEL_WIDTH,
    ROW_SPACING,
    SITE_ALTITUDE,
    SITE_LATITUDE,
    SITE_LONGITUDE,
)


def _load_device_registry() -> dict:
    """Import device registry from ThingsBoard client (extended with `position`)."""
    from src.data.thingsboard_client import DEVICE_REGISTRY
    devices = {}
    for name, info in DEVICE_REGISTRY.items():
        devices[name] = {
            "uuid": info.uuid,
            "area": info.area.value,
            "row": info.row,
            "label": info.label,
            "position": info.position,
        }
    return devices


# Row layout (2026):
#   Reference: row 202 (single row, open sky, lat ~30.9791)
#   Treatment: rows 501/502/503/504/509 (under solar panels, lat ~30.9794)
# Reference plot sits ~30 m south of the treatment plot; modelled as a
# negative Y offset on a separate block.
_TREATMENT_ROWS = [501, 502, 503, 504, 509]
_REFERENCE_ROWS = [202]


def _assign_row_positions(devices: dict) -> list[dict]:
    """Assign 3D positions to every vine row that has at least one device.

    X-axis = row-perpendicular (panel rows are stacked along X).
    Y-axis = along the row.
    Z-axis: ground = 0, panel top = PANEL_HEIGHT.
    """
    row_types: dict[int, str] = {r: "treatment" for r in _TREATMENT_ROWS}
    row_types.update({r: "reference" for r in _REFERENCE_ROWS})

    rows = []
    for i, row_id in enumerate(_TREATMENT_ROWS):
        x = i * ROW_SPACING
        row_devices = [name for name, d in devices.items() if d.get("row") == row_id]
        rows.append({
            "row_id": row_id,
            "type": row_types[row_id],
            "panel_center_x_m": round(x, 2),
            "panel_center_y_m": 0.0,
            "panel_height_m": PANEL_HEIGHT,
            "devices": sorted(row_devices),
        })

    for row_id in _REFERENCE_ROWS:
        row_devices = [name for name, d in devices.items() if d.get("row") == row_id]
        rows.append({
            "row_id": row_id,
            "type": "reference",
            "panel_center_x_m": 0.0,
            "panel_center_y_m": -30.0,    # ~30 m south of treatment block
            "panel_height_m": 0.0,        # reference has no panel
            "devices": sorted(row_devices),
        })

    return rows


# Map `position` strings → (y_offset_m, x_offset_m) relative to row center.
# Y runs along the row (north positive), X is cross-row (east positive).
_POSITION_OFFSETS = {
    "north":       (+5.0,  0.0),
    "south":       (-5.0,  0.0),
    "center":      ( 0.0,  0.0),
    "east":        ( 0.0, +0.5),
    "west":        ( 0.0, -0.5),
    "north-east":  (+5.0, +0.5),
    "north-west":  (+5.0, -0.5),
    "south-east":  (-5.0, +0.5),
    "south-west":  (-5.0, -0.5),
    "center-east": ( 0.0, +0.5),
    "center-west": ( 0.0, -0.5),
}


def _assign_device_positions(devices: dict, rows: list[dict]) -> dict:
    """Assign approximate 3D positions to each device.

    Z by device family:
      Crop_2Soil   : 0.6 m (fruiting zone — IRT/leaf temp/spectrometer/NDVI)
      Thermocouples: PANEL_HEIGHT (panel surface)
      Tracker      : PANEL_HEIGHT (panel pivot)
    """
    row_center = {r["row_id"]: (r["panel_center_x_m"], r["panel_center_y_m"]) for r in rows}

    positioned = {}
    for name, d in devices.items():
        row = d.get("row")
        x0, y0 = row_center.get(row, (0.0, 0.0)) if row else (0.0, 0.0)
        dy, dx = _POSITION_OFFSETS.get((d.get("position") or "").strip().lower(), (0.0, 0.0))
        x = x0 + dx
        y = y0 + dy

        if name.startswith("Crop_2Soil"):
            z = 0.6
        elif name.startswith("Thermocouples"):
            z = PANEL_HEIGHT
        elif name.startswith("Tracker"):
            z = PANEL_HEIGHT
        else:
            z = 0.0

        positioned[name] = {
            **d,
            "position_m": [round(x, 3), round(y, 3), round(z, 3)],
        }

    return positioned


def build_layout() -> dict:
    """Build the complete site layout dictionary."""
    devices = _load_device_registry()
    rows = _assign_row_positions(devices)
    positioned_devices = _assign_device_positions(devices, rows)

    layout = {
        "site": {
            "latitude": SITE_LATITUDE,
            "longitude": SITE_LONGITUDE,
            "altitude_m": SITE_ALTITUDE,
            "timezone": "Asia/Jerusalem",
            "name": "Sde Boker Agrivoltaic Research Site",
        },
        "panel_geometry": {
            "width_m": PANEL_WIDTH,
            "height_m": PANEL_HEIGHT,
            "row_spacing_m": ROW_SPACING,
            "row_azimuth_deg": 315.0,
            "tilt_axis": "single_axis_NS",
        },
        "canopy_geometry": {
            "height_m": CANOPY_HEIGHT,
            "width_m": CANOPY_WIDTH,
            "trellis_type": "VSP",
            "n_vertical_zones": 3,
            "zone_labels": ["basal_trunk", "fruiting_zone", "apical_canopy"],
            "zone_heights_m": [0.2, 0.6, 1.0],
            "lai_distribution": [0.15, 0.35, 0.50],
        },
        "rows": rows,
        "devices": positioned_devices,
    }

    return layout


def main():
    parser = argparse.ArgumentParser(description="Generate Data/layout.json")
    parser.add_argument("--print", action="store_true", help="Print to stdout instead of writing file")
    parser.add_argument("--output", type=str, default=None, help="Custom output path")
    args = parser.parse_args()

    layout = build_layout()
    json_str = json.dumps(layout, indent=2, ensure_ascii=False)

    if args.print:
        print(json_str)
    else:
        out_path = Path(args.output) if args.output else _PROJECT_ROOT / "Data" / "layout.json"
        out_path.parent.mkdir(parents=True, exist_ok=True)
        out_path.write_text(json_str, encoding="utf-8")
        print(f"Layout written to {out_path}")
        print(f"  {len(layout['rows'])} rows, {len(layout['devices'])} devices")


if __name__ == "__main__":
    main()