import numpy as np

class OpticalFDTD:
    """
    2D FDTD Solver (TMz Mode)
    Simulates E_z, H_x, H_y fields.
    """
    def __init__(self, size_x=100, size_y=100, steps=200):
        self.size_x = size_x
        self.size_y = size_y
        self.steps = steps

        # Fields
        self.Ez = np.zeros((size_x, size_y))
        self.Hx = np.zeros((size_x, size_y))
        self.Hy = np.zeros((size_x, size_y))

        # Materials (Permittivity)
        self.epsilon = np.ones((size_x, size_y))

    def solve(self, source_pos=(25, 25)):
        """
        Run simulation with a point source.
        """
        for t in range(self.steps):
            # Update Magnetic Fields
            self.Hx[:, :-1] -= 0.5 * (self.Ez[:, 1:] - self.Ez[:, :-1])
            self.Hy[:-1, :] += 0.5 * (self.Ez[1:, :] - self.Ez[:-1, :])

            # Update Electric Field
            self.Ez[1:-1, 1:-1] += (0.5 / self.epsilon[1:-1, 1:-1]) * (
                (self.Hy[1:-1, 1:-1] - self.Hy[:-2, 1:-1]) -
                (self.Hx[1:-1, 1:-1] - self.Hx[1:-1, :-2])
            )

            # Source Injection (Gaussian Pulse)
            pulse = np.exp(-((t - 30) ** 2) / 100)
            self.Ez[source_pos] += pulse

        return self.Ez

def simulate_wave_propagation(size=100):
    fdtd = OpticalFDTD(size, size)
    field = fdtd.solve()
    return np.sum(np.abs(field)) # Total energy metric
