src/layout_generator/visualizer.py

# src/layout_generator/visualizer.py

"""
Модуль для отрисовки и сохранения 2D-схем планировок и JSON-артефактов.
Генерирует векторные SVG-файлы с видом сверху (Floor Plan) и JSON с точными координатами.
Все артефакты (перезаписываемые и исторические) сохраняются в директорию logs.
"""

import datetime
import json
import logging
from pathlib import Path
from typing import Dict, Any, Tuple, List

import matplotlib.pyplot as plt
import matplotlib.patches as patches
import matplotlib.transforms as transforms
from omegaconf import OmegaConf

logger: logging.Logger = logging.getLogger(__name__)

# Простой кэш размеров для визуализатора
_SIZE_CACHE: Dict[str, Tuple[float, float]] = {}

def _get_fixture_size(fixture_type: str, project_dir: Path) -> Tuple[float, float]:
    """Пытается получить габариты оборудования из кэша или напрямую из геометрии."""
    if fixture_type in _SIZE_CACHE:
        return _SIZE_CACHE[fixture_type]

    w, d = 1.0, 0.5
    try:
        asset_cfg = OmegaConf.load(project_dir / "configs" / "asset_config.yaml")
        asset_cfg_raw = OmegaConf.to_container(asset_cfg, resolve=False)
        fixtures_cfg = asset_cfg_raw.get("fixtures", asset_cfg_raw.get("equipment", {}))
        base_assets_dir = asset_cfg_raw.get("assets_dir_path", "assets")

        if fixture_type in fixtures_cfg:
            fix_data = fixtures_cfg[fixture_type]
            raw_path = fix_data.get("asset_file_path", "")
            if raw_path:
                clean_path = raw_path.replace("${assets.assets_dir_path}", base_assets_dir)
                full_path = project_dir / clean_path
                import trimesh
                scene = trimesh.load(str(full_path), force='mesh')
                w = round(float(scene.extents[0]), 2)
                d = round(float(scene.extents[1]), 2)
    except Exception as e:
        logger.debug(f"⚠️ Визуализатор: не удалось извлечь размер для {fixture_type}: {e}")
        
    _SIZE_CACHE[fixture_type] = (w, d)
    return w, d

def _render_and_save_svg(
    items_data: List[Dict[str, Any]],
    size_n: float,
    size_m: float,
    title: str,
    file_prefix: str,
    project_dir: Path,
    timestamp: str
) -> None:
    """Внутренняя функция для рендеринга и сохранения SVG."""
    fig, ax = plt.subplots(figsize=(12, 12 * (size_m / size_n)))
    ax.set_xlim(0, size_n)
    ax.set_ylim(0, size_m)
    ax.set_aspect('equal')
    ax.grid(True, linestyle='--', alpha=0.5, color='#BDC3C7')
    
    room_rect = patches.Rectangle((0, 0), size_n, size_m, linewidth=3, edgecolor='#2C3E50', facecolor='none')
    ax.add_patch(room_rect)
    
    door_w = 1.5
    door_rect = patches.Rectangle((size_n/2 - door_w/2, -0.1), door_w, 0.2, linewidth=2, edgecolor='#E74C3C', facecolor='#FDEDEC')
    ax.add_patch(door_rect)
    ax.text(size_n/2, -0.3, "ВХОД", ha='center', va='center', fontsize=10, color='#E74C3C', fontweight='bold')

    for item in items_data:
        cx, cy = item['cx'], item['cy']
        w, d = item['w'], item['d']
        rot = item['rot']
        label = item['label']
        facecolor = item['facecolor']
        edgecolor = item['edgecolor']
        
        rect = patches.Rectangle((cx - w/2, cy - d/2), w, d, linewidth=2, edgecolor=edgecolor, facecolor=facecolor, alpha=0.8)
        t = transforms.Affine2D().rotate_deg_around(cx, cy, rot) + ax.transData
        rect.set_transform(t)
        ax.add_patch(rect)
        
        ax.text(cx, cy, label, ha='center', va='center', fontsize=7, color='#17202A', 
                bbox=dict(facecolor='white', alpha=0.7, edgecolor='none', boxstyle='round,pad=0.2'))

    current_time_str = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    plt.title(f"{title}\nСгенерировано: {current_time_str}", fontsize=14, pad=15)
    plt.xlabel("Ширина (X, метры)")
    plt.ylabel("Глубина (Y, метры)")
    
    # Сохраняем ТОЛЬКО в папку logs
    logs_path = project_dir / "logs"
    logs_path.mkdir(exist_ok=True)
    
    base_file = logs_path / f"{file_prefix}.svg"
    timestamp_file = logs_path / f"{file_prefix}_{timestamp}.svg"
    
    plt.savefig(base_file, format='svg', bbox_inches='tight')
    plt.savefig(timestamp_file, format='svg', bbox_inches='tight')
    
    plt.close()
    logger.info(f"🗺️ Векторный план '{title}' сохранен в {base_file.name} и {timestamp_file.name}")

def draw_llm_draft_plan(graph_data: Dict[str, Any], project_dir: Path) -> None:
    """Отрисовывает черновик (Draft) от нейросети и сохраняет его JSON."""
    if not graph_data or "nodes" not in graph_data:
        return
        
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    
    try:
        logs_path = project_dir / "logs"
        logs_path.mkdir(exist_ok=True)
        
        with open(logs_path / "layout_draft.json", "w", encoding="utf-8") as f:
            json.dump(graph_data, f, ensure_ascii=False, indent=4)
        with open(logs_path / f"layout_draft_{timestamp}.json", "w", encoding="utf-8") as f:
            json.dump(graph_data, f, ensure_ascii=False, indent=4)
        logger.info("📄 JSON черновика (LLM Draft) успешно сохранен в папку logs.")

        layout_cfg = OmegaConf.load(project_dir / "configs" / "layout_config.yaml")
        size_n = float(layout_cfg.layout.size_n)
        size_m = float(layout_cfg.layout.size_m)
        
        items_data = []
        for node in graph_data.get("nodes", []):
            cx, cy = node.get("draft_cx"), node.get("draft_cy")
            if cx is None or cy is None:
                continue
                
            fixture_type = node.get("type", "unknown")
            w, d = _get_fixture_size(fixture_type, project_dir)
            rot = float(node.get("draft_rotation", 0.0))
            
            if "fridge" in fixture_type.lower() or "showcase" in fixture_type.lower():
                fc, ec = "#D6EAF8", "#2E86C1"
            else:
                fc, ec = "#D5F5E3", "#28B463"
                
            items_dict = node.get("items", {})
            items_info = [f"{k}: {v}" for k, v in items_dict.items() if v > 0]
            label = f"{node.get('id', 'unknown')}\n" + "\n".join(items_info)
            
            items_data.append({'cx': float(cx), 'cy': float(cy), 'w': w, 'd': d, 'rot': rot, 'label': label, 'facecolor': fc, 'edgecolor': ec})
            
        _render_and_save_svg(items_data, size_n, size_m, "Черновой 2D-план (LLM Draft AABB)", "layout_draft", project_dir, timestamp)
    except Exception as e:
        logger.error(f"❌ Ошибка обработки LLM Draft: {e}", exc_info=True)


def draw_final_topology_plan(placed_fixtures: List[Any], size_n: float, size_m: float, project_dir: Path) -> None:
    """Отрисовывает финальную расстановку (Final) и сохраняет ее JSON."""
    if not placed_fixtures:
        return
        
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    
    try:
        final_data: Dict[str, Any] = {"nodes": []}
        for f in placed_fixtures:
            final_data["nodes"].append({
                "id": getattr(f, 'name', 'unknown'),
                "type": getattr(f, 'asset_name', 'unknown'),
                "cx": round(f.x, 3),
                "cy": round(f.y, 3),
                "w": round(f.l, 3),
                "d": round(f.w, 3),
                "rotation": round(getattr(f, '_rotation_deg', 0.0), 2)
            })
            
        logs_path = project_dir / "logs"
        logs_path.mkdir(exist_ok=True)
        
        with open(logs_path / "layout_final.json", "w", encoding="utf-8") as f:
            json.dump(final_data, f, ensure_ascii=False, indent=4)
        with open(logs_path / f"layout_final_{timestamp}.json", "w", encoding="utf-8") as f:
            json.dump(final_data, f, ensure_ascii=False, indent=4)
        logger.info("📄 JSON финальной топологии успешно сохранен в папку logs.")

        items_data = []
        for f in placed_fixtures:
            cx, cy = f.x, f.y
            w, d = f.l, f.w  
            rot = getattr(f, '_rotation_deg', 0.0)
            
            name = getattr(f, 'name', 'unknown')
            asset_name = getattr(f, 'asset_name', 'unknown')
            
            if "fridge" in asset_name.lower() or "showcase" in asset_name.lower():
                fc, ec = "#FDEBD0", "#D68910"
            else:
                fc, ec = "#D4E6F1", "#2471A3"
                
            label = f"{name}"
            
            items_data.append({'cx': cx, 'cy': cy, 'w': w, 'd': d, 'rot': rot, 'label': label, 'facecolor': fc, 'edgecolor': ec})
            
        _render_and_save_svg(items_data, size_n, size_m, "Финальный 2D-план (После физики Topology)", "layout_final", project_dir, timestamp)
    except Exception as e:
        logger.error(f"❌ Ошибка обработки финального плана: {e}", exc_info=True)