classdef Burgers2D < scheme.Scheme
    properties
        grid % Physical grid
        order % Order accuracy for the approximation
        
        D % Non-stabilized scheme operator
        H % Discrete norm
        H_inv  % Norm inverse
        halfnorm_inv % Cell array halfnorm operators
        e_l % Cell array of left boundary operators
        e_r % Cell array of right boundary operators
        d_l % Cell array of left boundary derivative operators
        d_r % Cell array of right boundary derivative operators
    end

    methods
        function obj = Burgers2D(g, operator_type, order, dissipation)
            if ~isa(g, 'grid.Cartesian') || g.D() ~= 2
                error('Grid must be 2d cartesian');
            end

            obj.grid = g;
            obj.order = order;

            % Create cell array of 1D operators. For example D1_1d{1} = D1_x, D1_1d{2} = D1_y. 
            [Dp_1d, Dm_1d, H_1d, H_inv_1d, d_l_1d, d_r_1d, e_l_1d, e_r_1d, I, DissipationOp_1d] = ...
                obj.assemble1DOperators(g, operator_type, order, dissipation);
            
            %% 2D-operators
            % D1 
            D1_1d{1} = (Dp_1d{1} + Dp_1d{1})/2;
            D1_1d{2} = (Dp_1d{2} + Dp_1d{2})/2;
            D1_2d = obj.extendOperatorTo2D(D1_1d, I);
            D1 = D1_2d{1} + D1_2d{2}; 
            % D2

            Dp_2d = obj.extendOperatorTo2D(Dp_1d, I);
            Dm_2d = obj.extendOperatorTo2D(Dm_1d, I);
            D2 = @(viscosity) Dm_2d{1}*spdiag(viscosity)*Dp_2d{1} + Dm_2d{2}*spdiag(viscosity)*Dp_2d{2};
            % m = g.size();
            % ind = grid.funcToMatrix(g, 1:g.N());
            % for i = 1:g.D()
            %     D2_2d{i} = sparse(zeros(g.N()));
            % end
            % % x-direction
            % for i = 1:m(2)
            %     p = ind(:,i);
            %     D2_2d{1}(p,p) = @(viscosity) D2_1d{1}(viscosity(p));
            % end
            % % y-direction
            % for i = 1:m(1)
            %     p = ind(i,:);
            %     D2_2d{2}(p,p) = @(viscosity) D2_1d{2}(viscosity(p));
            % end
            % D2 = D2_2d{1} + D2_2d{2}; 

            obj.d_l = obj.extendOperatorTo2D(d_l_1d, I);
            obj.d_r = obj.extendOperatorTo2D(d_r_1d, I);
            obj.e_l = obj.extendOperatorTo2D(e_l_1d, I);
            obj.e_r = obj.extendOperatorTo2D(e_r_1d, I);
            obj.H = kron(H_1d{1},H_1d{2});
            obj.H_inv = kron(H_inv_1d{1},H_inv_1d{2});
            obj.halfnorm_inv = obj.extendOperatorTo2D(H_inv_1d, I);

            % Dissipation operator
            switch dissipation
                case 'on'
                    DissOp_2d = obj.extendOperatorTo2D(DissipationOp_1d, I);
                    DissOp = DissOp_2d{1} + DissOp_2d{2};
                    obj.D = @(v, viscosity) -1/2*D1*v.^2 + (D2(viscosity) + max(abs(v))*DissOp)*v;
                case 'off'
                    obj.D = @(v, viscosity) -1/2*D1*v.^2 + D2(viscosity)*v;
            end
        end

        % Closure functions return the operators applied to the own doamin 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 domain.
        %       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,data)
            default_arg('type','robin');
            default_arg('data',0);
            [e, d, halfnorm_inv, i_b, s] = obj.get_boundary_ops(boundary);
            switch type
                % Stable robin-like boundary conditions ((u+-abs(u))*u/3 - eps*u_x)) with +- at left/right boundary
                case {'R','robin'}
                    p = s*halfnorm_inv*e;
                    closure = @(v, viscosity) p*(((v(i_b)-s*abs(v(i_b)))/3).*(v(i_b)) - ((viscosity(i_b)).*d*v));
                    switch class(data)
                        case 'double'
                            penalty = s*p*data;
                        case 'function_handle'
                            penalty = @(t) s*p*data(t);
                        otherwise
                            error('Wierd data argument!')
                    end
                otherwise
                    error('No such boundary condition: type = %s',type);
            end
        end

        % Ruturns the boundary ops, half-norm, boundary indices and sign for the boundary specified by the string boundary.
        % The right boundary for each coordinate direction is considered the positive boundary
        function [e, d, halfnorm_inv, ind_boundary, s] = get_boundary_ops(obj,boundary)
            ind = grid.funcToMatrix(obj.grid, 1:obj.grid.N());
            switch boundary
                case 'w'
                    e = obj.e_l{1};
                    d = obj.d_l{1};
                    halfnorm_inv = obj.halfnorm_inv{1};
                    ind_boundary = ind(1,:);
                    s = -1;
                case 'e'
                    e = obj.e_r{1};
                    d = obj.d_r{1};
                    halfnorm_inv = obj.halfnorm_inv{1};
                    
                    ind_boundary = ind(end,:);
                    s = 1;
                case 's'
                    e = obj.e_l{2};
                    d = obj.d_l{2};
                    halfnorm_inv = obj.halfnorm_inv{2};
                    ind_boundary = ind(:,1);
                    s = -1;
                case 'n'
                    e = obj.e_r{2};
                    d = obj.d_r{2};
                    halfnorm_inv = obj.halfnorm_inv{2};
                    ind_boundary = ind(:,end);
                    s = 1;
                otherwise
                    error('No such boundary: boundary = %s',boundary);
            end
        end

        function [closure, penalty] = interface(obj,boundary,neighbour_scheme,neighbour_boundary)
            error('An interface function does not exist yet');
        end

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

    methods(Static)
        function [Dp, Dm, H, Hi, d_l, d_r, e_l, e_r, I, DissipationOp] = assemble1DOperators(g, operator_type, order, dissipation)
            dim = g.D(); 
            I = cell(dim,1);
            D1 = cell(dim,1);
            D2 = cell(dim,1);
            H = cell(dim,1);
            Hi = cell(dim,1);
            e_l = cell(dim,1);
            e_r = cell(dim,1);
            d1_l = cell(dim,1);
            d1_r = cell(dim,1);
            DissipationOp = cell(dim,1);
            for i = 1:dim
                switch operator_type
                    % case 'narrow'
                    %     ops = sbp.D4Variable(g.m(i), g.lim{i}, order);
                    %     D1{i} = ops.D1;
                    %     D2{i} = ops.D2;
                    %     d_l{i} = ops.d1_l';
                    %     d_r{i} = ops.d1_r';
                    %     if (strcmp(dissipation,'on'))
                    %         DissipationOp{i} = -1*sbp.dissipationOperator(g.m(i), order, ops.HI);
                    %     end
                    % case 'upwind-'
                    %     ops = sbp.D1Upwind(g.m(i), g.lim{i}, order);
                    %     D1{i} = (ops.Dp + ops.Dm)/2;
                    %     D2{i} = @(viscosity) ops.Dp*spdiag(viscosity)*ops.Dm;
                    %     d_l{i} = ops.e_l'*ops.Dm;
                    %     d_r{i} = ops.e_r'*ops.Dm;
                    %     if (strcmp(dissipation,'on'))
                    %         DissipationOp{i} = (ops.Dp-ops.Dm)/2;
                    %     end
                    case 'upwind+'
                        ops = sbp.D1Upwind(g.m(i), g.lim{i}, order);
                        Dp{i} = ops.Dp;
                        Dm{i} = ops.Dm;
                        % D1{i} = (ops.Dp + ops.Dm)/2;
                        % D2{i} = @(viscosity) ops.Dm*spdiag(viscosity)*ops.Dp;
                        d_l{i} = ops.e_l'*ops.Dp;
                        d_r{i} = ops.e_r'*ops.Dp;
                        if (strcmp(dissipation,'on'))
                            DissipationOp{i} = (ops.Dp-ops.Dm)/2;
                        end
                    % case 'upwind+-'
                    %     ops = sbp.D1Upwind(g.m(i), g.lim{i}, order);
                    %     D1{i} = (ops.Dp + ops.Dm)/2;
                    %     D2{i} = @(viscosity) (ops.Dp*spdiag(viscosity)*ops.Dm + ops.Dm*spdiag(viscosity)*ops.Dp)/2;
                    %     d_l{i} = ops.e_l'*D1;
                    %     d_r{i} = ops.e_r'*D1;
                    %     if (strcmp(dissipation,'on'))
                    %         DissipationOp{i} = (ops.Dp-ops.Dm)/2;
                    %     end
                    otherwise
                        error('Other operator types not yet supported', operator_type);
                end
                H{i} = ops.H;
                Hi{i} = ops.HI;
                e_l{i} = ops.e_l;
                e_r{i} = ops.e_r;
                I{i} = speye(g.m(i));
            end
        end
        function op_2d = extendOperatorTo2D(op, I)
            op_2d{1} = kr(op{1}, I{2});
            op_2d{2} = kr(I{1}, op{2});        
        end
    end
end