% GenerateSurface.m
% Copyright (C) by Yaron Hadad 2010.
%
% Generates one of the following surface types [SurfaceType]
%     plane
%     sphere
%     pseudosphere
%     cylinder
%     cone
%     hyperboloid
%     saddle
%     torus
%     mobius
%     parameterized - A parameterized surface according to a
%                     parameterization you entered before.
% with grid size [GridSize].
% The boolean variable [Geodesic] tells it whether to show an example of a
% geodesic on the surface, or not.

% To do:
%   * Compute geodesic using Runge-Kutta
%   * Add parameterized surfaces

function GenerateSurface(SurfaceType, GridSize, Geodesic)
    Dimensions = 2; % Number of spatial + time dimensions
    Time = 15;       % Number of steps that the particle takes
    dt = 0.025;     % Size of each time interval
    Steps = round(Time / dt);
    MetricThreshold = 100;
    CurvatureThreshold = 10;

    invg = zeros(GridSize, GridSize, Dimensions, Dimensions); % The inverse of the metric tensor, i.e. g^{ij}
    DD_g = zeros(GridSize, GridSize, Dimensions, Dimensions, Dimensions); % The partial derivatives of the metric tensor
    Gamma = zeros(GridSize, GridSize, Dimensions, Dimensions, Dimensions); % The Christoffel symbols
    DD_Gamma = zeros(GridSize, GridSize, Dimensions, Dimensions, Dimensions, Dimensions); % The partial derivatives of the Christoffel symbols
    Riemann = zeros(GridSize, GridSize, Dimensions, Dimensions, Dimensions, Dimensions); % The Riemann curvature tensor
    Ricci = zeros(GridSize, GridSize, Dimensions, Dimensions); % The Ricci curvature tensor
    R = zeros(GridSize, GridSize); % The scalar curvature
    ParticleX = zeros(Steps, Dimensions); % The position of particle in local coordinates
    ParticleV = zeros(Steps, Dimensions); % The velocity of particle in local coordinates
    ParticleTrajectory = zeros(Steps, 3); % The position of the particle in 3D space

    uMin = 0;
    uMax = 1;
    vMin = 0;
    vMax = 1;
    
    uPeriodic = false; % Is the u coordinate periodic?
    vPeriodic = false; % Is the v coordinate periodic?

    % du and dv are taken to be uniform along the surface
    U = 0:1/(GridSize-1):1;
    V = 0:1/(GridSize-1):1;
    
    [v, u] = meshgrid(V,U);
    
    % First step is to generate our surface according to the
    % parameterization.
    % For each surface we rescale u & v to be in the right ranges,
    % set x, y and z according to the parameterization and compute the
    % differentials du & dv
    if (strcmp(SurfaceType, 'plane'))
        u = 10 .* u - 5;
        v = 10 .* v - 5;
        uMin = -5; vMin = -5;
        uMax = 5; vMax = 5;
        
        x = u;
        y = v;
        z = u;
    elseif (strcmp(SurfaceType, 'sphere'))
        % We use spherical coordinates, were u=\theta goes from 0 to pi,
        % and v=\phi goes from 0 to 2*pi.
        u = pi .* u; %(pi - 0.5) * u + 0.25;
        v = 2 .* pi .* v;
        uMin = 0; vMin = 0;
        uMax = pi; vMax = 2 * pi;
        
        x = sin(u) .* cos(v);
        y = sin(u) .* sin(v);
        z = cos(u);

        vPeriodic = true;
    elseif (strcmp(SurfaceType, 'pseudosphere'))
        u = 10 .* u - 5;
        v = 2 .* pi .* v;
        uMin = -5; vMin = 0;
        uMax = 5; vMax = 2 * pi;
        
        x = sech(u) .* cos(v);
        y = sech(u) .* sin(v);
        z = (u - tanh(u));

        vPeriodic = true;
    elseif (strcmp(SurfaceType, 'cylinder'))
        u = 10 .* u - 5;
        v = 2 .* pi .* v;
        uMin = -5; vMin = 0;
        uMax = 5; vMax = 2 * pi;
        
        x = cos(v);
        y = sin(v);
        z = u;
        
        vPeriodic = true;
    elseif (strcmp(SurfaceType, 'cone'))
        u = 10 .* u - 5;
        v = 2 .* pi .* v;
        uMin = -5; vMin = 0;
        uMax = 5; vMax = 2 * pi;

        x = u .* cos(v);
        y = u .* sin(v);
        z = u;

        vPeriodic = true;
    elseif (strcmp(SurfaceType, 'hyperboloid'))
        u = 40 .* u - 20;
        v = 2 .* pi .* v;
        uMin = -20; vMin = 0;
        uMax = 20; vMax = 2 * pi;
        
        x = sqrt(u.^2 + 1) .* cos(v);
        y = sqrt(u.^2 + 1) .* sin(v);
        z = u;

        vPeriodic = true;
    elseif (strcmp(SurfaceType, 'saddle'))
        u = 10 .* u - 5;
        v = 10 .* v - 5;
        uMin = -5; vMin = -5;
        uMax = 5; vMax = 5;
        
        x = u;
        y = v;
        z = u.^2 - v.^2;
    elseif (strcmp(SurfaceType, 'torus'))
        u = 2 .* pi .* u;
        v = 2 .* pi .* v;
        uMin = 0; vMin = 0;
        uMax = 2 * pi; vMax = 2 * pi;

        x = (2 + 1 .* cos(v)) .* cos(u);
        y = (2 + 1 .* cos(v)) .* sin(u);
        z = sin(v);
        
        uPeriodic = true;
        vPeriodic = true;
    elseif (strcmp(SurfaceType, 'mobius'))
        u = 2 .* pi .* u;
        v = 2 .* v - 1;
        uMin = 0; vMin = 1;
        uMax = 2 * pi; vMax = 1;

        x = (1 + 0.5 .* v .* cos(u/2)) .* cos(u);
        y = (1 + 0.5 .* v .* cos(u/2)) .* sin(u);
        z = 0.5 .* v .* sin(u/2);
        
        uPeriodic = true;
    elseif (strcmp(SurfaceType, 'parameterized'))
        error 'Parameterized surfaces are still not supported in this version.'
    else
        error 'Unrecognized surface.'
    end
    
    du = (u(GridSize,GridSize)-u(1,1))/(GridSize - 1);
    dv = (v(GridSize,GridSize)-v(1,1))/(GridSize - 1);
    
    function [xCoordinate, yCoordinate, zCoordinate] = CoordinatesTo3D(uCoordinate, vCoordinate)
        if (strcmp(SurfaceType, 'plane'))
            xCoordinate = uCoordinate;
            yCoordinate = vCoordinate;
            zCoordinate = uCoordinate;
        elseif (strcmp(SurfaceType, 'sphere'))
            xCoordinate = sin(uCoordinate) .* cos(vCoordinate);
            yCoordinate = sin(uCoordinate) .* sin(vCoordinate);
            zCoordinate = cos(uCoordinate);
        elseif (strcmp(SurfaceType, 'pseudosphere'))
            xCoordinate = sech(uCoordinate) .* cos(vCoordinate);
            yCoordinate = sech(uCoordinate) .* sin(vCoordinate);
            zCoordinate = uCoordinate - tanh(uCoordinate);
        elseif (strcmp(SurfaceType, 'cylinder'))
            xCoordinate = cos(vCoordinate);
            yCoordinate = sin(vCoordinate);
            zCoordinate = uCoordinate;
        elseif (strcmp(SurfaceType, 'cone'))
            xCoordinate = uCoordinate .* cos(vCoordinate);
            yCoordinate = uCoordinate .* sin(vCoordinate);
            zCoordinate = uCoordinate;
        elseif (strcmp(SurfaceType, 'hyperboloid'))
            xCoordinate = sqrt(uCoordinate.^2 + 1) .* cos(vCoordinate);
            yCoordinate = sqrt(uCoordinate.^2 + 1) .* sin(vCoordinate);
            zCoordinate = uCoordinate;
        elseif (strcmp(SurfaceType, 'saddle'))
            xCoordinate = uCoordinate;
            yCoordinate = vCoordinate;
            zCoordinate = uCoordinate.^2 - vCoordinate.^2;
        elseif (strcmp(SurfaceType, 'torus'))
            xCoordinate = (2 + 1 .* cos(vCoordinate)) .* cos(uCoordinate);
            yCoordinate = (2 + 1 .* cos(vCoordinate)) .* sin(uCoordinate);
            zCoordinate = sin(vCoordinate);
        elseif (strcmp(SurfaceType, 'mobius'))
            xCoordinate = (1 + 0.5 .* vCoordinate .* cos(uCoordinate/2)) .* cos(uCoordinate);
            yCoordinate = (1 + 0.5 .* vCoordinate .* cos(uCoordinate/2)) .* sin(uCoordinate);
            zCoordinate = 0.5 .* vCoordinate .* sin(uCoordinate/2);
        else
            error 'Unrecognized surface.'
        end
    end

    [DxDu,DxDv] = Gradient(x, du, dv, uPeriodic, vPeriodic);
    [DyDu,DyDv] = Gradient(y, du, dv, uPeriodic, vPeriodic);
    [DzDu,DzDv] = Gradient(z, du, dv, uPeriodic, vPeriodic);

    for iSum = 1:GridSize
        for jSum = 1:GridSize
            DDu1(1) = DxDu(iSum,jSum); DDv1(1) = DxDv(iSum,jSum);
            DDu1(2) = DyDu(iSum,jSum); DDv1(2) = DyDv(iSum,jSum);
            DDu1(3) = DzDu(iSum,jSum); DDv1(3) = DzDv(iSum,jSum);

            % Metric coefficients.  Call them a, b, c, d for the moment; we'll
            % insert them in the g and ingv arrays a few lines below.
            a = DDu1 * DDu1'; b = DDu1 * DDv1';
            c = DDv1 * DDu1'; d = DDv1 * DDv1';

            g(iSum,jSum,1,1) = a; g(iSum,jSum,1,2) = b;
            g(iSum,jSum,2,1) = c; g(iSum,jSum,2,2) = d;

            % Inverse metric coefficients
            gdet = a*d - b*c;

            if (gdet ~= 0)
                invg(iSum,jSum,1,1) =  d/gdet; invg(iSum,jSum,1,2) = -b/gdet;
                invg(iSum,jSum,2,1) = -c/gdet; invg(iSum,jSum,2,2) =  a/gdet;
                
                if (abs(invg(iSum,jSum,1,1)) > MetricThreshold)
                    invg(iSum,jSum,1,1) = sign(invg(iSum,jSum,1,1)) * MetricThreshold;
                end

                if (abs(invg(iSum,jSum,1,2)) > MetricThreshold)
                    invg(iSum,jSum,1,2) = sign(invg(iSum,jSum,1,2)) * MetricThreshold;
                end

                if (abs(invg(iSum,jSum,2,1)) > MetricThreshold)
                    invg(iSum,jSum,2,1) = sign(invg(iSum,jSum,2,1)) * MetricThreshold;
                end

                if (abs(invg(iSum,jSum,2,2)) > MetricThreshold)
                    invg(iSum,jSum,2,2) = sign(invg(iSum,jSum,2,2)) * MetricThreshold;
                end
            else
                invg(iSum,jSum,1,1) = MetricThreshold;
                invg(iSum,jSum,1,2) = MetricThreshold;
                invg(iSum,jSum,2,1) = MetricThreshold;
                invg(iSum,jSum,2,2) = MetricThreshold;
%                sprintf('The metric is singular at %d x %d', iSum, jSum)
            end

            % Unit normals
            N1    = cross(DDu1, DDv1);
            N1hat = N1 / norm(N1);

            % Store the tangent basis and unit normals for later.
            DDu (iSum,jSum,:) = DDu1;
            DDv (iSum,jSum,:) = DDv1;
            Nhat(iSum,jSum,:) = N1hat;
        end
    end
    
    % Here we compute the partial derivative of g
    for i=1:Dimensions
        for j=1:Dimensions
            Dummy = zeros(GridSize, GridSize);

            % TO DO: Rewrite it in a slicker way, there must be one!
            for iSum=1:GridSize
                for jSum=1:GridSize
                    Dummy(iSum,jSum) = g(iSum,jSum,i,j);
                end
            end

            DgDu = zeros(GridSize, GridSize);
            DgDv = zeros(GridSize, GridSize);

            [DgDu, DgDv] = Gradient(Dummy, du, dv, uPeriodic, vPeriodic);

            % TO DO: Rewrite it in a slicker way, there must be one!
            for iSum=1:GridSize
                for jSum=1:GridSize
                    DD_g(:,:,1,i,j) = DgDu(:,:);
                    DD_g(:,:,2,i,j) = DgDv(:,:);
                end
            end
        end
    end

    % And next, we compute the Christoffel symbols
    for i=1:Dimensions
        for j=1:Dimensions
            for k=1:Dimensions
                for iSum=1:Dimensions
                    Gamma(:,:,i,j,k) = Gamma(:,:,i,j,k) + 0.5 * invg(:,:,k,iSum) .* (DD_g(:,:,i,j,iSum) + DD_g(:,:,j,i,iSum) - DD_g(:,:,iSum,i,j));
                end
            end
        end
    end


%    sprintf('%d, %d, %d, %d', 1, 1, 2, DD_g(15,30,2,2,2))
    
    % We compute the partial derivatives of the Christoffel symbols
    for i=1:Dimensions
        for j=1:Dimensions
            for k=1:Dimensions
                Dummy = zeros(GridSize, GridSize);

                % TO DO: Rewrite it in a slicker way, there must be one!
                for iSum=1:GridSize
                    for jSum=1:GridSize
                        Dummy(iSum,jSum) = Gamma(iSum,jSum,i,j,k);
                    end
                end

                DGammaDu = zeros(GridSize, GridSize);
                DGammaDv = zeros(GridSize, GridSize);

                [DGammaDu, DGammaDv] = Gradient(Dummy, du, dv, uPeriodic, vPeriodic);

                % TO DO: Rewrite it in a slicker way, there must be one!
                for iSum=1:GridSize
                    for jSum=1:GridSize
                        DD_Gamma(:,:,1,i,j,k) = DGammaDu(:,:);
                        DD_Gamma(:,:,2,i,j,k) = DGammaDv(:,:);
                    end
                end
            end
        end
    end

    % And now the Riemann curvature tensor
    for i=1:Dimensions
        for j=1:Dimensions
            for k=1:Dimensions
                for l=1:Dimensions
                    Riemann(:,:,i,j,k,l) = DD_Gamma(:,:,j,i,k,l) - DD_Gamma(:,:,i,j,k,l);

                    for iSum=1:Dimensions
                        Riemann(:,:,i,j,k,l) = Riemann(:,:,i,j,k,l) + Gamma(:,:,i,k,iSum) .* Gamma(:,:,iSum,j,l) - Gamma(:,:,j,k,iSum) .* Gamma(:,:,iSum,i,l);
                    end
                end
            end
        end
    end

    % And the Ricci curvature tensor
    for i=1:Dimensions
        for j=1:Dimensions
            for iSum=1:Dimensions
    %            for jSum=1:Dimensions
    %            Ricci(:,:,i,j) = Ricci(:,:,i,j) + Riemann(:,:,i,iSum,j,iSum);
    %                Ricci(:,:,i,j) = Ricci(:,:,i,j) + invg(:,:,iSum,jSum) .* Riemann(:,:,iSum,i,j,jSum);
                    Ricci(:,:,i,j) = Ricci(:,:,i,j) + Riemann(:,:,i,iSum,j,iSum);
    %            end
            end
        end
    end

    % And finally - the scalar curvature
    for iSum=1:Dimensions
        for jSum=1:Dimensions
            R(:,:) = R(:,:) + invg(:,:,iSum,jSum) .* Ricci(:,:,iSum,jSum);
        end
    end

    % Setup an initial position an d velocity for the particle
    if (Geodesic)
        ParticleX(1,1) = 3;
        ParticleX(1,2) = -3;
        ParticleV(1,1) = -4;
        ParticleV(1,2) = 3;

        for i=1:Steps
            % First we make sure that our particle is stil ON the surface,
            % and that if the parameterization is periodic, we take the
            % coordinate 'mod' the period
            if ((ParticleX(i,1) > uMax) || (ParticleX(i,1) < uMin))
                if (uPeriodic)
                    ParticleX(i,1) = uMin + mod(ParticleX(i,1) - uMin, uMax - uMin);
                else
                    ParticleV(i,1) = 0;
                    ParticleV(i,2) = 0;
                end
            end

            if ((ParticleX(i,2) > vMax) || (ParticleX(i,2) < vMin))
                if (vPeriodic)
                    ParticleX(i,2) = vMin + mod(ParticleX(i,2) - vMin, vMax - vMin);
                else
                    ParticleV(i,1) = 0;
                    ParticleV(i,2) = 0;
                end
            end

            ParticleV(i+1,1) = ParticleV(i,1);
            ParticleV(i+1,2) = ParticleV(i,2);
            
            if (~(ParticleV(i+1,1) == 0 && ParticleV(i+1,2) == 0))
                for j=1:Dimensions
                    for iSum=1:Dimensions
                        for jSum=1:Dimensions
                            iIndex = (ParticleX(i,1) - uMin) / (uMax - uMin);
                            jIndex = (ParticleX(i,2) - vMin) / (vMax - vMin);
                            % Interpolating by taking the closest point
                            % ParticleV(i+1,j) = ParticleV(i+1,j) - Gamma(round((GridSize-1) * iIndex) + 1, round((GridSize-1) * jIndex) + 1,iSum,jSum,j) * ParticleV(i,iSum) * ParticleV(i,jSum) * dt;

                            % Interpolation by taking a weighted average
                            ParticleV(i+1,j) = ParticleV(i+1,j) - (1 - iIndex) * (1 - jIndex) * Gamma(floor((GridSize-1) * iIndex) + 1, floor((GridSize-1) * jIndex) + 1,iSum,jSum,j) * ParticleV(i,iSum) * ParticleV(i,jSum) * dt;
                            ParticleV(i+1,j) = ParticleV(i+1,j) - (1 - iIndex) * jIndex * Gamma(floor((GridSize-1) * iIndex) + 1, ceil((GridSize-1) * jIndex) + 1,iSum,jSum,j) * ParticleV(i,iSum) * ParticleV(i,jSum) * dt;
                            ParticleV(i+1,j) = ParticleV(i+1,j) - iIndex * (1 - jIndex) * Gamma(ceil((GridSize-1) * iIndex) + 1, floor((GridSize-1) * jIndex) + 1,iSum,jSum,j) * ParticleV(i,iSum) * ParticleV(i,jSum) * dt;
                            ParticleV(i+1,j) = ParticleV(i+1,j) - iIndex * jIndex * Gamma(ceil((GridSize-1) * iIndex) + 1, ceil((GridSize-1) * jIndex) + 1,iSum,jSum,j) * ParticleV(i,iSum) * ParticleV(i,jSum) * dt;
                        end
                    end
                end
            end
            
            % if the particle's speed is too high, we stop it to avoid
            % memory problems...
            if (abs(ParticleV(i+1,1) > 10^60) || abs(ParticleV(i+1,2)) > 10^60)
                ParticleV(i+1,1) = 0;
                ParticleV(i+1,2) = 0;
                sprintf('Particle was stopped aritifically due to an high-velocity.')
            end

            ParticleX(i+1,1) = ParticleX(i,1) + dt * ParticleV(i+1,1);
            ParticleX(i+1,2) = ParticleX(i,2) + dt * ParticleV(i+1,2);
        end
    end

    CurvatureColor = R;
    
    for i=1:GridSize
        for j=1:GridSize
            if (abs(CurvatureColor(i,j)) > CurvatureThreshold)
                CurvatureColor(i,j) = sign(CurvatureColor(i,j)) * CurvatureThreshold;
            end
        end
    end
    
    h = surf(x, y, z, CurvatureColor, 'FaceColor', 'interp', 'EdgeColor', 'black', 'FaceLighting', 'phong');
    xlabel('x');
    ylabel('y');
    zlabel('z');
    axis square;
    colorbar;

    hold on;

%    quiver3(x, y, z, DxDu, DyDu, DzDu);
%    quiver3(x, y, z, DxDv, DyDv, DzDv);    
    
    if (Geodesic)
        for i=1:Steps
            [ParticleTrajectory(i,1),ParticleTrajectory(i,2),ParticleTrajectory(i,3)] = CoordinatesTo3D(ParticleX(i,1), ParticleX(i,2));
        end

        %[ParticleTrajectory(:,1),ParticleTrajectory(:,2),ParticleTrajectory(:,3)] = CoordinatesTo3D(ParticleX(:,1),ParticleX(:,2));
        comet3(ParticleTrajectory(:,1),ParticleTrajectory(:,2),ParticleTrajectory(:,3));
    end

end
