classdef ViscoElastic2d < scheme.Scheme

% Discretizes the visco-elastic wave equation in curvilinear coordinates.
% 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
        % J, Ji
        RHO % Density
        C   % Elastic stiffness tensor
        ETA % Effective viscosity, used in strain rate eq

        D % Total operator
        Delastic        % Elastic operator (momentum balance)
        Dviscous        % Acts on viscous strains in momentum balance
        DstrainRate     % Acts on u and gamma, returns strain rate gamma_t

        D1, D1Tilde % Physical derivatives
        sigma % Cell matrix of physical stress operators

        % Inner products
        H

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

        % Restriction operators
        Eu, Egamma  % Pick out all components of u/gamma
        eU, eGamma  % Pick out one specific component

        % Bundary restriction ops
        e_scalar_w, e_scalar_e, e_scalar_s, e_scalar_n

        n_w, n_e, n_s, n_n % Physical normals
        tangent_w, tangent_e, tangent_s, tangent_n % Physical tangents

        elasticObj

    end

    methods

        % The coefficients can either be function handles or grid functions
        function obj = ViscoElastic2d(g, order, rho, C, eta)
            default_arg('rho', @(x,y) 0*x+1);
            default_arg('eta', @(x,y) 0*x+1);
            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 ;
                        end
                    end
                end
            end
            default_arg('C', C_default);

            assert(isa(g, 'grid.Curvilinear'));

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

            if isa(eta, 'function_handle')
                eta = grid.evalOn(g, eta);
            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;

            elasticObj = scheme.Elastic2dCurvilinearAnisotropic(g, order, rho, C);

            % Construct a pair of first derivatives
            K = elasticObj.K;
            for i = 1:dim
                for j = 1:dim
                    K{i,j} = spdiag(K{i,j});
                end
            end
            J = elasticObj.J;
            Ji = elasticObj.Ji;
            D_ref = elasticObj.refObj.D1;
            D1 = cell(dim, 1);
            D1Tilde = cell(dim, 1);
            for i = 1:dim
                D1{i} = 0*D_ref{i};
                D1Tilde{i} = 0*D_ref{i};
                for j = 1:dim
                    D1{i} = D1{i} + K{i,j}*D_ref{j};
                    D1Tilde{i} = D1Tilde{i} + Ji*D_ref{j}*J*K{i,j};
                end
            end
            obj.D1 = D1;
            obj.D1Tilde = D1Tilde;

            eU = elasticObj.E;

            % Storage order for gamma: 11-12-21-22.
            I = speye(g.N(), g.N());
            eGamma = cell(dim, dim);
            e = cell(dim, dim);
            e{1,1} = [1;0;0;0];
            e{1,2} = [0;1;0;0];
            e{2,1} = [0;0;1;0];
            e{2,2} = [0;0;0;1];
            for i = 1:dim
                for j = 1:dim
                    eGamma{i,j} = kron(I, e{i,j});
                end
            end

            % Store u first, then gamma
            mU = dim*g.N();
            mGamma = dim^2*g.N();
            Iu = speye(mU, mU);
            Igamma = speye(mGamma, mGamma);

            Eu = cell2mat({Iu, sparse(mU, mGamma)})';
            Egamma = cell2mat({sparse(mGamma, mU), Igamma})';

            for i = 1:dim
                eU{i} = Eu*eU{i};
            end
            for i = 1:dim
                for j = 1:dim
                    eGamma{i,j} = Egamma*eGamma{i,j};
                end
            end

            obj.eGamma = eGamma;
            obj.eU = eU;
            obj.Egamma = Egamma;
            obj.Eu = Eu;

            % Build stress operator
            sigma = cell(dim, dim);
            C = obj.C;
            for i = 1:dim
                for j = 1:dim
                    sigma{i,j} = spalloc(g.N(), (dim^2 + dim)*g.N(), order^2*g.N());
                    for k = 1:dim
                        for l = 1:dim
                            sigma{i,j} = sigma{i,j} + C{i,j,k,l}*(D1{k}*eU{l}' - eGamma{k,l}');
                        end
                    end
                end
            end

            % Elastic operator
            Delastic = Eu*elasticObj.D*Eu';

            % Add viscous strains to momentum balance
            RHOi = spdiag(1./rho);
            Dviscous = spalloc((dim^2 + dim)*g.N(), (dim^2 + dim)*g.N(), order^2*(dim^2 + dim)*g.N());
            for i = 1:dim
                for j = 1:dim
                    for k = 1:dim
                        for l = 1:dim
                            Dviscous = Dviscous - eU{j}*RHOi*D1Tilde{i}*C{i,j,k,l}*eGamma{k,l}';
                        end
                    end
                end
            end

            ETA = spdiag(eta);
            DstrainRate = 0*Delastic;
            for i = 1:dim
                for j = 1:dim
                    DstrainRate = DstrainRate + eGamma{i,j}*(ETA\sigma{i,j});
                end
            end

            obj.D = Delastic + Dviscous + DstrainRate;
            obj.Delastic = Delastic;
            obj.Dviscous = Dviscous;
            obj.DstrainRate = DstrainRate;
            obj.sigma = sigma;

            %---- Set remaining object properties ------
            obj.RHO = elasticObj.RHO;
            obj.ETA = ETA;
            obj.H = elasticObj.H;

            obj.n_w = elasticObj.n_w;
            obj.n_e = elasticObj.n_e;
            obj.n_s = elasticObj.n_s;
            obj.n_n = elasticObj.n_n;

            obj.tangent_w = elasticObj.tangent_w;
            obj.tangent_e = elasticObj.tangent_e;
            obj.tangent_s = elasticObj.tangent_s;
            obj.tangent_n = elasticObj.tangent_n;

            obj.H_w = elasticObj.H_w;
            obj.H_e = elasticObj.H_e;
            obj.H_s = elasticObj.H_s;
            obj.H_n = elasticObj.H_n;

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

            % Misc.
            obj.elasticObj = elasticObj;
            obj.m = elasticObj.m;
            obj.h = elasticObj.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.

        % For displacement bc:
        % bc = {comp, 'd', dComps},
        % where
        % dComps = vector of components with displacement BC. Default: 1:dim.
        % In this way, we can specify one BC at a time even though the SATs depend on all BC.
        function [closure, penalty] = boundary_condition(obj, boundary, bc, tuning)
            default_arg('tuning', 1.0);
            assert( iscell(bc), 'The BC type must be a 2x1 or 3x1 cell array' );

            component = bc{1};
            type = bc{2};
            dim = obj.dim;

            n       = obj.getNormal(boundary);
            H_gamma = obj.getBoundaryQuadratureForScalarField(boundary);
            e       = obj.getBoundaryOperatorForScalarField('e', boundary);

            H       = obj.H;
            RHO     = obj.RHO;
            ETA     = obj.ETA;
            C       = obj.C;
            Eu      = obj.Eu;
            eU      = obj.eU;
            eGamma  = obj.eGamma;

            switch type
            case {'F','f','Free','free','traction','Traction','t','T'}

                % Get elastic closure and penalty
                [closure, penalty] = obj.elasticObj.boundary_condition(boundary, bc, tuning);
                closure = Eu*closure*Eu';
                penalty = Eu*penalty;

                switch component
                case 't'
                    dir = obj.getTangent(boundary);
                case 'n'
                    dir = obj.getNormal(boundary);
                case 1
                    dir = {1, 0};
                case 2
                    dir = {0, 1};
                end

                % Add viscous part of closure
                for j = 1:dim
                    for i = 1:dim
                        for k = 1:dim
                            for l = 1:dim
                                closure = closure + eU{j}*( (RHO*H)\(C{i,j,k,l}*e*dir{j}^2*H_gamma*n{i}*e'*eGamma{k,l}') );
                            end
                        end
                    end
                end

            case {'D','d','dirichlet','Dirichlet','displacement','Displacement'}

                % Get elastic closure and penalty
                [closure, penalty] = obj.elasticObj.boundary_condition(boundary, bc, tuning);
                closure = Eu*closure*Eu';
                penalty = Eu*penalty;

                % Add penalty to strain rate eq
                l = component;
                for i = 1:dim
                    for j = 1:dim
                        for k = 1:dim
                            closure = closure - eGamma{i,j}*( (H*ETA)\(C{i,j,k,l}*e*H_gamma*n{k}*e'*eU{l}') );
                            penalty = penalty + eGamma{i,j}*( (H*ETA)\(C{i,j,k,l}*e*H_gamma*n{k}) );
                        end
                    end
                end

            end

        end

        function [closure, penalty] = displacementBCNormalTangential(obj, boundary, bc, tuning)
            disp('WARNING: displacementBCNormalTangential is only guaranteed to work for displacement BC on one component and traction on the other');
            u = obj;

            component = bc{1};
            type = bc{2};

            switch component
            case 'n'
                en      = u.getBoundaryOperator('en', boundary);
                tau_n   = u.getBoundaryOperator('tau_n', boundary);
                N       = u.getNormal(boundary);
            case 't'
                en      = u.getBoundaryOperator('et', boundary);
                tau_n   = u.getBoundaryOperator('tau_t', boundary);
                N       = u.getTangent(boundary);
            end

            % Operators
            e       = u.getBoundaryOperatorForScalarField('e', boundary);
            h11     = u.getBorrowing(boundary);
            n      = u.getNormal(boundary);

            C = u.C;
            Ji = u.Ji;
            s = spdiag(u.(['s_' boundary]));
            m_tot = u.grid.N();

            Hi      = u.E{1}*inv(u.H)*u.E{1}' + u.E{2}*inv(u.H)*u.E{2}';
            RHOi    = u.E{1}*inv(u.RHO)*u.E{1}' + u.E{2}*inv(u.RHO)*u.E{2}';

            H_gamma         = u.getBoundaryQuadratureForScalarField(boundary);
            dim             = u.dim;

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

            % Term 1: The symmetric term
            Z = sparse(m_int, m_int);
            for i = 1:dim
                for j = 1:dim
                    for l = 1:dim
                        for k = 1:dim
                            Z = Z + n{i}*N{j}*e'*Ji*C{i,j,k,l}*e*n{k}*N{l};
                        end
                    end
                end
            end

            Z = -tuning*dim*1/h11*s*Z;
            closure = closure + en*H_gamma*Z*en';
            penalty = penalty - en*H_gamma*Z;

            % Term 2: The symmetrizing term
            closure = closure + tau_n*H_gamma*en';
            penalty = penalty - tau_n*H_gamma;

            % Multiply all normal component terms by inverse of density x quadraure
            closure = RHOi*Hi*closure;
            penalty = RHOi*Hi*penalty;
        end

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

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

            switch type.type
            case 'standard'
                [closure, penalty] = obj.refObj.interface(boundary,neighbour_scheme.refObj,neighbour_boundary,type);
            case 'frictionalFault'
                [closure, penalty] = obj.interfaceFrictionalFault(boundary,neighbour_scheme,neighbour_boundary,type);
            end

        end

        function [closure, penalty] = interfaceFrictionalFault(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

            u = obj;
            v = neighbour_scheme;

            % Operators, u side
            e_u       = u.getBoundaryOperatorForScalarField('e', boundary);
            en_u       = u.getBoundaryOperator('en', boundary);
            tau_n_u     = u.getBoundaryOperator('tau_n', boundary);
            h11_u     = u.getBorrowing(boundary);
            n_u      = u.getNormal(boundary);

            C_u = u.C;
            Ji_u = u.Ji;
            s_u = spdiag(u.(['s_' boundary]));
            m_tot_u = u.grid.N();

            % Operators, v side
            e_v       = v.getBoundaryOperatorForScalarField('e', neighbour_boundary);
            en_v       = v.getBoundaryOperator('en', neighbour_boundary);
            tau_n_v     = v.getBoundaryOperator('tau_n', neighbour_boundary);
            h11_v     = v.getBorrowing(neighbour_boundary);
            n_v      = v.getNormal(neighbour_boundary);

            C_v = v.C;
            Ji_v = v.Ji;
            s_v = spdiag(v.(['s_' neighbour_boundary]));
            m_tot_v = v.grid.N();

            % Operators that are only required for own domain
            Hi      = u.E{1}*inv(u.H)*u.E{1}' + u.E{2}*inv(u.H)*u.E{2}';
            RHOi    = u.E{1}*inv(u.RHO)*u.E{1}' + u.E{2}*inv(u.RHO)*u.E{2}';

            % Shared operators
            H_gamma         = u.getBoundaryQuadratureForScalarField(boundary);
            dim             = u.dim;

            % Preallocate
            [~, m_int] = size(H_gamma);
            closure = sparse(dim*m_tot_u, dim*m_tot_u);
            penalty = sparse(dim*m_tot_u, dim*m_tot_v);

            % Continuity of normal displacement, term 1: The symmetric term
            Z_u = sparse(m_int, m_int);
            Z_v = sparse(m_int, m_int);
            for i = 1:dim
                for j = 1:dim
                    for l = 1:dim
                        for k = 1:dim
                            Z_u = Z_u + n_u{i}*n_u{j}*e_u'*Ji_u*C_u{i,j,k,l}*e_u*n_u{k}*n_u{l};
                            Z_v = Z_v + n_v{i}*n_v{j}*e_v'*Ji_v*C_v{i,j,k,l}*e_v*n_v{k}*n_v{l};
                        end
                    end
                end
            end

            Z = -tuning*dim*( 1/(4*h11_u)*s_u*Z_u + 1/(4*h11_v)*s_v*Z_v );
            closure = closure + en_u*H_gamma*Z*en_u';
            penalty = penalty + en_u*H_gamma*Z*en_v';

            % Continuity of normal displacement, term 2: The symmetrizing term
            closure = closure + 1/2*tau_n_u*H_gamma*en_u';
            penalty = penalty + 1/2*tau_n_u*H_gamma*en_v';

            % Continuity of normal traction
            closure = closure - 1/2*en_u*H_gamma*tau_n_u';
            penalty = penalty + 1/2*en_u*H_gamma*tau_n_v';

            % Multiply all normal component terms by inverse of density x quadraure
            closure = RHOi*Hi*closure;
            penalty = RHOi*Hi*penalty;

            % ---- Tangential tractions are imposed just like traction BC ------
            closure = closure + obj.boundary_condition(boundary, {'t', 't'});

        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.refObj.h11{1};
            case {'s', 'n'}
                h11 = obj.refObj.h11{2};
            end
        end

        % Returns the outward unit normal vector for the boundary specified by the string boundary.
        % n is a cell of diagonal matrices for each normal component, n{1} = n_1, n{2} = n_2.
        function n = getNormal(obj, boundary)
            assertIsMember(boundary, {'w', 'e', 's', 'n'})

            n = obj.(['n_' boundary]);
        end

        % Returns the unit tangent vector for the boundary specified by the string boundary.
        % t is a cell of diagonal matrices for each normal component, t{1} = t_1, t{2} = t_2.
        function t = getTangent(obj, boundary)
            assertIsMember(boundary, {'w', 'e', 's', 'n'})

            t = obj.(['tangent_' boundary]);
        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', 'en', 'et', 'tau_n', 'tau_t'})

            o = obj.([op, '_', boundary]);

        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 + obj.dim^2)*prod(obj.m);
        end
    end
end