classdef Elastic2dVariableAnisotropic < scheme.Scheme

% Discretizes the elastic wave equation:
% rho u_{i,tt} = dj C_{ijkl} dk u_j
% opSet should be cell array of opSets, one per dimension. This
% is useful if we have periodic BC in one direction.
% Assumes fully compatible operators

    properties
        m % Number of points in each direction, possibly a vector
        h % Grid spacing

        grid
        dim

        order % Order of accuracy for the approximation

        % Diagonal matrices for variable coefficients
        RHO, RHOi, RHOi_kron % Density
        C                    % Elastic stiffness tensor

        D  % Total operator
        D1 % First derivatives
        D2 % Second derivatives

        % Boundary operators in cell format, used for BC
        T_w, T_e, T_s, T_n

        % Traction operators
        tau_w, tau_e, tau_s, tau_n      % Return vector field
        tau1_w, tau1_e, tau1_s, tau1_n  % Return scalar field
        tau2_w, tau2_e, tau2_s, tau2_n  % Return scalar field

        % Inner products
        H, Hi, Hi_kron, H_1D

        % Boundary inner products (for scalar field)
        H_w, H_e, H_s, H_n

        % Boundary restriction operators
        e_w, e_e, e_s, e_n      % Act on vector field, return vector field at boundary
        e1_w, e1_e, e1_s, e1_n  % Act on vector field, return scalar field at boundary
        e2_w, e2_e, e2_s, e2_n  % Act on vector field, return scalar field at boundary
        e_scalar_w, e_scalar_e, e_scalar_s, e_scalar_n; % Act on scalar field, return scalar field

        % E{i}^T picks out component i
        E

        % Borrowing constants of the form gamma*h, where gamma is a dimensionless constant.
        h11 % First entry in norm matrix

    end

    methods

        % The coefficients can either be function handles or grid functions
        % optFlag -- if true, extra computations are performed, which may be helpful for optimization.
        function obj = Elastic2dVariableAnisotropic(g, order, rho, C, opSet, optFlag)
            default_arg('rho', @(x,y) 0*x+1);
            default_arg('opSet',{@sbp.D2VariableCompatible, @sbp.D2VariableCompatible});
            default_arg('optFlag', false);
            dim = 2;

            C_default = cell(dim,dim,dim,dim);
            for i = 1:dim
                for j = 1:dim
                    for k = 1:dim
                        for l = 1:dim
                            C_default{i,j,k,l} = @(x,y) 0*x + 1;
                        end
                    end
                end
            end
            default_arg('C', C_default);
            assert(isa(g, 'grid.Cartesian'))

            if isa(rho, 'function_handle')
                rho = grid.evalOn(g, rho);
            end

            C_mat = cell(dim,dim,dim,dim);
            for i = 1:dim
                for j = 1:dim
                    for k = 1:dim
                        for l = 1:dim
                            if isa(C{i,j,k,l}, 'function_handle')
                                C{i,j,k,l} = grid.evalOn(g, C{i,j,k,l});
                            end
                            C_mat{i,j,k,l} = spdiag(C{i,j,k,l});
                        end
                    end
                end
            end
            obj.C = C_mat;

            m = g.size();
            m_tot = g.N();
            lim = g.lim;
            if isempty(lim)
                x = g.x;
                lim = cell(length(x),1);
                for i = 1:length(x)
                    lim{i} = {min(x{i}), max(x{i})};
                end
            end

            % 1D operators
            ops = cell(dim,1);
            h = zeros(dim,1);
            for i = 1:dim
                ops{i} = opSet{i}(m(i), lim{i}, order);
                h(i) = ops{i}.h;
            end

            % Borrowing constants
            for i = 1:dim
                obj.h11{i} = h(i)*ops{i}.borrowing.H11;
            end

            I = cell(dim,1);
            D1 = cell(dim,1);
            D2 = cell(dim,1);
            H = cell(dim,1);
            Hi = cell(dim,1);
            e_0 = cell(dim,1);
            e_m = cell(dim,1);
            d1_0 = cell(dim,1);
            d1_m = cell(dim,1);

            for i = 1:dim
                I{i} = speye(m(i));
                D1{i} = ops{i}.D1;
                D2{i} = ops{i}.D2;
                H{i} =  ops{i}.H;
                Hi{i} = ops{i}.HI;
                e_0{i} = ops{i}.e_l;
                e_m{i} = ops{i}.e_r;
                d1_0{i} = ops{i}.d1_l;
                d1_m{i} = ops{i}.d1_r;
            end

            %====== Assemble full operators ========
            RHO = spdiag(rho);
            obj.RHO = RHO;
            obj.RHOi = inv(RHO);

            obj.D1 = cell(dim,1);
            obj.D2 = cell(dim,dim,dim);

            % D1
            obj.D1{1} = kron(D1{1},I{2});
            obj.D1{2} = kron(I{1},D1{2});

            % Boundary restriction operators
            e_l = cell(dim,1);
            e_r = cell(dim,1);
            e_l{1} = kron(e_0{1}, I{2});
            e_l{2} = kron(I{1}, e_0{2});
            e_r{1} = kron(e_m{1}, I{2});
            e_r{2} = kron(I{1}, e_m{2});

            e_scalar_w = e_l{1};
            e_scalar_e = e_r{1};
            e_scalar_s = e_l{2};
            e_scalar_n = e_r{2};

            I_dim = speye(dim, dim);
            e_w = kron(e_scalar_w, I_dim);
            e_e = kron(e_scalar_e, I_dim);
            e_s = kron(e_scalar_s, I_dim);
            e_n = kron(e_scalar_n, I_dim);

            % E{i}^T picks out component i.
            E = cell(dim,1);
            I = speye(m_tot,m_tot);
            for i = 1:dim
                e = sparse(dim,1);
                e(i) = 1;
                E{i} = kron(I,e);
            end
            obj.E = E;

            e1_w = (e_scalar_w'*E{1}')';
            e1_e = (e_scalar_e'*E{1}')';
            e1_s = (e_scalar_s'*E{1}')';
            e1_n = (e_scalar_n'*E{1}')';

            e2_w = (e_scalar_w'*E{2}')';
            e2_e = (e_scalar_e'*E{2}')';
            e2_s = (e_scalar_s'*E{2}')';
            e2_n = (e_scalar_n'*E{2}')';


            % D2
            for i = 1:dim
                for k = 2:dim
                    for l = 2:dim
                        obj.D2{i,k,l} = sparse(m_tot);
                    end
                end
            end
            ind = grid.funcToMatrix(g, 1:m_tot);

            k = 1;
            for r = 1:m(2)
                p = ind(:,r);
                for i = 1:dim
                    for l = 1:dim
                        coeff = C{i,k,k,l};
                        D_kk = D2{1}(coeff(p));
                        obj.D2{i,k,l}(p,p) = D_kk;
                    end
                end
            end

            k = 2;
            for r = 1:m(1)
                p = ind(r,:);
                for i = 1:dim
                    for l = 1:dim
                        coeff = C{i,k,k,l};
                        D_kk = D2{2}(coeff(p));
                        obj.D2{i,k,l}(p,p) = D_kk;
                    end
                end
            end

            % Quadratures
            obj.H = kron(H{1},H{2});
            obj.Hi = inv(obj.H);
            obj.H_w = H{2};
            obj.H_e = H{2};
            obj.H_s = H{1};
            obj.H_n = H{1};
            obj.H_1D = {H{1}, H{2}};

            % Differentiation matrix D (without SAT)
            D2 = obj.D2;
            D1 = obj.D1;
            D = sparse(dim*m_tot,dim*m_tot);
            d = @kroneckerDelta;    % Kronecker delta
            db = @(i,j) 1-d(i,j); % Logical not of Kronecker delta
            for i = 1:dim
                for j = 1:dim
                    for k = 1:dim
                        for l = 1:dim
                            D = D + E{i}*inv(RHO)*( d(j,k)*D2{i,k,l}*E{l}' +...
                                                    db(j,k)*D1{j}*C_mat{i,j,k,l}*D1{k}*E{l}' ...
                                                  );
                        end
                    end
                end
            end
            obj.D = D;
            %=========================================%'

            % Numerical traction operators for BC.
            %
            % Formula at boundary j: % tau^{j}_i = sum_l T^{j}_{il} u_l
            %
            T_l = cell(dim,1);
            T_r = cell(dim,1);
            tau_l = cell(dim,1);
            tau_r = cell(dim,1);

            D1 = obj.D1;

            % Boundary j
            for j = 1:dim
                T_l{j} = cell(dim,dim);
                T_r{j} = cell(dim,dim);
                tau_l{j} = cell(dim,1);
                tau_r{j} = cell(dim,1);

                [~, n_l] = size(e_l{j});
                [~, n_r] = size(e_r{j});

                % Traction component i
                for i = 1:dim
                    tau_l{j}{i} = sparse(dim*m_tot, n_l);
                    tau_r{j}{i} = sparse(dim*m_tot, n_r);

                    % Displacement component l
                    for l = 1:dim
                        T_l{j}{i,l} = sparse(m_tot, n_l);
                        T_r{j}{i,l} = sparse(m_tot, n_r);

                        % Derivative direction k
                        for k = 1:dim
                            T_l{j}{i,l} = T_l{j}{i,l} ...
                                        - (e_l{j}'*C_mat{i,j,k,l}*D1{k})';
                            T_r{j}{i,l} = T_r{j}{i,l} ...
                                        + (e_r{j}'*C_mat{i,j,k,l}*D1{k})';
                        end
                        tau_l{j}{i} = tau_l{j}{i} + (T_l{j}{i,l}'*E{l}')';
                        tau_r{j}{i} = tau_r{j}{i} + (T_r{j}{i,l}'*E{l}')';
                    end
                end
            end

            % Traction tensors, T_ij
            obj.T_w = T_l{1};
            obj.T_e = T_r{1};
            obj.T_s = T_l{2};
            obj.T_n = T_r{2};

            % Restriction operators
            obj.e_w = e_w;
            obj.e_e = e_e;
            obj.e_s = e_s;
            obj.e_n = e_n;

            obj.e1_w = e1_w;
            obj.e1_e = e1_e;
            obj.e1_s = e1_s;
            obj.e1_n = e1_n;

            obj.e2_w = e2_w;
            obj.e2_e = e2_e;
            obj.e2_s = e2_s;
            obj.e2_n = e2_n;

            obj.e_scalar_w = e_scalar_w;
            obj.e_scalar_e = e_scalar_e;
            obj.e_scalar_s = e_scalar_s;
            obj.e_scalar_n = e_scalar_n;

            % First component of traction
            obj.tau1_w = tau_l{1}{1};
            obj.tau1_e = tau_r{1}{1};
            obj.tau1_s = tau_l{2}{1};
            obj.tau1_n = tau_r{2}{1};

            % Second component of traction
            obj.tau2_w = tau_l{1}{2};
            obj.tau2_e = tau_r{1}{2};
            obj.tau2_s = tau_l{2}{2};
            obj.tau2_n = tau_r{2}{2};

            % Traction vectors
            obj.tau_w = (e_w'*e1_w*obj.tau1_w')' + (e_w'*e2_w*obj.tau2_w')';
            obj.tau_e = (e_e'*e1_e*obj.tau1_e')' + (e_e'*e2_e*obj.tau2_e')';
            obj.tau_s = (e_s'*e1_s*obj.tau1_s')' + (e_s'*e2_s*obj.tau2_s')';
            obj.tau_n = (e_n'*e1_n*obj.tau1_n')' + (e_n'*e2_n*obj.tau2_n')';

            % Kroneckered norms and coefficients
            obj.RHOi_kron = kron(obj.RHOi, I_dim);
            obj.Hi_kron = kron(obj.Hi, I_dim);

            % Misc.
            obj.m = m;
            obj.h = h;
            obj.order = order;
            obj.grid = g;
            obj.dim = dim;

        end


        % Closure functions return the operators applied to the own domain to close the boundary
        % Penalty functions return the operators to force the solution. In the case of an interface it returns the operator applied to the other doamin.
        %       boundary            is a string specifying the boundary e.g. 'l','r' or 'e','w','n','s'.
        %       bc                  is a cell array of component and bc type, e.g. {1, 'd'} for Dirichlet condition
        %                           on the first component. Can also be e.g.
        %                           {'normal', 'd'} or {'tangential', 't'} for conditions on
        %                           tangential/normal component.
        %       data                is a function returning the data that should be applied at the boundary.
        %       neighbour_scheme    is an instance of Scheme that should be interfaced to.
        %       neighbour_boundary  is a string specifying which boundary to interface to.
        function [closure, penalty] = boundary_condition(obj, boundary, bc, tuning)
            default_arg('tuning', 1.0);

            assert( iscell(bc), 'The BC type must be a 2x1 cell array' );
            comp = bc{1};
            type = bc{2};
            if ischar(comp)
                comp = obj.getComponent(comp, boundary);
            end

            e       = obj.getBoundaryOperatorForScalarField('e', boundary);
            tau     = obj.getBoundaryOperator(['tau' num2str(comp)], boundary);
            T       = obj.getBoundaryTractionOperator(boundary);
            h11     = obj.getBorrowing(boundary);
            H_gamma = obj.getBoundaryQuadratureForScalarField(boundary);
            nu      = obj.getNormal(boundary);

            E = obj.E;
            Hi = obj.Hi;
            RHOi = obj.RHOi;
            C = obj.C;

            dim = obj.dim;
            m_tot = obj.grid.N();

            % Preallocate
            [~, col] = size(tau);
            closure = sparse(dim*m_tot, dim*m_tot);
            penalty = sparse(dim*m_tot, col);

            j = comp;
            switch type

            % Dirichlet boundary condition
            % OBS! Cannot yet set one component at a time unless one assumes Displacement for all components
            case {'D','d','dirichlet','Dirichlet','displacement','Displacement'}

                % Loop over components that Dirichlet penalties end up on
                % Y: symmetrizing part of penalty
                % Z: symmetric part of penalty
                % X = Y + Z.
                for i = 1:dim
                    Y = T{j,i}';

                    Z = sparse(m_tot, m_tot);
                    for l = 1:dim
                        for k = 1:dim
                            Z = Z + nu(l)*C{i,l,k,j}*nu(k);
                        end
                    end
                    Z = -tuning*dim/h11*Z;
                    X = Z + e*Y;
                    closure = closure + E{i}*RHOi*Hi*X'*e*H_gamma*(e'*E{j}' );
                    penalty = penalty - E{i}*RHOi*Hi*X'*e*H_gamma;
                end

            % Free boundary condition
            case {'F','f','Free','free','traction','Traction','t','T'}
                    closure = closure - E{j}*RHOi*Hi*e*H_gamma*tau';
                    penalty = penalty + E{j}*RHOi*Hi*e*H_gamma;

            % Unknown boundary condition
            otherwise
                error('No such boundary condition: type = %s',type);
            end
        end

        % type     Struct that specifies the interface coupling.
        %          Fields:
        %          -- tuning:           penalty strength, defaults to 1.2
        %          -- interpolation:    type of interpolation, default 'none'
        function [closure, penalty] = interface(obj,boundary,neighbour_scheme,neighbour_boundary,type)

            defaultType.tuning = 1.2;
            defaultType.interpolation = 'none';
            default_struct('type', defaultType);

            switch type.interpolation
            case {'none', ''}
                [closure, penalty] = interfaceStandard(obj,boundary,neighbour_scheme,neighbour_boundary,type);
            case {'op','OP'}
                [closure, penalty] = interfaceNonConforming(obj,boundary,neighbour_scheme,neighbour_boundary,type);
            otherwise
                error('Unknown type of interpolation: %s ', type.interpolation);
            end
        end

        function [closure, penalty] = interfaceStandard(obj,boundary,neighbour_scheme,neighbour_boundary,type)
            tuning = type.tuning;

            % u denotes the solution in the own domain
            % v denotes the solution in the neighbour domain
            % Operators without subscripts are from the own domain.

            % Get boundary operators
            e   = obj.getBoundaryOperator('e', boundary);
            tau = obj.getBoundaryOperator('tau', boundary);

            e_v   = neighbour_scheme.getBoundaryOperator('e', neighbour_boundary);
            tau_v = neighbour_scheme.getBoundaryOperator('tau', neighbour_boundary);

            H_gamma = obj.getBoundaryQuadrature(boundary);

            % Operators and quantities that correspond to the own domain only
            Hi = obj.Hi_kron;
            RHOi = obj.RHOi_kron;

            % Penalty strength operators
            alpha_u = 1/4*tuning*obj.getBoundaryOperator('alpha', boundary);
            alpha_v = 1/4*tuning*neighbour_scheme.getBoundaryOperator('alpha', neighbour_boundary);

            closure = -RHOi*Hi*e*H_gamma*(alpha_u' + alpha_v'*e_v*e');
            penalty = RHOi*Hi*e*H_gamma*(alpha_u'*e*e_v' + alpha_v');

            closure = closure - 1/2*RHOi*Hi*e*H_gamma*tau';
            penalty = penalty - 1/2*RHOi*Hi*e*H_gamma*tau_v';

            closure = closure + 1/2*RHOi*Hi*tau*H_gamma*e';
            penalty = penalty - 1/2*RHOi*Hi*tau*H_gamma*e_v';

        end

        function [closure, penalty] = interfaceNonConforming(obj,boundary,neighbour_scheme,neighbour_boundary,type)
            error('Non-conforming interfaces not implemented yet.');
        end

        % Returns the component number that is the tangential/normal component
        % at the specified boundary
        function comp = getComponent(obj, comp_str, boundary)
            assertIsMember(comp_str, {'normal', 'tangential'});
            assertIsMember(boundary, {'w', 'e', 's', 'n'});

            switch boundary
            case {'w', 'e'}
                switch comp_str
                case 'normal'
                    comp = 1;
                case 'tangential'
                    comp = 2;
                end
            case {'s', 'n'}
                switch comp_str
                case 'normal'
                    comp = 2;
                case 'tangential'
                    comp = 1;
                end
            end
        end

        % Returns h11 for the boundary specified by the string boundary.
        % op -- string
        function h11 = getBorrowing(obj, boundary)
            assertIsMember(boundary, {'w', 'e', 's', 'n'})

            switch boundary
            case {'w','e'}
                h11 = obj.h11{1};
            case {'s', 'n'}
                h11 = obj.h11{2};
            end
        end

        % Returns the outward unit normal vector for the boundary specified by the string boundary.
        function nu = getNormal(obj, boundary)
            assertIsMember(boundary, {'w', 'e', 's', 'n'})

            switch boundary
            case 'w'
                nu = [-1,0];
            case 'e'
                nu = [1,0];
            case 's'
                nu = [0,-1];
            case 'n'
                nu = [0,1];
            end
        end

        % Returns the boundary operator op for the boundary specified by the string boundary.
        % op -- string
        function o = getBoundaryOperator(obj, op, boundary)
            assertIsMember(boundary, {'w', 'e', 's', 'n'})
            assertIsMember(op, {'e', 'e1', 'e2', 'tau', 'tau1', 'tau2'})

            switch op
                case {'e', 'e1', 'e2', 'tau', 'tau1', 'tau2'}
                    o = obj.([op, '_', boundary]);
            end

        end

        % Returns the boundary operator op for the boundary specified by the string boundary.
        % op -- string
        function o = getBoundaryOperatorForScalarField(obj, op, boundary)
            assertIsMember(boundary, {'w', 'e', 's', 'n'})
            assertIsMember(op, {'e'})

            switch op

                case 'e'
                    o = obj.(['e_scalar', '_', boundary]);
            end

        end

        % Returns the boundary operator T_ij (cell format) for the boundary specified by the string boundary.
        % Formula: tau_i = T_ij u_j
        % op -- string
        function T = getBoundaryTractionOperator(obj, boundary)
            assertIsMember(boundary, {'w', 'e', 's', 'n'})

            T = obj.(['T', '_', boundary]);
        end

        % Returns square boundary quadrature matrix, of dimension
        % corresponding to the number of boundary unknowns
        %
        % boundary -- string
        function H = getBoundaryQuadrature(obj, boundary)
            assertIsMember(boundary, {'w', 'e', 's', 'n'})

            H = obj.getBoundaryQuadratureForScalarField(boundary);
            I_dim = speye(obj.dim, obj.dim);
            H = kron(H, I_dim);
        end

        % Returns square boundary quadrature matrix, of dimension
        % corresponding to the number of boundary grid points
        %
        % boundary -- string
        function H_b = getBoundaryQuadratureForScalarField(obj, boundary)
            assertIsMember(boundary, {'w', 'e', 's', 'n'})

            H_b = obj.(['H_', boundary]);
        end

        function N = size(obj)
            N = obj.dim*prod(obj.m);
        end
    end
end
