embodied/interactive_objects.py

"""
InteractiveObjects — fyzické předměty pro haptické a kognitivní úlohy.

BallObject       – gumový míč s vysokou elasticitou
WaterCupObject   – hrnek s hladinou vody (kloní se = přelévá)
RubiksCubeObject – zjednodušená Rubikova kostka (fyzické těleso + logický stav)
ShapeSorterWall  – zeď s otvory (čtverec, trojúhelník, hvězda) + odpovídající tvar
"""
from __future__ import annotations
import sys, os
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))

import numpy as np
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Tuple

from engine.physics.rigid_body import (
    PhysicsWorld, RigidBody, BodyType,
    vec3, Vec3, length, normalize,
)


# ── Základní třída ────────────────────────────────────────────────────────────

class PhysicsObject:
    """Základní třída pro interaktivní objekty."""

    def __init__(self, phys: PhysicsWorld, obj_id: str):
        self.phys   = phys
        self.obj_id = obj_id
        self.body: Optional[RigidBody] = None

    @property
    def position(self) -> Vec3:
        return self.body.position if self.body else vec3()

    @property
    def velocity(self) -> Vec3:
        return self.body.velocity if self.body else vec3()

    def distance_to(self, pos: Vec3) -> float:
        return float(length(self.position - pos))


# ── Míč ───────────────────────────────────────────────────────────────────────

class BallObject(PhysicsObject):
    """
    Gumový míč — malé těleso s vysokou elasticitou a nízkou hmotností.
    Odraz: restitution=0.82 (≈ gumový míč, ne super-ball).
    """
    RADIUS   = 0.12   # m
    MASS     = 0.25   # kg

    def __init__(self, phys: PhysicsWorld, pos: Vec3, obj_id: str = "ball"):
        super().__init__(phys, obj_id)
        self.body = phys.create_body(
            obj_id, BodyType.DYNAMIC,
            position=pos,
            size=vec3(self.RADIUS * 2, self.RADIUS * 2, self.RADIUS * 2),
            mass=self.MASS,
            restitution=0.82,
            friction=0.35,
        )
        self._spawn_pos = pos

    def reset(self, pos: Optional[Vec3] = None):
        p = pos if pos is not None else self._spawn_pos
        self.body.position = p
        self.body.velocity = vec3()
        self.body.angular_velocity = vec3()
        self.body.rotation = vec3()

    def get_state(self) -> np.ndarray:
        """6 floats: [pos_x, pos_y, pos_z, vel_x, vel_y, vel_z]"""
        p = self.body.position
        v = self.body.velocity
        return np.array([p[0], p[1], p[2], v[0], v[1], v[2]], dtype=np.float32)


# ── Hrnek s vodou ─────────────────────────────────────────────────────────────

class WaterCupObject(PhysicsObject):
    """
    Keramický hrnek s hladinou vody.
    water_level klesá, když je hrnek nakloněný (tilt > threshold).
    Použití: agent se musí naučit nosit hrnek opatrně.
    """
    MASS               = 0.35
    SIZE               = vec3(0.09, 0.12, 0.09)
    TILT_SPILL_THRESH  = 0.40  # rad od svislé osy
    SPILL_RATE         = 0.03  # úbytek hladiny za krok při náklonu

    def __init__(self, phys: PhysicsWorld, pos: Vec3, obj_id: str = "water_cup"):
        super().__init__(phys, obj_id)
        self.body = phys.create_body(
            obj_id, BodyType.DYNAMIC,
            position=pos,
            size=self.SIZE,
            mass=self.MASS,
            restitution=0.08,
            friction=0.70,
        )
        self._spawn_pos   = pos
        self.water_level  = 1.0   # 0.0 = prázdný, 1.0 = plný
        self._spilled     = False

    def step(self):
        """Volat každý fyzikální krok."""
        if self.water_level <= 0.0:
            return
        rot = self.body.rotation
        tilt = float(np.sqrt(rot[0] ** 2 + rot[2] ** 2))
        if tilt > self.TILT_SPILL_THRESH:
            self.water_level = max(0.0, self.water_level - self.SPILL_RATE)
            if self.water_level == 0.0:
                self._spilled = True

    def reset(self, pos: Optional[Vec3] = None):
        p = pos if pos is not None else self._spawn_pos
        self.body.position = p
        self.body.velocity = vec3()
        self.body.angular_velocity = vec3()
        self.body.rotation = vec3()
        self.water_level = 1.0
        self._spilled = False

    def get_state(self) -> np.ndarray:
        """4 floats: [pos_y_rel_spawn, tilt, water_level, spilled]"""
        rot   = self.body.rotation
        tilt  = float(np.sqrt(float(rot[0]) ** 2 + float(rot[2]) ** 2))
        pos_y = float(self.body.position[1])
        return np.array([pos_y, tilt, self.water_level, float(self._spilled)],
                        dtype=np.float32)


# ── Rubikova kostka ───────────────────────────────────────────────────────────

# Každá stěna má 9 políček → 6 stěn × 9 = 54 hodnot (0-5 = 6 barev)
# Zjednodušení: jedna fyzická krychle, logický stav jako int array
_SOLVED_FACES = np.array([
    [i] * 9 for i in range(6)   # každá stěna má jednu barvu
], dtype=np.int8).flatten()


class RubiksCubeObject(PhysicsObject):
    """
    Fyzická krychle s logickým stavem 6×9 políček (6 barev).
    Rotace stěn: rotate_face(face_idx, clockwise).
    Stav: stupeň zamíchanosti (0=vyřešeno, 1=přeskupeno).
    """
    SIZE = vec3(0.12, 0.12, 0.12)
    MASS = 0.15

    # Indexy stěn
    FRONT, BACK, LEFT, RIGHT, TOP, BOTTOM = 0, 1, 2, 3, 4, 5

    # Otáčení: stěna 0-5 → seznam indexů políček (zjednodušeno — ring rotation)
    _FACE_RING = {
        0: ([6,7,8],[0,3,6],[2,1,0],[8,5,2],  [BACK:=1,LEFT:=2,TOP:=4,RIGHT:=3]),
    }

    def __init__(self, phys: PhysicsWorld, pos: Vec3, obj_id: str = "rubiks"):
        super().__init__(phys, obj_id)
        self.body = phys.create_body(
            obj_id, BodyType.DYNAMIC,
            position=pos,
            size=self.SIZE,
            mass=self.MASS,
            restitution=0.10,
            friction=0.80,
        )
        self._spawn_pos  = pos
        self.faces       = _SOLVED_FACES.copy()    # 54 políček
        self._moves      = 0

    def scramble(self, n_moves: int = 20, rng: np.random.Generator = None):
        if rng is None:
            rng = np.random.default_rng()
        for _ in range(n_moves):
            face = int(rng.integers(0, 6))
            self._rotate_face_logical(face, bool(rng.integers(0, 2)))
        self._moves = 0

    def _rotate_face_logical(self, face: int, clockwise: bool):
        """Otočí políčka na jedné stěně o 90°."""
        f = self.faces[face * 9:(face + 1) * 9].copy().reshape(3, 3)
        if clockwise:
            f = np.rot90(f, k=3)
        else:
            f = np.rot90(f, k=1)
        self.faces[face * 9:(face + 1) * 9] = f.flatten()
        self._moves += 1

    def is_solved(self) -> bool:
        return np.array_equal(self.faces, _SOLVED_FACES)

    def solve_fraction(self) -> float:
        """Poměr správně umístěných políček (0.0–1.0)."""
        return float(np.mean(self.faces == _SOLVED_FACES))

    def reset(self, pos: Optional[Vec3] = None, scramble_n: int = 0,
              rng: np.random.Generator = None):
        p = pos if pos is not None else self._spawn_pos
        self.body.position = p
        self.body.velocity = vec3()
        self.body.angular_velocity = vec3()
        self.body.rotation = vec3()
        self.faces = _SOLVED_FACES.copy()
        self._moves = 0
        if scramble_n > 0:
            self.scramble(scramble_n, rng)

    def get_state(self) -> np.ndarray:
        """3 floats: [solve_fraction, n_moves_norm, is_held_approx(pos_y)]"""
        return np.array([
            self.solve_fraction(),
            min(1.0, self._moves / 100.0),
            float(np.clip((float(self.body.position[1]) - 0.5) / 1.5, 0.0, 1.0)),
        ], dtype=np.float32)


# ── Třídicí zeď ───────────────────────────────────────────────────────────────

SHAPE_SQUARE   = 0
SHAPE_TRIANGLE = 1
SHAPE_STAR     = 2

# Fyzické rozměry tvarových předmětů
_SHAPE_SIZES = {
    SHAPE_SQUARE:   vec3(0.09, 0.09, 0.04),
    SHAPE_TRIANGLE: vec3(0.09, 0.09, 0.04),
    SHAPE_STAR:     vec3(0.10, 0.10, 0.04),
}

# Barvy tvarů (R, G, B) pro renderer
SHAPE_COLORS = {
    SHAPE_SQUARE:   (0.9, 0.2, 0.2),   # červená
    SHAPE_TRIANGLE: (0.2, 0.8, 0.2),   # zelená
    SHAPE_STAR:     (0.9, 0.8, 0.1),   # žlutá
}

# Pozice otvorů (x-offset od středu zdi)
_HOLE_X_OFFSETS = {
    SHAPE_SQUARE:   -0.35,
    SHAPE_TRIANGLE:  0.00,
    SHAPE_STAR:      0.35,
}

# Tolerance vložení (vzdálenost + orientace)
_INSERT_DIST_OK  = 0.15   # m — střed tvaru musí být do 15 cm od otvoru
_INSERT_TILT_OK  = 0.30   # rad — maximální náklon tvaru při vkládání


@dataclass
class ShapePiece:
    shape:   int           # SHAPE_* konstanta
    body:    RigidBody
    inserted: bool = False
    spawn_pos: Vec3 = field(default_factory=lambda: vec3())


class ShapeSorterWall:
    """
    Statická zeď s třemi otvory (čtverec, trojúhelník, hvězda).
    Tvarové předměty jsou dynamické tělesa, agent je musí vložit do správného otvoru.
    Kontrola vložení: vzdálenost od otvoru < 15 cm + přibližně správná orientace.
    """
    WALL_POS   = vec3(0.0, 1.0, -2.0)   # zeď je v záporném Z (před agentem)
    WALL_SIZE  = vec3(1.2, 1.0, 0.08)
    HOLE_Y     = 1.0   # výška středu otvorů (v absolutních souřadnicích)

    def __init__(self, phys: PhysicsWorld, origin: Vec3 = None, rng=None):
        self.phys   = phys
        self._rng   = rng if rng is not None else np.random.default_rng()
        self.origin = origin if origin is not None else vec3()
        self._wall_body: Optional[RigidBody] = None
        self.pieces: List[ShapePiece] = []
        self._spawn()

    def _spawn(self):
        wall_pos = self.origin + self.WALL_POS
        self._wall_body = self.phys.create_body(
            "shape_wall", BodyType.STATIC,
            position=wall_pos,
            size=self.WALL_SIZE,
            mass=0.0, restitution=0.05, friction=0.90,
        )

        # Spawn tvarové předměty před zdi
        spawn_z  = float(self.WALL_POS[2]) + 1.2   # 1.2 m před zdi
        shapes   = [SHAPE_SQUARE, SHAPE_TRIANGLE, SHAPE_STAR]
        # Zamíchej pořadí spawnu
        order    = self._rng.permutation(len(shapes))
        x_slots  = [-0.35, 0.0, 0.35]
        for i, shape_idx in enumerate(order):
            shape = shapes[shape_idx]
            pos   = self.origin + vec3(
                x_slots[i],
                0.80,       # výška stolu/podlahy
                spawn_z,
            )
            body = self.phys.create_body(
                f"shape_{shape}_{i}", BodyType.DYNAMIC,
                position=pos,
                size=_SHAPE_SIZES[shape],
                mass=0.08, restitution=0.10, friction=0.75,
            )
            self.pieces.append(ShapePiece(shape=shape, body=body, spawn_pos=pos.copy()))

    def hole_position(self, shape: int) -> Vec3:
        """Absolutní pozice středu otvoru pro daný tvar."""
        wx = float(self.WALL_POS[0]) + float(self.origin[0])
        wz = float(self.WALL_POS[2]) + float(self.origin[2]) + 0.05
        return vec3(wx + _HOLE_X_OFFSETS[shape], self.HOLE_Y, wz)

    def check_insertions(self) -> int:
        """Zkontroluje vložení předmětů; vrací počet správně vložených."""
        count = 0
        for piece in self.pieces:
            if piece.inserted:
                count += 1
                continue
            hole = self.hole_position(piece.shape)
            dist = float(length(piece.body.position - hole))
            rot  = piece.body.rotation
            tilt = float(np.sqrt(float(rot[0]) ** 2 + float(rot[1]) ** 2))
            if dist < _INSERT_DIST_OK and tilt < _INSERT_TILT_OK:
                piece.inserted = True
                # Zafixuj předmět na místě
                piece.body.body_type = BodyType.STATIC
                piece.body.velocity  = vec3()
                count += 1
        return count

    def all_inserted(self) -> bool:
        return all(p.inserted for p in self.pieces)

    def nearest_piece(self, hand_pos: Vec3) -> Optional[ShapePiece]:
        """Vrátí nejbližší nevložený předmět k ruce agenta."""
        best_d = float("inf")
        best_p = None
        for p in self.pieces:
            if p.inserted:
                continue
            d = float(length(p.body.position - hand_pos))
            if d < best_d:
                best_d, best_p = d, p
        return best_p

    def reset(self):
        for piece in self.pieces:
            piece.inserted = False
            piece.body.body_type = BodyType.DYNAMIC
            piece.body.position  = piece.spawn_pos.copy()
            piece.body.velocity  = vec3()
            piece.body.angular_velocity = vec3()
            piece.body.rotation  = vec3()

    def get_state(self) -> np.ndarray:
        """
        9 floats: pro každý předmět (3×):
          [dist_to_hole_norm, inserted, piece_pos_y_norm]
        """
        out = []
        for piece in self.pieces:
            hole = self.hole_position(piece.shape)
            dist = float(length(piece.body.position - hole))
            out.extend([
                float(np.clip(dist / 3.0, 0.0, 1.0)),
                float(piece.inserted),
                float(np.clip(float(piece.body.position[1]) / 2.0, 0.0, 1.0)),
            ])
        return np.array(out, dtype=np.float32)