"""
volumetric.py — Optimized Gaussian density mapper and volumetric data loaders.
"""
from __future__ import annotations
import numpy as np
from dataclasses import dataclass
from typing import Tuple, Optional

@dataclass
class VolumetricData:
    grid: np.ndarray    # 3D float32 array
    origin: np.ndarray  # (3,) origin in Å
    spacing: np.ndarray # (3,) dx, dy, dz spacing in Å
    shape: Tuple[int, int, int]

def generate_density_from_atoms(pos: np.ndarray, 
                                radii: np.ndarray, 
                                voxels_per_a: float = 10.0,
                                sigma: float = 0.5) -> VolumetricData:
    """
    Highly optimized O(N) Gaussian splatting for atomic density mapping.
    Maps atoms into a localized voxel grid to avoid N^3 complexity.
    """
    if len(pos) == 0:
        return VolumetricData(np.zeros((1,1,1), dtype=np.float32), 
                              np.zeros(3), np.ones(3)*0.1, (1,1,1))

    # 1. Coordinate Space
    mn, mx = pos.min(axis=0) - 4.0, pos.max(axis=0) + 4.0
    dims = np.ceil((mx - mn) * voxels_per_a).astype(int)
    nx, ny, nz = dims
    spacing = (mx - mn) / (dims.astype(np.float32) + 1e-9)
    origin = mn

    grid = np.zeros((nx, ny, nz), dtype=np.float32)

    # 2. Gaussian Splatting (Localized)
    # rho = exp(-0.5 * (r/sigma)^2)
    cut = 3.0 * sigma
    cut_vox = np.ceil(cut / spacing).astype(int)

    for i in range(len(pos)):
        p = pos[i]
        # Map center to grid index
        idx = ((p - origin) / spacing).astype(int)
        
        # Local bounding box in voxels
        x0, x1 = max(0, idx[0]-cut_vox[0]), min(nx, idx[0]+cut_vox[0]+1)
        y0, y1 = max(0, idx[1]-cut_vox[1]), min(ny, idx[1]+cut_vox[1]+1)
        z0, z1 = max(0, idx[2]-cut_vox[2]), min(nz, idx[2]+cut_vox[2]+1)

        if x1 <= x0 or y1 <= y0 or z1 <= z0: continue

        # Local coordinate grids
        gx = np.arange(x0, x1) * spacing[0] + origin[0] - p[0]
        gy = np.arange(y0, y1) * spacing[1] + origin[1] - p[1]
        gz = np.arange(z0, z1) * spacing[2] + origin[2] - p[2]

        X, Y, Z = np.meshgrid(gx, gy, gz, indexing='ij')
        R2 = X*X + Y*Y + Z*Z
        
        weight = np.exp(-0.5 * R2 / (sigma**2))
        grid[x0:x1, y0:y1, z0:z1] += weight.astype(np.float32)

    # Normalize roughly to 0..1
    g_max = grid.max()
    if g_max > 0:
        grid /= g_max

    return VolumetricData(grid, origin, spacing, (nx, ny, nz))

def load_cube_file(path: str) -> Optional[VolumetricData]:
    """Basic Gaussian Cube file loader."""
    try:
        with open(path, 'r') as f:
            lines = f.readlines()
        
        natoms = int(lines[2].split()[0])
        origin = np.array(lines[3].split()[1:4], dtype=np.float32)
        
        nx, dx, dy, dz = lines[4].split()
        nx = int(nx)
        spacing = np.array([dx, dy, dz], dtype=np.float32) # Simplification for orthogonal
        
        ny = int(lines[5].split()[0])
        nz = int(lines[6].split()[0])
        
        # Skip atom lines
        data_start = 7 + abs(natoms)
        flat_data = []
        for l in lines[data_start:]:
            flat_data.extend([float(x) for x in l.split()])
        
        grid = np.array(flat_data, dtype=np.float32).reshape((nx, ny, nz))
        return VolumetricData(grid, origin, spacing, (nx, ny, nz))
    except Exception:
        return None
