import numpy as np
from beamz.const import EPS_0

class GaussianSource:
    """Gaussian spatial source for FDTD simulations.
    
    Injects a Gaussian spatial profile into the Ez field (and other E components in 3D).
    Useful for dipole-like excitations.
    """
    
    def __init__(self, position, width, signal):
        """Initialize the Gaussian source.
        
        Args:
            position: (x, y) for 2D or (x, y, z) for 3D - center of Gaussian
            width: Standard deviation of Gaussian profile
            signal: Time-dependent signal function s(t) or array
        """
        self.position = position
        self.width = width
        self.signal = signal
        self._spatial_profile_ez = None
        self._grid_indices = None
        
    def _get_signal_value(self, time, dt):
        """Interpolate signal value at arbitrary time."""
        # Handle array signal
        if isinstance(self.signal, (list, np.ndarray)):
            idx_float = float(time / dt)
            idx_low = int(np.floor(idx_float))
            idx_high = idx_low + 1
            frac = idx_float - idx_low
            
            if 0 <= idx_low < len(self.signal) - 1:
                return (1.0 - frac) * self.signal[idx_low] + frac * self.signal[idx_high]
            elif idx_low == len(self.signal) - 1:
                return self.signal[idx_low]
            else:
                return 0.0
        # Handle callable signal
        elif callable(self.signal):
            return self.signal(time)
        else:
            return 0.0

    def inject(self, fields, t, dt, current_step, resolution, design):
        """Inject source fields directly into the simulation grid before the FDTD update step.
        
        Args:
            fields: The Fields object containing E and H arrays
            t: Current simulation time
            dt: Time step size
            current_step: Current time step index
            resolution: Spatial resolution (dx, dy)
            design: The simulation Design object
        """
        dx = dy = resolution
        
        # Check dimensionality from fields
        if fields.permittivity.ndim == 3:
            self._inject_3d(fields, t, dt, resolution)
        else:
            self._inject_2d(fields, t, dt, resolution)
            
    def _inject_2d(self, fields, t, dt, resolution):
        """Inject into 2D grid (Ez component)."""
        ny, nx = fields.Ez.shape
        
        # Initialize spatial profile if needed (do this once)
        if self._spatial_profile_ez is None:
            x0, y0 = self.position
            
            # Determine bounding box for Gaussian (e.g., +/- 4 sigma) to save computation
            # Convert position to grid indices
            sigma_grid = self.width / resolution
            radius_grid = int(np.ceil(4 * sigma_grid))
            
            cx = int(round(x0 / resolution))
            cy = int(round(y0 / resolution))
            
            # Define ROI limits
            x_start = max(0, cx - radius_grid)
            x_end = min(nx, cx + radius_grid + 1)
            y_start = max(0, cy - radius_grid)
            y_end = min(ny, cy + radius_grid + 1)
            
            self._grid_indices = (slice(y_start, y_end), slice(x_start, x_end))
            
            # Generate coordinate grids for the ROI
            # Ez is at (i+0.5, j+0.5) * resolution? Or integer?
            # Assuming standard integer/half-integer staggering for Ez
            # Beamz convention: Ez at (i, j) corresponds to integer steps usually?
            # Let's use meshgrid relative to center
            
            # Coordinate arrays for ROI
            x_coords = (np.arange(x_start, x_end) + 0.5) * resolution
            y_coords = (np.arange(y_start, y_end) + 0.5) * resolution
            
            X, Y = np.meshgrid(x_coords, y_coords)
            
            # Compute Gaussian
            dist_sq = (X - x0)**2 + (Y - y0)**2
            profile = np.exp(-dist_sq / (2 * self.width**2))
            self._spatial_profile_ez = profile

        # Get signal value
        # Inject at t + 0.5 dt because Ez is updated at n+1 from n
        # Soft source: J_z is added. 
        # Ez_new = Ez_old + ... - dt/eps * J_z
        # We want to add to Ez.
        # Typically soft source adds to E directly or J term.
        # ModeSource adds: fields.Ez += injection
        # injection = -jz * dt / (eps * eps0)
        # Here we treat GaussianSource as a J_z source.
        
        signal_val = self._get_signal_value(t + 0.5 * dt, dt)
        
        # Get permittivity in the region
        eps_region = fields.permittivity[self._grid_indices]
        
        # Calculate injection term
        # J(x, t) = Profile(x) * s(t)
        # Update: E += -J * dt / (eps * eps0)
        # We can absorb the negative sign into the signal definition if we want E to follow signal
        # But strictly J is current.
        # Let's just add it as a "forcing function" to E.
        # If we want E ~ Signal, we might just add Profile * Signal * Scaling
        # Let's follow ModeSource physics: inject current J.
        
        term = self._spatial_profile_ez * signal_val
        injection = -term * dt / (EPS_0 * eps_region)
        
        # Inject
        fields.Ez[self._grid_indices] += injection

    def _inject_3d(self, fields, t, dt, resolution):
        """Inject into 3D grid (Ez component, could be expanded)."""
        # Similar to 2D but with Z coordinate
        # Only implementing Ez injection for dipole-like behavior along Z
        nz, ny, nx = fields.Ez.shape
        
        if self._spatial_profile_ez is None:
            x0, y0, z0 = self.position
            
            sigma_grid = self.width / resolution
            radius_grid = int(np.ceil(4 * sigma_grid))
            
            cx, cy, cz = int(round(x0/resolution)), int(round(y0/resolution)), int(round(z0/resolution))
            
            x_start, x_end = max(0, cx - radius_grid), min(nx, cx + radius_grid + 1)
            y_start, y_end = max(0, cy - radius_grid), min(ny, cy + radius_grid + 1)
            z_start, z_end = max(0, cz - radius_grid), min(nz, cz + radius_grid + 1)
            
            self._grid_indices = (slice(z_start, z_end), slice(y_start, y_end), slice(x_start, x_end))
            
            x_coords = (np.arange(x_start, x_end) + 0.5) * resolution
            y_coords = (np.arange(y_start, y_end) + 0.5) * resolution
            z_coords = (np.arange(z_start, z_end) + 0.5) * resolution # Ez staggered in Z?
            
            # Ez in 3D is at (i, j, k+0.5)? 
            # Assuming cell centers for simplicity for Gaussian "blob"
            
            Z, Y, X = np.meshgrid(z_coords, y_coords, x_coords, indexing='ij')
            
            dist_sq = (X - x0)**2 + (Y - y0)**2 + (Z - z0)**2
            self._spatial_profile_ez = np.exp(-dist_sq / (2 * self.width**2))

        signal_val = self._get_signal_value(t + 0.5 * dt, dt)
        eps_region = fields.permittivity[self._grid_indices]
        
        term = self._spatial_profile_ez * signal_val
        injection = -term * dt / (EPS_0 * eps_region)
        
        fields.Ez[self._grid_indices] += injection
