%==============================================================================
% This code is part of the Finite Element Method app for the Matlab-based toolbox
%  FAIR - Flexible Algorithms for Image Registration. 
% For details see 
% - https://github.com/C4IR/FAIRFEM 
%==============================================================================
%
% classdef TriMesh3 < handle
%
% Finite Element Mesh based on triangular subdivision of rectangular mesh.
%
% Each Cell is divided into four triangles:
%
%  o---------o---------o---------o---------o
%  | \     / | \     / | \     / | \     / |
%  |   \ /   |   \ /   |   \ /   |   \ /   |
%  |    X    |    X    |    X    |    X    |
%  |  /  \   |  /  \   |  /  \   |  /  \   |
%  |/      \ |/      \ |/      \ |/      \ |
%  o---------o---------o---------o---------o
%  | \     / | \     / | \     / | \     / |
%  |   \ /   |   \ /   |   \ /   |   \ /   |
%  |    X    |    X    |    X    |    X    |
%  |  /  \   |  /  \   |  /  \   |  /  \   |
%  |/      \ |/      \ |/      \ |/      \ |
%  o---------o---------o---------o---------o
%  | \     / | \     / | \     / | \     / |
%  |   \ /   |   \ /   |   \ /   |   \ /   |
%  |    X    |    X    |    X    |    X    |
%  |  /  \   |  /  \   |  /  \   |  /  \   |
%  |/      \ |/      \ |/      \ |/      \ |
%  o---------o---------o---------o---------o
%
%  To construct an instance of this class type:
%
%  >> Mesh = TriMesh3(omega,m)
%
% Input:
% 	omega - description of spatial domain
%   m     - number of cells
%
% Properties:
%   xn     - node list
%   tri    - triangle list
%   dim    - space dimension
%   omega  - description of spatial domain
%   m      - number of cells
%   type   - type of partition
%   vol    - volume of triangles
%   nnodes - number of nodes
%   ntri   - number of triangles
%   dx1    - partial derivative operator
%   dx2    - partial derivative operator
%   GRAD   - gradient operator
%   P1     - projection operator for node 1
%   P2     - projection operator for node 2
%   P3     - projection operator for node 3
%   PC     - projection operator for Barycentrum
%   P      - prolongation operator
%   Pt     - restriction operator
%
%  Methods:
%   mfPu   - matrix free prolongation/restriction
%   mfPi   - matrix free edge projector
%   getP   - builds prolongation operator
%   tri2cc - averaging
%
%
% see also
% =========================================================================
classdef TriMesh3 < handle
    
    properties
        % ===================================
        % node list
        % ===================================
        xn
        % ===================================
        % triangle list
        % ===================================
        tri
        % ===================================
        % space dimension
        % ===================================
        dim  = 2;
        % ===================================
        % description of computational domain
        % ===================================
        omega
        % ===================================
        % number of cells
        % ===================================
        m
        % ===================================
        % type of partition
        % ===================================
        type = 3;
        % ===================================
        % function handle to myself
        % ===================================
        me = @TriMesh3
        % ===================================
        % volume of triangles
        % ===================================
        vol
        % ===================================
        % number of nodes
        % ===================================
        nnodes
        % ===================================
        % number of triangles in mesh
        % ===================================
        ntri
    end
    
    properties (Access = public, Dependent) % These will be created when first callend and stores persistently
        % ===================================
        % dx1  - partial derivative operator
        % ===================================
        dx1
        % ===================================
        % dx2  - Partial derivative operator
        % ===================================
        dx2
        % ===================================
        % GRAD - Gradient operator
        %
        %         | dx1 |
        %  GRAD = |     |
        %         | dx2 |
        %
        % ===================================
        GRAD
        % ===================================
        % B - Vector gradient operator
        %
        %         | GRAD  0    |
        %  B =    |            |
        %         | 0     GRAD |
        %
        % ===================================
        B
        % ===================================
        % P1 - Projection operator on Node 1
        % ===================================
        P1
        % ===================================
        % P2 - Projection operator on Node 2
        % ===================================
        P2
        % ===================================
        % P3 - Projection operator on Node 3
        % ===================================
        P3
        % ===================================
        % PC - Projection operator on Barycenter
        % ===================================
        PC
        % ===================================
        % P  - Prolongation operator
        % ===================================
        P
        % ===================================
        % Pt - Restriction operator
        % ===================================
        Pt
        % ===================================
        % Boundary indices
        % ===================================
        boundaryIdx
        % ===================================
        % Boundary projector
        % ===================================
        boundaryProj
        AvN
        mfdx1
        mfdx2
        mfGRAD
    end
    
    properties (Access = private)
        % These are where the dependent data is actually stored
        dx1_
        dx2_
        GRAD_
        B_
        P1_
        P2_
        P3_
        PC_
        P_
        Pt_
        mfdx1_
        mfdx2_
        mfGRAD_
        boundaryIdx_
        boundaryProj_
        AvN_
    end
    
    
    methods
        function this = TriMesh3(omega,m)
            if nargin==0,
                help(mfilename);
                this.runMinimalExample;
                return;
            end
            this.omega = omega;
            this.m     = m;
            
            this.xn = [ reshape(getNodalGrid(omega,m),[],2);...
               reshape(getCellCenteredGrid(omega,m),[],2)];
  
            % get indices of bottom left vertices
            nodes = reshape(1:prod(m+1),m+1);
            nodes = reshape(nodes(1:end-1,1:end-1),1,[]);
            % get indices of cell-centered points
            cc    = prod(m+1)+(1:prod(m));
            % specify triangles
            this.tri   = [...
                nodes;          nodes+1;        cc; ...
                nodes+1;        nodes+m(1)+2;   cc; ...
                nodes+m(1)+2;   nodes+m(1)+1;   cc; ...
                nodes+m(1)+1;   nodes;          cc; ...
                ];
            this.tri   = reshape(this.tri,3,[])';
            
            this.nnodes = prod(m+1)+prod(m);
            this.ntri  = size(this.tri,1);
            this.vol = prod((omega(2:2:end)-omega(1:2:end))./m)/4*ones(this.ntri,1);
            
        end
        
        function runMinimalExample(~)
            omega = [0 4 2 6]; m = [3 6];
            Mesh  = feval(mfilename,omega,m);
            figure(3); clf;
            set(gcf,'Name',sprintf('%s, m=[%d,%d]',mfilename,m));
            plotTriMesh(Mesh,Mesh.xn);
        end
        
        function x = mfPi(this,x,i)
            % =============================================================
            % function x = mfPi(this,x,i)
            %
            % matrix free edge projector
            % =============================================================
            switch i
                case 1
                    P = this.P1;
                case 2
                    P = this.P2;
                case 3
                    P = this.P3;
                case 'C'
                    P = this.PC;
            end
            if size(x,1) == this.ntri,
                % ajoint
                x = P'*x;
            else
                x = P * x;
            end
        end
        
        function x = tri2cc(this,x)
            % =============================================================
            % function x = tri2cc(x)
            %
            % averaging or adjoint
            % =============================================================
            if numel(x)==this.ntri,
                x = mean(reshape(x,4,[]),1);
            else
                x = reshape(x,1,[]);
                x = .25*repmat(x,[1 4]);
                x = x(:);
            end
            
        end
        
        % ===================================
        % x = getNodalGridData(this,x)
        %
        % gets values at triangle nodes that belong the underlying nodal
        % grid
        % ===================================
        function x = getNodalGridData(this,x)
            x = reshape(x,prod(this.nnodes),[]);
            x = x(1:prod(this.m+1),:);
            x = x(:);
        end
        
        % ===================================
        % x = getCCGridData(this,x)
        %
        % gets values at cell-centers that belong the underlying nodal
        % grid
        % ===================================
        function x = getCCGridData(this,x)
            x = reshape(x,prod(this.nnodes),[]);
            x = x(prod(this.m+1)+1:end,:);
            x = x(:);
        end
        
        
        
        function Pu = mfPuNodal(~,yn,m,flag)
            % =============================================================
            % function Pu = mfPuNodal(~,yn,m,flag)
            %
            % matrix free prolongation/restriction for nodal quantities
            % =============================================================
            switch flag
                case 'Pu' % coarse --> fine
                    yn = reshape(yn,[],2);
                    yc = reshape(yn(prod(m+1)+1:end,:),[m 2]);
                    yn = reshape(yn(1:prod(m+1),:),[m+1 2]);
                    
                    Pyn = zeros([2*m+1 2]);
                    Pyc = zeros([2*m 2]);
                    
                    % include existing nodes
                    Pyn(1:2:end,1:2:end,:) = yn;
                    % prolongate in x direction
                    Pyn(2:2:end-1,:,:) =  .5*(Pyn(1:2:end-2,:,:) + Pyn(3:2:end,:,:));
                    % prolongate in y direction
                    Pyn(1:2:end,2:2:end-1,:) =  .5*(Pyn(1:2:end,1:2:end-2,:) + Pyn(1:2:end,3:2:end,:));
                    % midpoints
                    Pyn(2:2:end,2:2:end,:)   = yc;
                    
                    % compute cell-centers
                    Pyc(1:2:end-1,1:2:end-1,:) = .5*(yn(1:end-1,1:end-1,:) + yc);
                    Pyc(2:2:end  ,1:2:end-1,:) = .5*(yn(2:end  ,1:end-1,:) + yc);
                    Pyc(1:2:end-1,2:2:end  ,:) = .5*(yn(1:end-1,2:end  ,:) + yc);
                    Pyc(2:2:end  ,2:2:end  ,:) = .5*(yn(2:end  ,2:end  ,:) + yc);
                    
                    Pu = [reshape(Pyn,[],2); reshape(Pyc,[],2)];
                    Pu = Pu(:);
                case 'PTu'
                    yn = reshape(yn,[],2);
                    yc = reshape(yn(prod(m+1)+1:end,:),[m 2]);
                    yn = reshape(yn(1:prod(m+1),:),[m+1 2]);
                    
                    % include parent nodes
                    Pyn = yn(1:2:end,1:2:end,:);
                    % distribute in x direction
                    t = .5* yn(2:2:end-1,1:2:end,:);
                    Pyn(1:end-1,:,:)   = Pyn(1:end-1,:,:) + t;
                    Pyn(2:end  ,:,:)   = Pyn(2:end  ,:,:) + t;
                    
                    % distribute in y direction
                    t = .5* yn(1:2:end,2:2:end-1,:);
                    Pyn(:,1:end-1,:) = Pyn(:,1:end-1,:) + t;
                    Pyn(:,2:end  ,:) = Pyn(:,2:end  ,:) + t;
                    
                    % distribute cell-centers
                    t1 = .5*yc(1:2:end-1,1:2:end-1,:);
                    t2 = .5*yc(2:2:end  ,1:2:end-1,:);
                    t3 = .5*yc(1:2:end-1,2:2:end  ,:);
                    t4 = .5*yc(2:2:end  ,2:2:end  ,:);
                    Pyc = yn(2:2:end,2:2:end,:) + (t1+t2+t3+t4);
                    Pyn(1:end-1,1:end-1,:) = Pyn(1:end-1,1:end-1,:) + t1;
                    Pyn(2:end  ,1:end-1,:) = Pyn(2:end  ,1:end-1,:) + t2;
                    Pyn(1:end-1,2:end  ,:) = Pyn(1:end-1,2:end  ,:) + t3;
                    Pyn(2:end  ,2:end  ,:) = Pyn(2:end  ,2:end  ,:) + t4;
                    
                    Pu = [reshape(Pyn,[],2); reshape(Pyc,[],2)];
                    Pu = Pu(:);
            end
        end
        
        function P = getPuNodal(~,m)
            % =============================================================
            % function P = getPuNodal(~,m)
            %
            % returns prolongation operator for input of cell-width m
            % =============================================================
            
            mf = 2*m;
            % indices of fine and coarse grid nodes
            indfn = reshape(1:prod(mf+1),mf+1);
            indfc = reshape(prod(mf+1)+(1:prod(mf)),mf);
            indcn = reshape(1:prod(m+1),m+1);
            indcc = reshape(prod(m+1)+(1:prod(m)),m);
            
            % allocate space
            I = []; J = []; W = [];
            
            % include existing nodes
            ii = indfn(1:2:end,1:2:end);
            jj = indcn(:);
            ww = ones(size(jj));
            I = [I; ii(:)]; J = [J; jj(:)]; W = [W; ww(:)];
            
            % include midpoint as node
            ii = indfn(2:2:end,2:2:end);
            jj = indcc(:);
            ww = ones(size(jj));
            I = [I; ii(:)]; J = [J; jj(:)]; W = [W; ww(:)];
            
            % prolongate in x direction
            ii   = indfn(2:2:end-1,1:2:end);
            jj   = [reshape(indcn(1:end-1,:),[],1); reshape(indcn(2:end,:),[],1)];
            ww = .5*ones(size(jj));
            I = [I; ii(:);ii(:)]; J = [J; jj(:)]; W = [W; ww(:)];
            
            % prolongate in y direction
            ii   = indfn(1:2:end,2:2:end-1);
            jj   = [reshape(indcn(:,1:end-1),[],1); reshape(indcn(:,2:end),[],1)];
            ww   = .5*ones(size(jj));
            I = [I; ii(:); ii(:)]; J = [J; jj(:)]; W = [W; ww(:)];
            
            % prolongate from bottom-left to top-right
            ii   = indfc(1:2:end,1:2:end);
            jj   = [reshape(indcn(1:end-1,1:end-1),[],1); reshape(indcc,[],1)];
            ww   = .5*ones(size(jj));
            I = [I; ii(:); ii(:)]; J = [J; jj(:)]; W = [W; ww(:)];
            ii   = indfc(2:2:end,2:2:end);
            jj   = [reshape(indcn(2:end,2:end),[],1); reshape(indcc,[],1)];
            ww   = .5*ones(size(jj));
            I = [I; ii(:); ii(:)]; J = [J; jj(:)]; W = [W; ww(:)];
            
            
            % prolongate from top-left to bottom-right
            ii   = indfc(2:2:end,1:2:end);
            jj   = [reshape(indcn(2:end,1:end-1),[],1); reshape(indcc,[],1)];
            ww   = .5*ones(size(jj));
            I = [I; ii(:); ii(:)]; J = [J; jj(:)]; W = [W; ww(:)];
            ii   = indfc(1:2:end,2:2:end);
            jj   = [reshape(indcn(1:end-1,2:end),[],1); reshape(indcc,[],1)];
            ww   = .5*ones(size(jj));
            I = [I; ii(:); ii(:)]; J = [J; jj(:)]; W = [W; ww(:)];
            
            P = sparse(I,J,W,prod(mf+1)+prod(mf),prod(m+1)+prod(m));
        end
        
        
        
        % ========== get methods ========================================
        function dx1 = get.dx1(this)
            if isempty(this.dx1_),
                [this.dx1_, this.dx2_] = getGradientMatrixFEM(this,0);
            end
            dx1 = this.dx1_;
        end
        
        function dx2 = get.dx2(this)
            if isempty(this.dx2_),
                [this.dx1_, this.dx2_] = getGradientMatrixFEM(this,0);
            end
            dx2 = this.dx2_;
        end
        
        function GRAD = get.GRAD(this)
            if isempty(this.GRAD_),
                this.GRAD_ = [this.dx1;this.dx2];
            end
            GRAD = this.GRAD_;
        end
        function B = get.B(this)
            if isempty(this.B_),
                this.B_ = blkdiag(this.GRAD,this.GRAD);
            end
            B = this.B_;
        end
        
        function P1 = get.P1(this)
            if isempty(this.P1_),
                A = speye(this.nnodes);
                this.P1_ = A(this.tri(:,1),:);
            end
            P1 = this.P1_;
        end
        function P2 = get.P2(this)
            if isempty(this.P2_),
                A = speye(this.nnodes);
                this.P2_ = A(this.tri(:,2),:);
            end
            P2 = this.P2_;
        end
        
        function P3 = get.P3(this)
            if isempty(this.P3_),
                A = speye(this.nnodes);
                this.P3_ = A(this.tri(:,3),:);
            end
            P3 = this.P3_;
        end
        function PC = get.PC(this)
            if isempty(this.PC_),
                this.PC_ = (this.P1+this.P2+this.P3)/3;
            end
            PC = this.PC_;
        end
        function P = get.P(this)
            if isempty(this.P_),
                this.P_ = this.getPuNodal(this.m);
            end
            P = this.P_;
        end
        function Pt = get.Pt(this)
            if isempty(this.Pt_),
                this.Pt_ = this.getPuNodal(this.m/2);
            end
            Pt = this.Pt_;
        end
        
        function mfdx1 = get.mfdx1(this)
            if isempty(this.mfdx1_),
                [this.mfdx1_, this.mfdx2_] = getGradientMatrixFEM(this,1);
            end
            mfdx1 = this.mfdx1_;
        end
        
        function mfdx2 = get.mfdx2(this)
            if isempty(this.mfdx2_),
                [this.mfdx1_, this.mfdx2_] = getGradientMatrixFEM(this,1);
            end
            mfdx2 = this.mfdx2_;
        end
        
        function AvN = get.AvN(this)
            if isempty(this.AvN_),
                av  = @(n) spdiags(ones(n+1,1)*[0.5 0.5],[0,1],n,n+1);
                this.AvN_ = kron(speye(2), [speye(prod(this.m+1)); kron(av(this.m(2)),av(this.m(1)))]);
            end
            AvN = this.AvN_;
        end
        
        function mfGRAD = get.mfGRAD(this)
            if isempty(this.mfGRAD_),
                this.mfGRAD_ = getGradientMatrixFEM(this,1);
            end
            mfGRAD = this.mfGRAD_;
        end
        function idx = get.boundaryIdx(this)
            if isempty(this.boundaryIdx_),
                id = reshape(1:prod(this.m+1),this.m+1);
                idx = [reshape(id([1,end],:),[],1);reshape(id(2:end-1,[1,end]),[],1)] ;
                this.boundaryIdx_ = idx;
            end
            idx = this.boundaryIdx_;
        end
        
        function idx = get.boundaryProj(this)
            if isempty(this.boundaryProj_),
                idx = this.boundaryIdx;
                
                P = speye(prod(this.m+1));
                P = P(idx,:);
                P = [P,sparse(size(P,1),prod(this.m))];
                P = kron(speye(2),P);
                
                this.boundaryProj_ = P;
            end
            idx = this.boundaryProj_;
        end 
        
        % ========== set methods ========================================
        function set.xn(this,xn)
            if numel(xn)~=(2*(prod(this.m+1)+prod(this.m))),
                error('Invalid number of nodes');
            end
            if isempty(this.xn),
                this.xn = xn;
            else
                
                this.xn = xn;
                % delete operators that are sensitive to nodes
                this.dx1_  = [];
                this.dx2_  = [];
                this.GRAD_ = [];
                this.B_    = [];
                
                this.vol = volTetraGrid(this,this.xn,'matrixFree',1);
            end
        end
    end
end

