% 2D Acoustic Wave Equation with Variable Coefficients
% Non-homogeneous media with circular obstacle

clear all; close all; clc;

% Simulation parameters
Lx = 10;              % Domain size in x (meters)
Ly = 10;              % Domain size in y (meters)
Nx = 400;             % Number of grid points in x
Ny = 400;             % Number of grid points in y
dx = Lx / (Nx - 1);   % Spatial step in x
dy = Ly / (Ny - 1);   % Spatial step in y
T = 20;               % Total simulation time (seconds)

% Create spatial grid
x = linspace(0, Lx, Nx);
y = linspace(0, Ly, Ny);
[X, Y] = meshgrid(x, y);

% Define spatially varying wave speed c(x,y)
% Create heterogeneous media with different layers and regions
c = ones(Ny, Nx) * 1.0;  % Base wave speed (m/s)

% Add horizontal layers with different speeds
c(Y < 2.5) = 1.5;          % Fast layer at bottom
c(Y > 7.5) = 0.8;          % Slow layer at top
c(Y >= 4 & Y <= 6) = 1.2;  % Medium layer in middle

% Add some random heterogeneity
rng(42);  % For reproducibility
noise = 0.1 * randn(Ny, Nx);
c = c .* (1 + noise);

% Add a high-speed region (lens effect)
lens_x = 7;
lens_y = 3;
lens_radius = 1.0;
lens_mask = sqrt((X - lens_x).^2 + (Y - lens_y).^2) <= lens_radius;
c(lens_mask) = 2.0;

% Create circular obstacle (rigid boundary)
obstacle_center_x = Lx / 2;
obstacle_center_y = Ly / 2;
obstacle_radius = 1.0;
obstacle = sqrt((X - obstacle_center_x).^2 + (Y - obstacle_center_y).^2) <= obstacle_radius;

% Set wave speed to zero in obstacle (will be handled separately)
c_original = c;
c(obstacle) = 0;

% Calculate time step based on minimum wave speed and CFL condition
c_max = max(c(:));
dt = 0.4 * min(dx, dy) / c_max;  % CFL condition
Nt = round(T / dt);

% Calculate c² and its gradients for variable coefficient PDE
c2 = c.^2;
[c2_x, c2_y] = gradient(c2, dx, dy);

% Initialize wave field (pressure)
p = zeros(Ny, Nx);        % Current time step
p_old = zeros(Ny, Nx);    % Previous time step
p_new = zeros(Ny, Nx);    % Next time step

% Initial condition: Gaussian pulse
source_x = 2.0;
source_y = Ly / 2;
sigma = 0.3;
p = exp(-((X - source_x).^2 + (Y - source_y).^2) / (2 * sigma^2));
p(obstacle) = 0;
p_old = p;

% Set up figure for animation
figure('Position', [100, 100, 1200, 500]);

% Time-stepping loop using variable coefficient wave equation
% ∂²p/∂t² = ∇·(c²∇p) = c²∇²p + ∇c²·∇p
for n = 1:Nt
    % Calculate spatial derivatives
    [p_x, p_y] = gradient(p, dx, dy);

    % Calculate Laplacian
    p_xx = zeros(Ny, Nx);
    p_yy = zeros(Ny, Nx);

    for i = 2:Ny-1
        for j = 2:Nx-1
            if ~obstacle(i, j)
                % Second derivatives
                p_xx(i, j) = (p(i, j+1) - 2*p(i, j) + p(i, j-1)) / dx^2;
                p_yy(i, j) = (p(i+1, j) - 2*p(i, j) + p(i-1, j)) / dy^2;

                % Variable coefficient wave equation
                % ∂²p/∂t² = c²(x,y)∇²p + ∂c²/∂x·∂p/∂x + ∂c²/∂y·∂p/∂y
                laplacian = p_xx(i, j) + p_yy(i, j);
                gradient_term = c2_x(i, j) * p_x(i, j) + c2_y(i, j) * p_y(i, j);

                % Time integration (leap-frog scheme)
                p_new(i, j) = 2*p(i, j) - p_old(i, j) + ...
                              dt^2 * (c2(i, j) * laplacian + gradient_term);
            end
        end
    end

    % Apply absorbing boundary conditions (first-order Mur ABC)
    % Left boundary
    p_new(:, 1) = p(:, 2) + ((c_max*dt - dx)/(c_max*dt + dx)) * (p_new(:, 2) - p(:, 1));
    % Right boundary
    p_new(:, Nx) = p(:, Nx-1) + ((c_max*dt - dx)/(c_max*dt + dx)) * (p_new(:, Nx-1) - p(:, Nx));
    % Bottom boundary
    p_new(1, :) = p(2, :) + ((c_max*dt - dy)/(c_max*dt + dy)) * (p_new(2, :) - p(1, :));
    % Top boundary
    p_new(Ny, :) = p(Ny-1, :) + ((c_max*dt - dy)/(c_max*dt + dy)) * (p_new(Ny-1, :) - p(Ny, :));

    % Apply rigid obstacle boundary
    p_new(obstacle) = 0;

    % Update for next time step
    p_old = p;
    p = p_new;

    % Plot every 15 time steps
    if mod(n, 15) == 0
        % Left subplot: Wave field
        subplot(1, 2, 1);
        surf(X, Y, p, 'EdgeColor', 'none');
        view(2);
        shading interp;
        axis equal tight;
        caxis([-0.3 0.3]);
        colormap(gca, jet);

        % Draw obstacle boundary
        hold on;
        theta = linspace(0, 2*pi, 100);
        plot3(obstacle_center_x + obstacle_radius*cos(theta), ...
              obstacle_center_y + obstacle_radius*sin(theta), ...
              ones(size(theta))*0.4, 'k-', 'LineWidth', 2);
        hold off;

        title(sprintf('Wave Field (t = %.2f s)', n*dt));
        xlabel('x (m)');
        ylabel('y (m)');
        colorbar;

        % Right subplot: Medium properties
        subplot(1, 2, 2);
        imagesc(x, y, c_original);
        axis equal tight;
        colormap(gca, hot);
        set(gca, 'YDir', 'normal');

        % Draw obstacle on medium plot
        hold on;
        viscircles([obstacle_center_x, obstacle_center_y], obstacle_radius, ...
                   'Color', 'b', 'LineWidth', 2);
        viscircles([lens_x, lens_y], lens_radius, ...
                   'Color', 'g', 'LineWidth', 1.5, 'LineStyle', '--');
        hold off;

        title('Wave Speed c(x,y)');
        xlabel('x (m)');
        ylabel('y (m)');
        colorbar;
        caxis([0.5 2.5]);

        drawnow;
    end
end

disp('Simulation complete!');
disp(['Total time steps: ', num2str(Nt)]);
disp(['Time step dt: ', num2str(dt), ' s']);
