Blame for env.py

from typing import Dict, List, Tuple
from collections import deque
import copy

from models import Observation, Action, EpisodeResult, PlacedFurniture, Furniture
from tasks import TASKS

NEAR_WINDOW_RADIUS = 2
NEAR_DOOR_RADIUS = 3

class InteriorEnv:
    def __init__(self):
        self.task_id = ""
        self.room_size = (0, 0)
        self.door_position = (0, 0)
        self.windows: List[Tuple[int, int]] = []
        self.item_defs: Dict[str, Furniture] = {}
        self.unplaced_items: Dict[str, Furniture] = {}
        self.placed_items: Dict[str, PlacedFurniture] = {}
        self.max_steps = 10
        self.step_count = 0
        self.success = False
        self.logs: List[str] = []
        self.solar_grid: List[List[bool]] = []
        self.thermal_grid: List[List[float]] = []

    def reset(self, task_config: dict, task_id: str) -> Observation:
        self.task_id = task_id
        self.room_size = task_config["room_size"]
        self.door_position = task_config["door_position"]
        self.windows = task_config.get("windows", [])
        self.item_defs = {item.id: item for item in task_config["items"]}
        self.unplaced_items = {item.id: item for item in task_config["items"]}
        self.placed_items = {}
        self.max_steps = task_config["max_steps"]
        self.step_count = 0
        self.success = False
        self.solar_grid = [[False] * self.room_size[0] for _ in range(self.room_size[1])]
        self.thermal_grid = [[0.0] * self.room_size[0] for _ in range(self.room_size[1])]
        self.logs = self._build_start_logs(task_config)
        return self._get_obs()

    def _build_start_logs(self, cfg: dict) -> List[str]:
        lines = [
            f"START: Task={self.task_id.upper()} | Room={self.room_size} | Door={self.door_position}",
            f"Windows at: {self.windows}",
            f"Items to place: {len(self.unplaced_items)}",
            "─── Design Rules ────────────────────────────────",
        ]
        for item in cfg["items"]:
            for rule in item.rules:
                lines.append(f"  [{item.id.upper()}]: {self._describe_rule(rule)}")
        lines.append("─────────────────────────────────────────────────")
        return lines

    def _describe_rule(self, rule: str) -> str:
        if rule == "must_touch_wall":
            return "Must be placed flush against any wall."
        if rule == "near_window":
            return f"Must be within {NEAR_WINDOW_RADIUS} tiles of a window."
        if rule == "all_sides_clear":
            return "Must have at least 1 free tile on ALL 4 sides (e.g., dining space)."
        if rule == "near_door":
            return f"Must be within {NEAR_DOOR_RADIUS} tiles of the main entry door."
        if rule.startswith("adjacent_to:"):
            target = rule.split(":")[1]
            return f"Must be placed directly adjacent (touching) the {target.upper()}."
        if rule.startswith("must_face:"):
            target = rule.split(":")[1]
            return f"Must align on the same axis as the {target.upper()} (facing it)."
        return rule

    # ─── Geometry ─────────────────────────────────────────────────────────────
    def _bbox(self, placed: PlacedFurniture, item_def: Furniture) -> Tuple[int, int, int, int]:
        w = item_def.width  if placed.rotation == 0 else item_def.height
        h = item_def.height if placed.rotation == 0 else item_def.width
        return (placed.x, placed.x + w - 1, placed.y, placed.y + h - 1)

    def _bbox_action(self, action: Action, item_def: Furniture) -> Tuple[int, int, int, int]:
        w = item_def.width  if action.rotation == 0 else item_def.height
        h = item_def.height if action.rotation == 0 else item_def.width
        return (action.x, action.x + w - 1, action.y, action.y + h - 1)

    def _build_grid(self) -> List[List[bool]]:
        rw, rh = self.room_size
        grid = [[False] * rw for _ in range(rh)]
        for iid, placed in self.placed_items.items():
            xb, xe, yb, ye = self._bbox(placed, self.item_defs[iid])
            for gy in range(yb, ye + 1):
                for gx in range(xb, xe + 1):
                    if 0 <= gx < rw and 0 <= gy < rh:
                        grid[gy][gx] = True
        return grid

    def _bfs_reachable(self, grid: List[List[bool]]) -> List[List[bool]]:
        rw, rh = self.room_size
        visited = [[False] * rw for _ in range(rh)]
        dx, dy = self.door_position
        if grid[dy][dx]:
            return visited
        queue = deque([(dx, dy)])
        visited[dy][dx] = True
        while queue:
            cx, cy = queue.popleft()
            for nx, ny in [(cx+1,cy),(cx-1,cy),(cx,cy+1),(cx,cy-1)]:
                if 0 <= nx < rw and 0 <= ny < rh and not grid[ny][nx] and not visited[ny][nx]:
                    visited[ny][nx] = True
                    queue.append((nx, ny))
        return visited

    def _all_neighbors(self, pid: str) -> List[Tuple[int, int]]:
        placed = self.placed_items[pid]
        xb, xe, yb, ye = self._bbox(placed, self.item_defs[pid])
        rw, rh = self.room_size
        result = []
        for gy in range(yb, ye + 1):
            for gx in range(xb, xe + 1):
                for nx, ny in [(gx+1,gy),(gx-1,gy),(gx,gy+1),(gx,gy-1)]:
                    if 0 <= nx < rw and 0 <= ny < rh:
                        result.append((nx, ny))
        return result

    def _build_solar_grid(self) -> List[List[bool]]:
        rw, rh = self.room_size
        solar = [[False] * rw for _ in range(rh)]
        tall_grid = [[False] * rw for _ in range(rh)]
        for iid, placed in self.placed_items.items():
            if any(r == "is_tall" for r in self.item_defs[iid].rules):
                xb, xe, yb, ye = self._bbox(placed, self.item_defs[iid])
                for gy in range(yb, ye + 1):
                    for gx in range(xb, xe + 1):
                        if 0 <= gx < rw and 0 <= gy < rh:
                            tall_grid[gy][gx] = True

        for wx, wy in self.windows:
            for dx, dy in [(1,0), (-1,0), (0,1), (0,-1)]:
                cx, cy = wx, wy
                while 0 <= cx < rw and 0 <= cy < rh:
                    if tall_grid[cy][cx]: break
                    solar[cy][cx] = True
                    cx += dx
                    cy += dy
        return solar

    def _build_thermal_grid(self) -> List[List[float]]:
        rw, rh = self.room_size
        thermal = [[0.0] * rw for _ in range(rh)]
        
        cool_sources = self.windows + [self.door_position]
        for cx, cy in cool_sources:
            for gy in range(rh):
                for gx in range(rw):
                    dist = abs(gx - cx) + abs(gy - cy)
                    if dist <= 3:
                        thermal[gy][gx] -= 2.0 / (dist + 1)
                        
        for iid, placed in self.placed_items.items():
            for rule in self.item_defs[iid].rules:
                if rule.startswith("emits_heat:") or rule.startswith("emits_cold:"):
                    val = float(rule.split(":")[1])
                    if rule.startswith("emits_cold:"): val = -val
                    xb, xe, yb, ye = self._bbox(placed, self.item_defs[iid])
                    for gy in range(rh):
                        for gx in range(rw):
                            dx = max(0, xb - gx, gx - xe)
                            dy = max(0, yb - gy, gy - ye)
                            dist = dx + dy
                            if dist <= 3:
                                thermal[gy][gx] += val / (dist + 1)
        return thermal

    # ─── Rule Checkers ────────────────────────────────────────────────────────
    def _check_must_touch_wall(self, pid: str) -> bool:
        rw, rh = self.room_size
        placed = self.placed_items[pid]
        xb, xe, yb, ye = self._bbox(placed, self.item_defs[pid])
        return xb == 0 or xe == rw - 1 or yb == 0 or ye == rh - 1

    def _check_near_window(self, pid: str) -> bool:
        if not self.windows:
            return True
        placed = self.placed_items[pid]
        xb, xe, yb, ye = self._bbox(placed, self.item_defs[pid])
        for wx, wy in self.windows:
            for gy in range(yb, ye + 1):
                for gx in range(xb, xe + 1):
                    if abs(gx - wx) + abs(gy - wy) <= NEAR_WINDOW_RADIUS:
                        return True
        return False

    def _check_near_door(self, pid: str) -> bool:
        dx, dy = self.door_position
        placed = self.placed_items[pid]
        xb, xe, yb, ye = self._bbox(placed, self.item_defs[pid])
        for gy in range(yb, ye + 1):
            for gx in range(xb, xe + 1):
                if abs(gx - dx) + abs(gy - dy) <= NEAR_DOOR_RADIUS:
                    return True
        return False

    def _check_must_face(self, pid: str, target_id: str) -> bool:
        if target_id not in self.placed_items:
            return False
        p = self.placed_items[pid]
        t = self.placed_items[target_id]
        pxb, pxe, pyb, pye = self._bbox(p, self.item_defs[pid])
        txb, txe, tyb, tye = self._bbox(t, self.item_defs[target_id])
        pcx = (pxb + pxe) / 2
        pcy = (pyb + pye) / 2
        tcx = (txb + txe) / 2
        tcy = (tyb + tye) / 2
        aligned_x = abs(pcx - tcx) <= 2
        aligned_y = abs(pcy - tcy) <= 2
        # True if they share one axis (same column or same row) but not both
        return (aligned_x and not aligned_y) or (aligned_y and not aligned_x)

    def _check_adjacent_to(self, pid: str, target_id: str) -> bool:
        if target_id not in self.placed_items:
            return False
        t = self.placed_items[target_id]
        tdef = self.item_defs[target_id]
        txb, txe, tyb, tye = self._bbox(t, tdef)
        for nx, ny in self._all_neighbors(pid):
            if txb <= nx <= txe and tyb <= ny <= tye:
                return True
        return False

    def _check_all_sides_clear(self, pid: str, grid: List[List[bool]]) -> bool:
        rw, rh = self.room_size
        placed = self.placed_items[pid]
        xb, xe, yb, ye = self._bbox(placed, self.item_defs[pid])
        top    = yb > 0     and any(not grid[yb-1][gx] for gx in range(xb, xe+1))
        bottom = ye+1 < rh  and any(not grid[ye+1][gx] for gx in range(xb, xe+1))
        left   = xb > 0     and any(not grid[gy][xb-1] for gy in range(yb, ye+1))
        right  = xe+1 < rw  and any(not grid[gy][xe+1] for gy in range(yb, ye+1))
        return top and bottom and left and right

    def _check_needs_sunlight(self, pid: str, solar_grid: List[List[bool]]) -> bool:
        placed = self.placed_items[pid]
        xb, xe, yb, ye = self._bbox(placed, self.item_defs[pid])
        for gy in range(yb, ye + 1):
            for gx in range(xb, xe + 1):
                if solar_grid[gy][gx]: return True
        return False

    def _check_thermal_comfort(self, pid: str, min_v: float, max_v: float, thermal_grid: List[List[float]]) -> bool:
        placed = self.placed_items[pid]
        xb, xe, yb, ye = self._bbox(placed, self.item_defs[pid])
        total = 0.0
        count = 0
        for gy in range(yb, ye + 1):
            for gx in range(xb, xe + 1):
                total += thermal_grid[gy][gx]
                count += 1
        avg = total / count if count > 0 else 0
        return min_v <= avg <= max_v

    # ─── Master Evaluator ─────────────────────────────────────────────────────
    def _evaluate(self):
        self.logs.clear()
        violations = []

        is_incomplete = len(self.unplaced_items) > 0
        if is_incomplete:
            self.success = False
            self.logs.append(f"Incomplete: {len(self.unplaced_items)} item(s) unplaced: {list(self.unplaced_items.keys())}")

        grid    = self._build_grid()
        self.solar_grid   = self._build_solar_grid()
        self.thermal_grid = self._build_thermal_grid()
        visited = self._bfs_reachable(grid)
        rw, rh  = self.room_size

        # 1. Human Flow
        for iid in self.placed_items:
            neighbour_reachable = any(
                0 <= nx < rw and 0 <= ny < rh and visited[ny][nx] and not grid[ny][nx]
                for (nx, ny) in self._all_neighbors(iid)
            )
            if not neighbour_reachable:
                violations.append(f"FLOW ERROR: Human path to '{iid}' is completely blocked.")

        # 2. Per-item architectural rules
        for iid, placed in self.placed_items.items():
            item_def = self.item_defs[iid]
            for rule in item_def.rules:
                err = None
                if rule == "must_touch_wall":
                    if not self._check_must_touch_wall(iid):
                        err = f"DESIGN ERROR [{iid.upper()}]: Must be flush against a wall."
                elif rule == "near_window":
                    if not self._check_near_window(iid):
                        err = f"DESIGN ERROR [{iid.upper()}]: Must be within {NEAR_WINDOW_RADIUS} tiles of a window."
                elif rule == "near_door":
                    if not self._check_near_door(iid):
                        err = f"DESIGN ERROR [{iid.upper()}]: Must be within {NEAR_DOOR_RADIUS} tiles of the door."
                elif rule == "all_sides_clear":
                    if not self._check_all_sides_clear(iid, grid):
                        err = f"DESIGN ERROR [{iid.upper()}]: Must have free space on all 4 sides."
                elif rule == "needs_sunlight":
                    if not self._check_needs_sunlight(iid, self.solar_grid):
                        err = f"DESIGN ERROR [{iid.upper()}]: Blocked from natural sunlight by a tall object."
                elif rule.startswith("thermal_comfort:"):
                    parts = rule.split(":")
                    if not self._check_thermal_comfort(iid, float(parts[1]), float(parts[2]), self.thermal_grid):
                        err = f"DESIGN ERROR [{iid.upper()}]: Outside thermal comfort limits ({parts[1]} to {parts[2]})."
                elif rule.startswith("must_face:"):
                    target_id = rule.split(":")[1]
                    if not self._check_must_face(iid, target_id):
                        err = f"DESIGN ERROR [{iid.upper()}]: Must align axis with {target_id.upper()}."
                elif rule.startswith("adjacent_to:"):
                    target_id = rule.split(":")[1]
                    if not self._check_adjacent_to(iid, target_id):
                        err = f"DESIGN ERROR [{iid.upper()}]: Must be placed adjacent to {target_id.upper()}."
                if err:
                    violations.append(err)

        if violations:
            self.logs.extend(violations)
            self.success = False
        elif not is_incomplete:
            self.logs = ["SUCCESS: All items placed. Human flow verified. All architectural constraints satisfied!"]
            self.success = True
            
        return len(violations)

    # ─── Step ─────────────────────────────────────────────────────────────────
    def step(self, action: Action) -> Tuple[Observation, float, bool, dict]:
        self.step_count += 1
        
        # Dense reward variables
        reward = 0.0
        
        # Calculate pre-action violations (to reward fixing errors)
        pre_violations = self._evaluate() if self.placed_items else 0

        if action.action_type == "place_item" and action.item:
            if action.item in self.unplaced_items:
                item_def = self.unplaced_items[action.item]
            elif action.item in self.placed_items:
                item_def = self.item_defs[action.item]
                del self.placed_items[action.item]
            else:
                self.logs = [f"ERROR: '{action.item}' not found in inventory."]
                return self._get_obs(), -0.5, False, {}

            xb, xe, yb, ye = self._bbox_action(action, item_def)
            rw, rh = self.room_size

            if xb < 0 or xe >= rw or yb < 0 or ye >= rh:
                self.logs = [f"BOUNDS: '{action.item}' at ({action.x},{action.y}) is outside the room ({rw}x{rh})."]
                self.unplaced_items[action.item] = item_def
                reward = -1.0
            else:
                overlap_id = None
                for oid, op in self.placed_items.items():
                    oxb, oxe, oyb, oye = self._bbox(op, self.item_defs[oid])
                    if not (xe < oxb or xb > oxe or ye < oyb or yb > oye):
                        overlap_id = oid
                        break

                if overlap_id:
                    self.logs = [f"COLLISION: '{action.item}' overlaps with '{overlap_id}'."]
                    self.unplaced_items[action.item] = item_def
                    reward = -1.0
                else:
                    if action.item in self.unplaced_items:
                        del self.unplaced_items[action.item]
                    self.placed_items[action.item] = PlacedFurniture(
                        id=action.item, x=action.x, y=action.y, rotation=action.rotation
                    )
                    
                    # Dense partial progress reward logic
                    post_violations = self._evaluate()
                    if post_violations < pre_violations:
                        reward = 1.0  # Fixed a violation!
                    elif post_violations == 0 and len(self.unplaced_items) < len(self.item_defs):
                        reward = 0.5  # Placed an item cleanly, no violations!
                    else:
                        reward = -0.2 * post_violations  # Penalize the current active violations

        elif action.action_type == "remove_item" and action.item in self.placed_items:
            self.unplaced_items[action.item] = self.item_defs[action.item]
            del self.placed_items[action.item]
            self._evaluate()
            reward = -0.1  # Slight penalty for backtracking

        done = self.success or self.step_count >= self.max_steps
        if self.success:
            reward = 10.0  # Big sparse bonus for completing the level
            
        return self._get_obs(), reward, done, {}

    def _get_obs(self) -> Observation:
        return Observation(
            room_size=self.room_size,
            door_position=self.door_position,
            windows=self.windows,
            unplaced_items=list(self.unplaced_items.values()),
            placed_items=copy.deepcopy(self.placed_items),
            solar_grid=copy.deepcopy(self.solar_grid),
            thermal_grid=copy.deepcopy(self.thermal_grid),
            feedback_logs=copy.deepcopy(self.logs),
            step_count=self.step_count,
            max_steps=self.max_steps,
        )

    def get_episode_result(self) -> EpisodeResult:
        # Meaningful score mapping (0.0 to 1.0)
        # 0.0 - 0.5: Partial progress based on placed valid items.
        # 0.5 - 1.0: Success efficiency.
        
        total_items = len(self.item_defs)
        placed_ratio = len(self.placed_items) / total_items if total_items > 0 else 0
        
        # Penalize for active violations
        active_violations = len([l for l in self.logs if "ERROR" in l])
        valid_placement_score = max(0.0, placed_ratio * 0.5 - (active_violations * 0.05))
        
        if self.success:
            # The baseline agent places everything perfectly in exactly `total_items` steps.
            wasted_steps = max(0, self.step_count - total_items)
            efficiency_bonus = max(0.0, 0.5 - (wasted_steps * 0.05))
            score = 0.5 + efficiency_bonus
        else:
            score = valid_placement_score

        # OpenEnv Validator Requirement: Strictly between (0.0, 1.0)
        final_score = min(0.99, max(0.01, score))

        return EpisodeResult(
            task_id=self.task_id,
            success=self.success,
            score=final_score,
            steps_taken=self.step_count,
        )