classdef Staggered1DAcousticsVariable < scheme.Scheme
   properties
        m % Number of points of primal grid in each direction, possibly a vector
        h % Grid spacing

        % Grids
        grid % Total grid object
        grid_primal
        grid_dual

        order % Order accuracy for the approximation

        H  % Combined norm
        Hi % Inverse
        D % Semi-discrete approximation matrix

        D1_primal
        D1_dual

        % Pick out left or right boundary
        e_l
        e_r

        % Initial data
        v0

        % Pick out primal or dual component
        e_primal
        e_dual

        % System matrices, cell matrices of function handles
        A
        B

        % Variable coefficient matrices evaluated at boundaries.
        A_l, A_r
        B_l, B_r
    end


    methods
         % Scheme for A*u_t + B u_x = 0,
         % u = [p, v];
         % A: Diagonal and A > 0,
         % A = [a1, 0;
         %      0,  a2]
         %
         % B = B^T and with diagonal entries = 0.
         % B = [0 b
         %     b 0]
         % Here we store p on the primal grid and v on the dual
         % A and B are cell matrices of function handles
         function obj = Staggered1DAcousticsVariable(g, order, A, B)
            default_arg('B',{@(x)0*x, @(x)0*x+1; @(x)0*x+1, @(x)0*x});
            default_arg('A',{@(x)0*x+1, @(x)0*x; @(x)0*x, @(x)0*x+1});

            obj.order = order;
            obj.A = A;
            obj.B = B;

            % Grids
            obj.m = g.size();
            xl = g.getBoundary('l');
            xr = g.getBoundary('r');
            xlim = {xl, xr};
            obj.grid = g;
            obj.grid_primal = g.grids{1};
            obj.grid_dual = g.grids{2};
            x_primal = obj.grid_primal.points();
            x_dual = obj.grid_dual.points();

            % If coefficients are function handles, evaluate them
            if isa(A{1,1}, 'function_handle')
                A{1,1} = A{1,1}(x_primal);
                A{1,2} = A{1,2}(x_dual);
                A{2,1} = A{2,1}(x_primal);
                A{2,2} = A{2,2}(x_dual);
            end

            if isa(B{1,1}, 'function_handle')
                B{1,1} = B{1,1}(x_primal);
                B{1,2} = B{1,2}(x_primal);
                B{2,1} = B{2,1}(x_dual);
                B{2,2} = B{2,2}(x_dual);
            end

            % If coefficents are vectors, turn them into diagonal matrices
            [m, n] = size(A{1,1});
            if m==1 || n == 1
                A{1,1} = spdiag(A{1,1}, 0);
                A{2,1} = spdiag(A{2,1}, 0);
                A{1,2} = spdiag(A{1,2}, 0);
                A{2,2} = spdiag(A{2,2}, 0);
            end

            [m, n] = size(B{1,1});
            if m==1 || n == 1
                B{1,1} = spdiag(B{1,1}, 0);
                B{2,1} = spdiag(B{2,1}, 0);
                B{1,2} = spdiag(B{1,2}, 0);
                B{2,2} = spdiag(B{2,2}, 0);
            end

            % Boundary matrices
            obj.A_l = full([A{1,1}(1,1), A{1,2}(1,1);....
                       A{2,1}(1,1), A{2,2}(1,1)]);
            obj.A_r = full([A{1,1}(end,end), A{1,2}(end,end);....
                       A{2,1}(end,end), A{2,2}(end,end)]);
            obj.B_l = full([B{1,1}(1,1), B{1,2}(1,1);....
                       B{2,1}(1,1), B{2,2}(1,1)]);
            obj.B_r = full([B{1,1}(end,end), B{1,2}(end,end);....
                       B{2,1}(end,end), B{2,2}(end,end)]);

            % Get operators
            ops = sbp.D1StaggeredUpwind(obj.m, xlim, order);
            obj.h = ops.h;

            % Build combined operators
            H_primal = ops.H_primal;
            H_dual = ops.H_dual;
            obj.H = blockmatrix.toMatrix( {H_primal, []; [], H_dual } );
            obj.Hi = inv(obj.H);

            D1_primal = ops.D1_primal;
            D1_dual = ops.D1_dual;
            A11_B12 = -A{1,1}\B{1,2};
            A22_B21 = -A{2,2}\B{2,1};
            D = {[],                A11_B12*D1_primal;...
                 A22_B21*D1_dual,     []};
            obj.D = blockmatrix.toMatrix(D);
            obj.D1_primal = D1_primal;
            obj.D1_dual = D1_dual;

            % Combined boundary operators
            e_primal_l = ops.e_primal_l;
            e_primal_r = ops.e_primal_r;
            e_dual_l = ops.e_dual_l;
            e_dual_r = ops.e_dual_r;
            e_l = {e_primal_l, [];...
                   []       ,  e_dual_l};
            e_r = {e_primal_r, [];...
                   []       ,  e_dual_r};
            obj.e_l = blockmatrix.toMatrix(e_l);
            obj.e_r = blockmatrix.toMatrix(e_r);

            % Pick out first or second component of solution
            N_primal = obj.grid_primal.N();
            N_dual = obj.grid_dual.N();
            obj.e_primal = [speye(N_primal, N_primal); sparse(N_dual, N_primal)];
            obj.e_dual = [sparse(N_primal, N_dual); speye(N_dual, N_dual)];


        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 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.
        %       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, type)

            default_arg('type','p');

            % type = 'p' => boundary condition for p
            % type = 'v' => boundary condition for v
            % No other types implemented yet

            % BC on the form Lu - g = 0;

            % Need a transformation T such that w = T^{−1}*u is a
            % meaningful change of variables and T^T*B*T is block-diagonal,
            % For linear acoustics, T = T_C meets these criteria.

            % Get boundary matrices
            switch boundary
                case 'l'
                    A = obj.A_l;
                    B = obj.B_l;
                case 'r'
                    A = obj.A_r;
                    B = obj.B_r;
            end
            C = inv(A)*B;

            % Diagonalize C and use T_C to diagonalize B
            [T, ~] = eig(C);
            Lambda = T'*B*T;
            lambda = diag(Lambda);

            % Identify in- and outgoing characteristic variables
            Iplus = lambda > 0;
            Iminus = lambda < 0;

            switch boundary
            case 'l'
                Iin = Iplus;
            case 'r'
                Iin = Iminus;
            end
            Tin = T(:,Iin);

            switch type
            case 'p'
                L = [1, 0];
            case 'v'
                L = [0, 1];
            case 'characteristic'
                % Diagonalize C
                [T_C, Lambda_C] = eig(C);
                lambda_C = diag(Lambda_C);
                % Identify in- and outgoing characteristic variables
                Iplus = lambda_C > 0;
                Iminus = lambda_C < 0;

                switch boundary
                case 'l'
                    Iin_C = Iplus;
                case 'r'
                    Iin_C = Iminus;
                end
                T_C_inv = inv(T_C);
                L = T_C_inv(Iin_C,:);
            otherwise
                error('Boundary condition not implemented.');
            end

            % Penalty parameters
            sigma = [0; 0];
            sigma(Iin) = lambda(Iin);

            % Sparsify
            A = sparse(A);
            T = sparse(T);
            sigma = sparse(sigma);
            L = sparse(L);
            Tin = sparse(Tin);

            switch boundary
            case 'l'
                tau = -1*obj.e_l * inv(A) * inv(T)' * sigma * inv(L*Tin);
                closure = obj.Hi*tau*L*obj.e_l';

            case 'r'
                tau = 1*obj.e_r * inv(A) * inv(T)' * sigma * inv(L*Tin);
                closure = obj.Hi*tau*L*obj.e_r';

            end

            penalty = -obj.Hi*tau;

         end

         % Uses the hat variable method for the interface coupling,
         % see hypsyst_varcoeff.pdf in the hypsyst repository.
         % Notation: u for left side of interface, v for right side.
         function [closure, penalty] = interface(obj,boundary,neighbour_scheme,neighbour_boundary)

            % Get boundary matrices
            switch boundary
                case 'l'
                    A_v = obj.A_l;
                    B_v = obj.B_l;
                    Hi_v = obj.Hi;
                    e_v = obj.e_l;

                    A_u = neighbour_scheme.A_r;
                    B_u = neighbour_scheme.B_r;
                    Hi_u = neighbour_scheme.Hi;
                    e_u = neighbour_scheme.e_r;
                case 'r'
                    A_u = obj.A_r;
                    B_u = obj.B_r;
                    Hi_u = obj.Hi;
                    e_u = obj.e_r;

                    A_v = neighbour_scheme.A_l;
                    B_v = neighbour_scheme.B_l;
                    Hi_v = neighbour_scheme.Hi;
                    e_v = neighbour_scheme.e_l;
            end
            C_u = inv(A_u)*B_u;
            C_v = inv(A_v)*B_v;

            % Diagonalize C
            [T_u, ~] = eig(C_u);
            Lambda_u = inv(T_u)*C_u*T_u;
            lambda_u = diag(Lambda_u);
            S_u = inv(T_u);

            [T_v, ~] = eig(C_v);
            Lambda_v = inv(T_v)*C_v*T_v;
            lambda_v = diag(Lambda_v);
            S_v = inv(T_v);

            % Identify in- and outgoing characteristic variables
            Im_u = lambda_u < 0;
            Ip_u = lambda_u > 0;
            Im_v = lambda_v < 0;
            Ip_v = lambda_v > 0;

            Tp_u = T_u(:,Ip_u);
            Tm_u = T_u(:,Im_u);
            Sp_u = S_u(Ip_u,:);
            Sm_u = S_u(Im_u,:);

            Tp_v = T_v(:,Ip_v);
            Tm_v = T_v(:,Im_v);
            Sp_v = S_v(Ip_v,:);
            Sm_v = S_v(Im_v,:);

            % Create S_tilde and T_tilde
            Stilde = [Sp_u; Sm_v];
            Ttilde = inv(Stilde);
            Ttilde_p = Ttilde(:,1);
            Ttilde_m = Ttilde(:,2);

            % Solve for penalty parameters theta_1,2
            LHS = Ttilde_m*Sm_v*Tm_u;
            RHS = B_u*Tm_u;
            th_u = RHS./LHS;
            TH_u = diag(th_u);

            LHS = Ttilde_p*Sp_u*Tp_v;
            RHS = -B_v*Tp_v;
            th_v = RHS./LHS;
            TH_v = diag(th_v);

            % Construct penalty matrices
            Z_u = TH_u*Ttilde_m*Sm_v;
            Z_u = sparse(Z_u);

            Z_v = TH_v*Ttilde_p*Sp_u;
            Z_v = sparse(Z_v);

            closure_u = Hi_u*e_u*(A_u\Z_u)*e_u';
            penalty_u = -Hi_u*e_u*(A_u\Z_u)*e_v';
            closure_v = Hi_v*e_v*(A_v\Z_v)*e_v';
            penalty_v = -Hi_v*e_v*(A_v\Z_v)*e_u';

             switch boundary
                 case 'l'
                     closure = closure_v;
                     penalty = penalty_v;
                 case 'r'
                     closure = closure_u;
                     penalty = penalty_u;
             end

         end

        function N = size(obj)
            N = 2*obj.m+1;
        end

    end

    methods(Static)
        % Calculates the matrices needed for the inteface coupling between boundary bound_u of scheme schm_u
        % and bound_v of scheme schm_v.
        %   [uu, uv, vv, vu] = inteface_coupling(A,'r',B,'l')
        function [uu, uv, vv, vu] = interface_coupling(schm_u,bound_u,schm_v,bound_v)
            [uu,uv] = schm_u.interface(bound_u,schm_v,bound_v);
            [vv,vu] = schm_v.interface(bound_v,schm_u,bound_u);
        end
    end
end