src/shading/vine_3d_scene.py

"""
3D scene data and HTML generator for vine, tracker, sun and photosynthesis.

Builds JSON-serializable scene data from ShadowModel + CanopyPhotosynthesisModel,
and renders an interactive Three.js scene showing which parts of the vine
are doing how much photosynthesis (A per zone, colored by rate).
"""

from __future__ import annotations

import json
from datetime import date
from typing import Any

import numpy as np
import pandas as pd


def build_scene_data(
    hour: int = 12,
    date_str: str | None = None,
    par: float = 1800.0,
    tleaf: float = 32.0,
    co2: float = 400.0,
    vpd: float = 2.5,
    tair: float = 33.0,
) -> dict[str, Any]:
    """
    Build scene data for the 3D visualization: sun, tracker, vine geometry,
    shadow mask, PAR and A per zone.

    Returns a dict suitable for JSON serialization and for build_scene_html().
    """
    from src.canopy_photosynthesis import CanopyPhotosynthesisModel
    from src.solar_geometry import ShadowModel

    shadow = ShadowModel()
    canopy = CanopyPhotosynthesisModel(shadow_model=shadow)

    dt_str = date_str or str(date.today())
    try:
        dt = pd.Timestamp(f"{dt_str} {hour:02d}:00:00", tz="Asia/Jerusalem")
    except Exception:
        dt = pd.Timestamp(f"{date.today()} {hour:02d}:00:00", tz="Asia/Jerusalem")

    solar_pos = shadow.get_solar_position(pd.DatetimeIndex([dt]))
    elev = float(solar_pos["solar_elevation"].iloc[0])
    azim = float(solar_pos["solar_azimuth"].iloc[0])

    # Sun direction (world: x=East, y=North, z=up), unit vector toward sun
    elev_rad = np.radians(elev)
    azim_rad = np.radians(azim)
    sun_x = np.cos(elev_rad) * np.sin(azim_rad)
    sun_y = np.cos(elev_rad) * np.cos(azim_rad)
    sun_z = np.sin(elev_rad)
    sun_dir = [float(sun_x), float(sun_y), float(sun_z)]

    if elev <= 2.0:
        # Night: still return geometry, zero A
        tracker_theta = 0.0
        shadow_mask = np.ones((shadow.n_vertical, shadow.n_horizontal), dtype=bool)
        par_zones = np.full((shadow.n_vertical, shadow.n_horizontal), par * 0.15)
        A_zones = np.zeros((shadow.n_vertical, shadow.n_horizontal))
        A_vine = 0.0
        sunlit_fraction = 0.0
    else:
        tracker = shadow.compute_tracker_tilt(azim, elev)
        tracker_theta = float(tracker["tracker_theta"])
        shadow_mask = shadow.project_shadow(elev, azim, tracker_theta)
        vine_result = canopy.compute_vine_A(
            par=par, Tleaf=tleaf, CO2=co2, VPD=vpd, Tair=tair,
            shadow_mask=shadow_mask, solar_elevation=elev,
            solar_azimuth=azim, tracker_tilt=tracker_theta,
        )
        par_zones = vine_result["par_zones"]
        A_zones = vine_result["A_zones"]
        A_vine = float(vine_result["A_vine"])
        sunlit_fraction = float(vine_result["sunlit_fraction"])

    # Panel and vine box in world coords (x=East, y=North, z=up)
    panel_corners = shadow.panel_corners_world(tracker_theta, row_offset=0.0)
    vine_box = shadow.vine_box_world(row_offset=0.0)

    # Grid for zone centres (for positioning vine cells in 3D)
    grid_v = shadow._grid_v.tolist()
    grid_z = shadow._grid_z.tolist()

    def to_list(a: np.ndarray) -> list:
        if a.dtype == bool:
            return [[bool(x) for x in row] for row in a.tolist()]
        return [[float(x) for x in row] for row in a.tolist()]

    return {
        "hour": hour,
        "date": dt_str,
        "sun_elevation": round(elev, 2),
        "sun_azimuth": round(azim, 2),
        "sun_direction": sun_dir,
        "tracker_theta": round(tracker_theta, 2),
        "panel_corners": panel_corners.tolist(),
        "vine_box": vine_box.tolist(),
        "n_vertical": shadow.n_vertical,
        "n_horizontal": shadow.n_horizontal,
        "grid_v": grid_v,
        "grid_z": grid_z,
        "canopy_width": shadow.canopy_width,
        "canopy_height": shadow.canopy_height,
        "shadow_mask": to_list(shadow_mask),
        "par_zones": to_list(par_zones),
        "A_zones": to_list(A_zones),
        "A_vine": round(A_vine, 3),
        "sunlit_fraction": round(sunlit_fraction, 3),
    }


def build_scene_html(scene_data: dict[str, Any], height_px: int = 480) -> str:
    """
    Generate a self-contained HTML file with a Three.js scene: sun, tracker panel,
    vine canopy grid colored by photosynthesis rate (A).
    """
    # Three.js uses Y-up; world is x=East, y=North, z=up → we use (x, z, y) for Three
    def w2t(w: list[float]) -> list[float]:
        return [w[0], w[2], w[1]]

    A_zones = scene_data["A_zones"]
    n_v = scene_data["n_vertical"]
    n_h = scene_data["n_horizontal"]
    grid_v = scene_data["grid_v"]
    grid_z = scene_data["grid_z"]
    cw = scene_data["canopy_width"]
    ch = scene_data["canopy_height"]
    sun_dir = scene_data["sun_direction"]
    panel_corners = scene_data["panel_corners"]
    vine_box = scene_data["vine_box"]
    shadow_mask = scene_data["shadow_mask"]

    A_flat = [A_zones[iz][ih] for iz in range(n_v) for ih in range(n_h)]
    A_min = min(A_flat) if A_flat else 0
    A_max = max(A_flat) if A_flat else 1
    A_range = (A_max - A_min) or 1

    # Color gradient: dark green (low A) -> bright green (high A); shaded can be darker
    def color_for(iz: int, ih: int) -> list[float]:
        a = A_zones[iz][ih]
        shaded = shadow_mask[iz][ih]
        t = (a - A_min) / A_range if A_range else 0
        # 0–1 green gradient; shaded dimmed
        g = 0.2 + 0.7 * t
        r = 0.1
        b = 0.1
        if shaded:
            g *= 0.6
            r *= 0.6
            b *= 0.6
        return [r, g, b]

    # Zone cell size
    dv = (cw / n_h) if n_h else 0.1
    dz = (ch / n_v) if n_v else 0.1
    half_len = 0.4

    cells_json = []
    for iz in range(n_v):
        for ih in range(n_h):
            v_c = grid_v[ih]
            z_c = grid_z[iz]
            # World position of cell centre (row-local v,z; u=0 at centre)
            # In world, row is along u; v is cross-row. We use row_offset=0 so vine at origin.
            # shadow._row_v, _row_u: world x = v*_row_v[0]+u*_row_u[0], same for y. z = z_c
            # For centre of row segment: u=0, v=v_c, z=z_c → world (v_c*_row_v[0], v_c*_row_v[1], z_c)
            # We don't have _row_v in scene_data; approximate: vine_box gives us extent.
            # Simpler: use local v,z and assume row_u points along -Y (315°), row_v along -X
            # So world x ≈ -v_c*cos(45°)= -v_c*0.707, y ≈ v_c*0.707, z=z_c. Actually from settings row_azimuth=315.
            # 315°: along-row = sin(315), cos(315) = -0.707, 0.707. So u direction in world is (-0.707, 0.707, 0).
            # v direction (cross-row) = cos(315), -sin(315) = 0.707, 0.707. So world = (v*0.707, v*0.707, z).
            wx = v_c * 0.707
            wy = v_c * 0.707
            wz = z_c
            cells_json.append({
                "pos": [wx, wz, wy],
                "color": color_for(iz, ih),
                "A": A_zones[iz][ih],
                "shaded": shadow_mask[iz][ih],
            })

    panel_t3 = [w2t(p) for p in panel_corners]
    sun_t3 = w2t(sun_dir)

    # Sun sphere position (far along sun direction)
    sun_dist = 8.0
    sun_pos = [sun_t3[0] * sun_dist, sun_t3[1] * sun_dist, sun_t3[2] * sun_dist]

    scene_json = json.dumps({
        "cells": cells_json,
        "panel": panel_t3,
        "sun_pos": sun_pos,
        "sun_dir": sun_t3,
        "vine_box": [w2t(v) for v in vine_box],
        "A_vine": scene_data["A_vine"],
        "sunlit_fraction": scene_data["sunlit_fraction"],
        "hour": scene_data["hour"],
        "date": scene_data["date"],
        "A_max": A_max,
        "A_min": A_min,
    })

    html = f"""<!DOCTYPE html>
<html lang="en">
<head>
  <meta charset="utf-8">
  <title>Vine photosynthesis 3D</title>
  <style>
    body {{ margin: 0; overflow: hidden; font-family: system-ui, sans-serif; }}
    #info {{ position: absolute; left: 8px; top: 8px; color: #eee; background: rgba(0,0,0,0.6); padding: 8px 12px; border-radius: 8px; font-size: 12px; pointer-events: none; }}
    #legend {{ position: absolute; right: 8px; top: 8px; color: #eee; background: rgba(0,0,0,0.6); padding: 8px 12px; border-radius: 8px; font-size: 11px; pointer-events: none; }}
  </style>
</head>
<body>
  <div id="info">Hour: {scene_data["hour"]:02d}:00 | Date: {scene_data["date"]} | A_vine: {scene_data["A_vine"]:.2f} µmol/m²/s | Sunlit: {scene_data["sunlit_fraction"]*100:.0f}%</div>
  <div id="legend">Green = photosynthesis rate (dark = low, bright = high). Shaded zones are dimmer.</div>
  <script type="importmap">
    {{ "imports": {{
      "three": "https://esm.sh/three",
      "three/addons/OrbitControls": "https://esm.sh/three/examples/jsm/controls/OrbitControls.js"
    }} }}
  </script>
  <script type="module">
    import * as THREE from 'three';
    import {{ OrbitControls }} from 'three/addons/OrbitControls';

    const SCENE = {scene_json};

    const scene = new THREE.Scene();
    scene.background = new THREE.Color(0x87ceeb);

    const camera = new THREE.PerspectiveCamera(50, window.innerWidth / window.innerHeight, 0.1, 100);
    camera.position.set(4, 3, 4);
    camera.lookAt(0, 0.6, 0);

    const renderer = new THREE.WebGLRenderer({{ antialias: true }});
    renderer.setSize(window.innerWidth, window.innerHeight);
    renderer.setPixelRatio(window.devicePixelRatio);
    document.body.appendChild(renderer.domElement);

    const controls = new OrbitControls(camera, renderer.domElement);
    controls.target.set(0, 0.6, 0);
    controls.enableDamping = true;

    // Sun directional light
    const sunLight = new THREE.DirectionalLight(0xffffcc, 1.2);
    sunLight.position.set(...SCENE.sun_pos);
    sunLight.castShadow = false;
    scene.add(sunLight);

    // Ambient
    scene.add(new THREE.AmbientLight(0x4444ff, 0.3));

    // Sun sphere
    const sunGeo = new THREE.SphereGeometry(0.15, 16, 16);
    const sunMat = new THREE.MeshBasicMaterial({{ color: 0xffdd00 }});
    const sunMesh = new THREE.Mesh(sunGeo, sunMat);
    sunMesh.position.set(...SCENE.sun_pos);
    scene.add(sunMesh);

    // Panel (quad)
    const panelGeom = new THREE.BufferGeometry();
    const panelVerts = new Float32Array([
      ...SCENE.panel[0], ...SCENE.panel[1], ...SCENE.panel[2],
      ...SCENE.panel[0], ...SCENE.panel[2], ...SCENE.panel[3],
    ]);
    panelGeom.setAttribute('position', new THREE.BufferAttribute(panelVerts, 3));
    panelGeom.computeVertexNormals();
    const panelMat = new THREE.MeshStandardMaterial({{ color: 0x333333, metalness: 0.6, roughness: 0.4 }});
    const panelMesh = new THREE.Mesh(panelGeom, panelMat);
    scene.add(panelMesh);

    // Vine zone cells (boxes colored by A)
    const cellGeoms = [];
    SCENE.cells.forEach((c, i) => {{
      const w = 0.08, h = 0.08, d = 0.5;
      const box = new THREE.BoxGeometry(w, h, d);
      const mat = new THREE.MeshStandardMaterial({{ color: new THREE.Color(c.color[0], c.color[1], c.color[2]) }});
      const mesh = new THREE.Mesh(box, mat);
      mesh.position.set(c.pos[0], c.pos[1], c.pos[2]);
      mesh.rotation.y = Math.PI / 4;
      scene.add(mesh);
      cellGeoms.push(mesh);
    }});

    window.addEventListener('resize', () => {{
      camera.aspect = window.innerWidth / window.innerHeight;
      camera.updateProjectionMatrix();
      renderer.setSize(window.innerWidth, window.innerHeight);
    }});

    function animate() {{
      requestAnimationFrame(animate);
      controls.update();
      renderer.render(scene, camera);
    }}
    animate();
  </script>
</body>
</html>"""
    return html