% Simple implementation of the heat equation on a square domain
clear;

domain_x = [0,1];
domain_y = [0,1];

n_x = 51;
n_y = 51;

coords_x = linspace( domain_x(1), domain_x(2), n_x);
coords_y = linspace( domain_y(1), domain_y(2), n_y);

n_vert = n_x*n_y;
vertices = zeros(n_vert,2);
k = 1;
for j = 1:n_y
    for i = 1:n_x
        coords = [coords_x(i), coords_y(j)];
        vertices(k, :) = coords;
        k = k+1;
    end
end

n_el = 2*(n_x-1)*(n_y-1);
elements = zeros(n_el, 3);
k = 1;
for j = 1:n_y-1
    for i = 1:n_x-1
        triangle_upper = [ j*n_x + i, (j-1)*n_x + i, j*n_x + i+1 ];
        triangle_lower = [ (j-1)*n_x + i+1, j*n_x + i+1, (j-1)*n_x + i ];
        elements(k, :) = triangle_upper;
        k = k+1;
        elements(k, :) = triangle_lower;
        k = k+1;
    end
end

is_boundary = zeros(n_vert,1);
bound_val = zeros(n_vert,1);
for i=1:n_x
    is_boundary(i) = 1;
    bound_val(i) = -1;
end
for j=2:n_y-1
    is_boundary(j*n_x) = 1;
    bound_val(j*n_x) = -1;
    is_boundary((j-1)*n_x+1) = 1;
    bound_val((j-1)*n_x+1) = -1;
end
for i=1:n_x
    is_boundary(n_x*(n_y-1) + i) = 1;
    bound_val(n_x*(n_y-1) + i) = -1;
end

dof = zeros(n_vert);
k = 0;
for i = 1:n_vert
    if is_boundary(i) ~= 1
        k = k+1;
        dof(i) = k;
    end
end

n_dof = nnz(dof);
E = sparse(n_dof, n_dof);
A = sparse(n_dof, n_dof);
h = zeros(n_dof,1);

gradN1 = [-1; -1]; gradN2 = [1; 0]; gradN3 = [0; 1];
ref_gradients = [gradN1, gradN2, gradN3];

eps = 0.1;

for k = 1:n_el
    x1 = vertices( elements(k,1), 1); y1 = vertices( elements(k,1), 2);
    x2 = vertices( elements(k,2), 1); y2 = vertices( elements(k,2), 2);
    x3 = vertices( elements(k,3), 1); y3 = vertices( elements(k,3), 2);
    B = [x2-x1, x3-x1; y2-y1 , y3-y1];
    size = 0.5*det(B);
    %center_x = 1/3*(x1+x2+x3); center_y = 1/3*(y1+y2+y3);
    
    for j = 1:3
        if is_boundary( elements(k, j) ) == 1, continue; end;
        
        dj = dof(elements(k,j));
        
        A(dj,dj) = A(dj,dj) + size*1/3;
        
        for i = 1:3
            
            gradi_ref = ref_gradients(:, i); gradj_ref = ref_gradients(:, j);
            
            gradi = (B')\gradi_ref; gradj = (B')\gradj_ref;
            integrand = eps*dot(gradi, gradj);
            integral = integrand*size;
            
            if is_boundary( elements(k, i) ) == 1
                c = bound_val(elements(k, i));
                h(dj) = h(dj) - c*integral;
            else
                di = dof(elements(k,i));
                E(di,dj) = E(di,dj) + integral;
            end
        end
        
        
    end
end

u = zeros(n_dof, 1);
for i = 1:n_vert
    if is_boundary(i) ~= 1
        x = vertices(i, 1); y = vertices(i, 2);
        r = sqrt((x-0.5)^2+(y-0.5)^2);
        theta = atan2(y-0.5,x-0.5);
        
        u(dof(i)) = -1;
        if r < (0.3+0.1*cos(5*theta))
            u(dof(i)) = +1;
        end
        
    end
end


T = 1;
dt = 0.00001;

n = round(T/dt);
Mov(n) = struct('cdata',[],'colormap',[]);
for t = 1:n
    
    f = zeros(n_dof, 1);
    for k = 1:n_el
        x1 = vertices( elements(k,1), 1); y1 = vertices( elements(k,1), 2);
        x2 = vertices( elements(k,2), 1); y2 = vertices( elements(k,2), 2);
        x3 = vertices( elements(k,3), 1); y3 = vertices( elements(k,3), 2);
        B = [x2-x1, x3-x1; y2-y1 , y3-y1];
        size = 0.5*det(B);
        %center_x = 1/3*(x1+x2+x3); center_y = 1/3*(y1+y2+y3);
        
        u_ctr = 0.0;
        for j = 1:3
             if is_boundary( elements(k, j) ) == 1
                 u_ctr = u_ctr + bound_val(elements(k,j));
             else
                 u_ctr = u_ctr + u(dof(elements(k,j)));
             end
        end
        u_ctr = u_ctr/3.0;
        force = -1/eps*(u_ctr^2-1) * u_ctr;
        
        for j = 1:3
            if is_boundary( elements(k, j) ) == 1, continue; end;
            dj = dof(elements(k,j));

            f(dj) = f(dj) + force*size* 1.0/3.0;
        end
    end
    
    
    b = A*u + dt/eps*(h+f);
    u_new = (dt/eps*E+A)\b;
    
    u = u_new;
    
    plot_z = zeros(n_vert,1);
    for i = 1:n_vert
        if is_boundary(i) == 1
            plot_z(i) = bound_val(i);
        else
            plot_z(i) = u(dof(i));
        end
    end
    
    trisurf(elements, vertices(:,1),vertices(:,2), plot_z);
    axis([0,1, 0,1, -1, 1]);
    %daspect manual
    %view(0,90)
    %shading interp
    drawnow;
    Mov(t) = getframe;
    
end



