classdef Hypsyst3d < scheme.Scheme
    properties
        grid
        order % Order accuracy for the approximation

        nEquations

        D % non-stabilized scheme operator
        A, B, C, E % Symbolic coefficient matrices
        Aevaluated,Bevaluated,Cevaluated, Eevaluated

        H % Discrete norm
        Hx, Hy, Hz  % Norms in the x, y and z directions
        Hxi,Hyi, Hzi % Kroneckerd norms. 1'*Hx*v corresponds to integration in the x dir.
        e_w, e_e, e_s, e_n, e_b, e_t
        params % Parameters for the coefficient matrices
    end


    methods
        % Solving Hyperbolic systems on the form u_t=-Au_x-Bu_y-Cu_z-Eu
        function obj = Hypsyst3d(cartesianGrid, order, A, B, C, E, params, opSet)
            assertType(cartesianGrid, 'grid.Cartesian');
            default_arg('E', [])
            default_arg('opSet', @sbp.D2Standard)
            obj.grid = cartesianGrid;

            obj.A = A;
            obj.B = B;
            obj.C = C;
            obj.E = E;
            obj.params = params;
            assert(obj.grid.d == 3, 'Dimensions not correct')
            % Construct 1D operators for each coordinate dir
            for i = 1:obj.grid.d
                ops{i} = opSet(obj.grid.m(i),obj.grid.lim{i},order);
                I{i} = speye(obj.grid.m(i));
            end
            x = obj.grid.x; %Grid points in each coordinate dir
            X = num2cell(obj.grid.points(),1); %Kroneckered grid values in each coordinate dir

            obj.Aevaluated = obj.evaluateCoefficientMatrix(A, X{:});
            obj.Bevaluated = obj.evaluateCoefficientMatrix(B, X{:});
            obj.Cevaluated = obj.evaluateCoefficientMatrix(C, X{:});
            obj.Eevaluated = obj.evaluateCoefficientMatrix(E, X{:});

            obj.nEquations = length(A(obj.params,0,0,0));
            I_n = speye(obj.nEquations);
            I_N = kr(I_n,I{:});

            obj.Hxi = kr(I_n, ops{1}.HI, I{2}, I{3});
            obj.Hx = ops{1}.H;
            obj.Hyi = kr(I_n, I{1}, ops{2}.HI, I{3});
            obj.Hy = ops{2}.H;
            obj.Hzi = kr(I_n, I{1},I{2}, ops{3}.HI);
            obj.Hz = ops{3}.H;

            obj.e_w = kr(I_n, ops{1}.e_l, I{2}, I{3});
            obj.e_e = kr(I_n, ops{1}.e_r, I{2}, I{3});
            obj.e_s = kr(I_n, I{1}, ops{2}.e_l, I{3});
            obj.e_n = kr(I_n, I{1}, ops{2}.e_r, I{3});
            obj.e_b = kr(I_n, I{1}, I{2}, ops{3}.e_l);
            obj.e_t = kr(I_n, I{1}, I{2}, ops{3}.e_r);

            obj.order = order;

            switch toString(opSet)
                case 'sbp.D1Upwind'
                    alphaA = max(abs(eig(A(params, x{1}(end), x{2}(end), x{3}(end)))));
                    alphaB = max(abs(eig(B(params, x{1}(end), x{2}(end), x{3}(end)))));
                    alphaC = max(abs(eig(C(params, x{1}(end), x{2}(end), x{3}(end)))));

                    Ap = (obj.Aevaluated+alphaA*I_N)/2;
                    Am = (obj.Aevaluated-alphaA*I_N)/2;
                    Dpx = kr(I_n, ops{1}.Dp, I{2}, I{3});
                    Dmx = kr(I_n, ops{1}.Dm, I{2}, I{3});
                    obj.D = -Am*Dpx;
                    temp = Ap*Dmx;
                    obj.D = obj.D-temp;
                    clear Ap Am Dpx Dmx

                    Bp = (obj.Bevaluated+alphaB*I_N)/2;
                    Bm = (obj.Bevaluated-alphaB*I_N)/2;
                    Dpy = kr(I_n, I{1}, ops{2}.Dp, I{3});
                    Dmy = kr(I_n, I{1}, ops{2}.Dm, I{3});
                    temp = Bm*Dpy;
                    obj.D = obj.D-temp;
                    temp = Bp*Dmy;
                    obj.D = obj.D-temp;
                    clear Bp Bm Dpy Dmy


                    Cp = (obj.Cevaluated+alphaC*I_N)/2;
                    Cm = (obj.Cevaluated-alphaC*I_N)/2;
                    Dpz = kr(I_n, I{1}, I{2}, ops{3}.Dp);
                    Dmz = kr(I_n, I{1}, I{2}, ops{3}.Dm);

                    temp = Cm*Dpz;
                    obj.D = obj.D-temp;
                    temp = Cp*Dmz;
                    obj.D = obj.D-temp;
                    clear Cp Cm Dpz Dmz
                    obj.D = obj.D-obj.Eevaluated;

                case {'sbp.D2Standard', 'sbp.D2Variable', 'sbp.D4Standard', 'sbp.D4Variable'}
                    D1_x = kr(I_n, ops{1}.D1, I{2}, I{3});
                    D1_y = kr(I_n, I{1}, ops{2}.D1, I{3});
                    D1_z = kr(I_n, I{1}, I{2}, ops{3}.D1);
                    obj.D = -obj.Aevaluated*D1_x-obj.Bevaluated*D1_y-obj.Cevaluated*D1_z-obj.Eevaluated;
                otherwise
                    error('Operator not supported');
            end
        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'.
        %       type                is a string specifying the type of boundary condition if there are several.
        %       data                is a function returning the data that should be applied at the boundary.
        function [closure, penalty] = boundary_condition(obj,boundary,type,L)
            default_arg('type','char');
            BM = boundary_matrices(obj,boundary);
            switch type
                case{'c','char'}
                    [closure,penalty] = boundary_condition_char(obj,BM);
                case{'general'}
                    [closure,penalty] = boundary_condition_general(obj,BM,boundary,L);
                otherwise
                    error('No such boundary condition')
            end
        end

        function [closure, penalty] = interface(obj, boundary, neighbour_scheme, neighbour_boundary, type)
            error('Not implemented');
        end

        % TODO: Implement! This function should potentially replace boundary_matrices.
        % Returns the boundary operator op for the boundary specified by the string boundary.
        % op        -- string or a cell array of strings
        % boundary  -- string
        function varargout = getBoundaryOperator(obj, op, boundary)
            error('Not implemented');
        end

        % TODO: Implement!
        % Returns square boundary quadrature matrix, of dimension
        % corresponding to the number of boundary points
        %
        % boundary -- string
        function H = getBoundaryQuadrature(obj, boundary)
            error('Not implemented');
        end

        function N = size(obj)
            N = obj.grid.m;
        end

        function [ret] = evaluateCoefficientMatrix(obj, mat, X, Y, Z)
            params = obj.params;
            side = max(length(X),length(Y));
            if isa(mat,'function_handle')
                [rows,cols] = size(mat(params,0,0,0));
                matVec = mat(params,X',Y',Z');
                matVec = sparse(matVec);
            else
                matVec = mat;
                [rows,cols] = size(matVec);
                side = max(length(X),length(Y));
                cols = cols/side;
            end

            ret = cell(rows,cols);
            for ii = 1:rows
                for jj = 1:cols
                    ret{ii,jj} = diag(matVec(ii,(jj-1)*side+1:jj*side));
                end
            end
            ret = cell2mat(ret);
        end

        function [BM] = boundary_matrices(obj,boundary)
            params = obj.params;
            bp = num2cell(obj.grid.getBoundary(boundary),1); %Kroneckered boundary points
            switch boundary
                case {'w'}
                    BM.e_ = obj.e_w;
                    mat = obj.A;
                    BM.boundpos = 'l';
                    BM.Hi = obj.Hxi;
                    BM.side = length(bp{2});
                case {'e'}
                    BM.e_ = obj.e_e;
                    mat = obj.A;
                    BM.boundpos = 'r';
                    BM.Hi = obj.Hxi;
                    BM.side = length(bp{2});
                case {'s'}
                    BM.e_ = obj.e_s;
                    mat = obj.B;
                    BM.boundpos = 'l';
                    BM.Hi = obj.Hyi;
                    BM.side = length(bp{1});
                case {'n'}
                    BM.e_ = obj.e_n;
                    mat = obj.B;
                    BM.boundpos = 'r';
                    BM.Hi = obj.Hyi;
                    BM.side = length(bp{1});
                case{'d'}
                    BM.e_ = obj.e_b;
                    mat = obj.C;
                    BM.boundpos = 'l';
                    BM.Hi = obj.Hzi;
                    BM.side = length(bp{1});
                case{'u'}
                    BM.e_ = obj.e_t;
                    mat = obj.C;
                    BM.boundpos = 'r';
                    BM.Hi = obj.Hzi;
                    BM.side = length(bp{1});
            end
            [BM.V,BM.Vi,BM.D,signVec] = obj.diagonalizeMatrix(mat,bp{:});
            BM.pos = signVec(1); BM.zeroval=signVec(2); BM.neg=signVec(3);
        end

        % Characteristic boundary conditions
        function [closure, penalty]=boundary_condition_char(obj,BM)
            side = BM.side;
            pos = BM.pos;
            neg = BM.neg;
            zeroval=BM.zeroval;
            V = BM.V;
            Vi = BM.Vi;
            Hi = BM.Hi;
            D = BM.D;
            e_ = BM.e_;

            switch BM.boundpos
                case {'l'}
                    tau = sparse(obj.nEquations*side,pos);
                    Vi_plus = Vi(1:pos,:);
                    tau(1:pos,:) = -abs(D(1:pos,1:pos));
                    closure = Hi*e_*V*tau*Vi_plus*e_';
                    penalty = -Hi*e_*V*tau*Vi_plus;
                case {'r'}
                    tau = sparse(obj.nEquations*side,neg);
                    tau((pos+zeroval)+1:obj.nEquations*side,:) = -abs(D((pos+zeroval)+1:obj.nEquations*side,(pos+zeroval)+1:obj.nEquations*side));
                    Vi_minus = Vi((pos+zeroval)+1:obj.nEquations*side,:);
                    closure = Hi*e_*V*tau*Vi_minus*e_';
                    penalty = -Hi*e_*V*tau*Vi_minus;
            end
        end

        % General boundary condition in the form Lu=g(x)
        function [closure,penalty] = boundary_condition_general(obj,BM,boundary,L)
            side = BM.side;
            pos = BM.pos;
            neg = BM.neg;
            zeroval=BM.zeroval;
            V = BM.V;
            Vi = BM.Vi;
            Hi = BM.Hi;
            D = BM.D;
            e_ = BM.e_;

            bp = num2cell(obj.grid.getBoundary(boundary),1); %Kroneckered boundary points
            L = obj.evaluateCoefficientMatrix(L,bp{:}); % General boundary condition in the form Lu=g(x)

            switch BM.boundpos
                case {'l'}
                    tau = sparse(obj.nEquations*side,pos);
                    Vi_plus = Vi(1:pos,:);
                    Vi_minus = Vi(pos+zeroval+1:obj.nEquations*side,:);
                    V_plus = V(:,1:pos);
                    V_minus = V(:,(pos+zeroval)+1:obj.nEquations*side);

                    tau(1:pos,:) = -abs(D(1:pos,1:pos));
                    R = -inv(L*V_plus)*(L*V_minus);
                    closure = Hi*e_*V*tau*(Vi_plus-R*Vi_minus)*e_';
                    penalty = -Hi*e_*V*tau*inv(L*V_plus)*L;
                case {'r'}
                    tau = sparse(obj.nEquations*side,neg);
                    tau((pos+zeroval)+1:obj.nEquations*side,:) = -abs(D((pos+zeroval)+1:obj.nEquations*side,(pos+zeroval)+1:obj.nEquations*side));
                    Vi_plus = Vi(1:pos,:);
                    Vi_minus = Vi((pos+zeroval)+1:obj.nEquations*side,:);

                    V_plus = V(:,1:pos);
                    V_minus = V(:,(pos+zeroval)+1:obj.nEquations*side);
                    R = -inv(L*V_minus)*(L*V_plus);
                    closure = Hi*e_*V*tau*(Vi_minus-R*Vi_plus)*e_';
                    penalty = -Hi*e_*V*tau*inv(L*V_minus)*L;
            end
        end

        % Function that diagonalizes a symbolic matrix A as A=V*D*Vi
        % D         is a diagonal matrix with the eigenvalues on A on the diagonal sorted by their sign
        %                                    [d+       ]
        %                               D =  [   d0    ]
        %                                    [       d-]
        % signVec   is a vector specifying the number of positive, zero and negative eigenvalues of D
        function [V, Vi, D, signVec] = diagonalizeMatrix(obj, mat, x, y, z)
            params = obj.params;
            syms xs ys zs
            [V, D] = eig(mat(params,xs,ys,zs));
            Vi=inv(V);
            xs = x;
            ys = y;
            zs = z;


            side = max(length(x),length(y));
            Dret = zeros(obj.nEquations,side*obj.nEquations);
            Vret = zeros(obj.nEquations,side*obj.nEquations);
            Viret= zeros(obj.nEquations,side*obj.nEquations);

            for ii=1:obj.nEquations
                for jj=1:obj.nEquations
                    Dret(jj,(ii-1)*side+1:side*ii) = eval(D(jj,ii));
                    Vret(jj,(ii-1)*side+1:side*ii) = eval(V(jj,ii));
                    Viret(jj,(ii-1)*side+1:side*ii) = eval(Vi(jj,ii));
                end
            end

            D = sparse(Dret);
            V = sparse(Vret);
            Vi = sparse(Viret);
            V = obj.evaluateCoefficientMatrix(V,x,y,z);
            Vi= obj.evaluateCoefficientMatrix(Vi,x,y,z);
            D = obj.evaluateCoefficientMatrix(D,x,y,z);
            DD = diag(D);

            poseig = (DD>0);
            zeroeig = (DD==0);
            negeig = (DD<0);

            D = diag([DD(poseig); DD(zeroeig); DD(negeig)]);
            V = [V(:,poseig) V(:,zeroeig) V(:,negeig)];
            Vi= [Vi(poseig,:); Vi(zeroeig,:); Vi(negeig,:)];
            signVec = [sum(poseig),sum(zeroeig),sum(negeig)];
        end
    end
end