%% FEM vs analytical solution for -(lambda(x) u')' = 0 with Dirichlet BC
clear; clc;

%% Parameters
n = 20;                 % number of elements (increase for nicer plot)
h = 1/n;
x = linspace(0,1,n+1)'; % node coordinates

% Jump location
x0 = 1/sqrt(2);

% Piecewise lambda function
lambda_fun = @(x) (x < x0).*1 + (x >= x0).*10;

%% Assemble FEM matrix and RHS (lifting via u_D(x)=x)
A = zeros(n+1,n+1);
F = zeros(n+1,1);

% Derivatives of linear shape functions are constant on each element
phi_prime = [-1/h; 1/h];

for e = 1:n
    % element nodes
    nodes = [e, e+1];
    xL = x(nodes(1)); 
    xR = x(nodes(2));

    % CASE A: no jump inside the element
    if (xR <= x0) || (xL >= x0)
        lambda_e = lambda_fun((xL + xR)/2);  % constant lambda
        A_e = lambda_e * (1/h) * [1 -1; -1 1];
        % load due to u_D'(x)=1
        F_e = [ lambda_e; -lambda_e ];

    % CASE B: jump inside the element -> split integrals
    else
        h1 = x0 - xL;
        h2 = xR - x0;

        % Local stiffness contributions
        A_e = [ (1*h1 + 10*h2)/h^2 , -(1*h1 + 10*h2)/h^2 ;
               -(1*h1 + 10*h2)/h^2 ,  (1*h1 + 10*h2)/h^2 ];

        % load vector components
        F_e = [ (h1 + 10*h2)/h ;
               -(h1 + 10*h2)/h ];
    end

    % Assemble
    A(nodes,nodes) = A(nodes,nodes) + A_e;
    F(nodes) = F(nodes) + F_e;
end

%% Apply Dirichlet BC: u(0)=0, u(1)=1
% Dirichlet BC: w(0)=0, w(1)=0
A(1,:) = 0; A(:,1) = 0; A(1,1) = 1;
A(end,:) = 0; A(:,end) = 0; A(end,end) = 1;
F(1) = 0;
F(end) = 0;

%% Solve FEM
w = A \ F;
u_fem = x + w;

%% Analytical solution with piecewise λ(x)
x0 = 1/sqrt(2);
J = 1 / (x0 + (1-x0)/10);

u_exact = @(x) (x <= x0) .* (J*x) + ...
               (x >  x0) .* ( J*x0 + (J/10)*(x - x0) );
u_an = u_exact(x);

%% Plot
plot(x,u_fem,'o-','LineWidth',2); hold on;
plot(x,u_an,'r--','LineWidth',2);
legend('FEM','Analytical','Location','Best');
xlabel('x'); ylabel('u(x)');
title('FEM vs Analytical Solution for (B)');
grid on;