# ============================
# scene.py — Structural State
# ============================
from __future__ import annotations
import numpy as np
from typing import List, Tuple, Optional, Sequence, Dict
from collections import Counter
from .constants import CPK_COLORS, ATOMIC_RADII, DEFAULT_COLOR, DEFAULT_RADIUS
from .analysis.geometry import compute_bounds, cell_corners

class Scene:
    """
    Manages the current structural state, including atom positions, 
    lattice, selection, and element-based properties.
    """
    def __init__(self):
        self.N = 0
        self.pos = np.array([], dtype=np.float32)
        self.elements: List[str] = []
        self.radii = np.array([], dtype=np.float32)
        self.colors = np.array([], dtype=np.float32)
        
        # Unit Cell
        self.lat = np.eye(3, dtype=np.float32)
        self.is_periodic = False
        
        # Metadata
        self.composition_str = "N/A"
        self.energy = 0.0
        
        # Selections
        self.selected_atoms: List[int] = []
        self.box_selecting = False
        self.box_start = None
        self.box_end = None
        
        # Bonds (Adjacency)
        self.bonds: List[Tuple[int, int]] = []
        self.bond_count = 0
        
        # Replicas (Supercells)
        self.replicas = (0, 0, 0)
        self.replica_offsets = np.zeros((1, 3), dtype=np.float32)

    def ingest(self, positions: np.ndarray, lattice: np.ndarray, energy: float, 
               elements: Optional[List[str]] = None, 
               colors: Optional[np.ndarray] = None, 
               radii: Optional[np.ndarray] = None):
        """Standardizes and loads new atomic data into the scene."""
        self.pos = np.asarray(positions, dtype=np.float32).reshape(-1, 3)
        self.N = self.pos.shape[0]
        self.energy = energy
        
        try:
            self.lat = np.asarray(lattice, dtype=np.float32).reshape(3, 3)
            self.is_periodic = (np.linalg.norm(self.lat) > 1e-4) # Simple non-zero check
        except Exception:
            self.lat = np.eye(3, dtype=np.float32)
            self.is_periodic = False
            
        self.elements = elements if elements else []
        
        # Update Composition string
        if self.elements:
            counts = Counter(self.elements)
            self.composition_str = " ".join([f"{k}{v}" for k, v in sorted(counts.items())])
        else:
            self.composition_str = "Unknown"

        # Apply default coloring/radii if not provided
        if radii is None:
            if self.elements:
                self.radii = np.array([ATOMIC_RADII.get(e, DEFAULT_RADIUS) for e in self.elements], dtype=np.float32)
            else:
                self.radii = np.full(self.N, DEFAULT_RADIUS, dtype=np.float32)
        else:
            self.radii = np.asarray(radii, dtype=np.float32)
            
        if colors is None:
            if self.elements:
                self.colors = np.array([CPK_COLORS.get(e, DEFAULT_COLOR) for e in self.elements], dtype=np.float32)
            else:
                self.colors = np.tile(DEFAULT_COLOR, (self.N, 1)).astype(np.float32)
        else:
            self.colors = np.asarray(colors, dtype=np.float32)
            
        # Re-calc bounding boxes for initial camera positioning helpers (exterior logic)
        return self.get_full_bounds()

    def get_full_bounds(self) -> Tuple[np.ndarray, np.ndarray]:
        """Returns the AABB (min, max) including atoms and the unit cell corners."""
        if self.N == 0:
            return np.array([-5, -5, -5]), np.array([5, 5, 5])
        
        all_pts = [self.pos]
        if self.is_periodic:
            all_pts.append(cell_corners(self.lat))
        
        full_cloud = np.vstack(all_pts)
        return compute_bounds(full_cloud)

    @property
    def unique_elements(self) -> List[str]:
        """Returns sorted list of unique chemical species present in the scene."""
        if not self.elements: return []
        return sorted(list(set(self.elements)))

    def set_replicas(self, rx: int, ry: int, rz: int):
        """Updates replica grid and pre-calculates the offset matrix."""
        self.replicas = (int(rx), int(ry), int(rz))
        a, b, c = self.lat[0], self.lat[1], self.lat[2]
        offs = [i*a + j*b + k*c for i in range(-rx, rx+1) for j in range(-ry, ry+1) for k in range(-rz, rz+1)]
        self.replica_offsets = np.array(offs, dtype=np.float32)
        return self.replica_offsets

    def clear_selection(self):
        self.selected_atoms = []
        self.box_selecting = False
