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

        grid
        xm, ym

        order % Order accuracy for the approximation

        D % non-stabalized scheme operator
        M % Derivative norm
        H % Discrete norm
        Hi
        H_u, H_v % Norms in the x and y directions
        Hu,Hv % Kroneckerd norms. 1'*Hx*v corresponds to integration in the x dir.
        Hi_u, Hi_v
        Hiu, Hiv
        D1_v, D1_u
        D2_v, D2_u
        Du, Dv
        x,y
        b1, b2
        b1_u,b2_v
        DU, DV, DUU, DUV, DVU, DVV 
        
        e_w, e_e, e_s, e_n
        du_w, dv_w
        du_e, dv_e
        du_s, dv_s
        du_n, dv_n
        g_1, g_2
        c
        ind
        t_up
        
        a11, a12, a22
        m_tot, m_u, m_v
        p
        Ji, J
        Hamiltonian
    end

    methods
        function obj = Schrodinger2dCurve(g ,order, opSet,p)
            default_arg('opSet',@sbp.D2Variable);
            default_arg('c', 1);

            obj.p=p;
            obj.c=1;
            
            m = g.size();
            obj.m_u = m(1);
            obj.m_v = m(2);
            obj.m_tot = g.N();
            obj.grid=g;
            
            h = g.scaling();
            h_u = h(1);
            h_v = h(2);

            % Operators
            ops_u = opSet(obj.m_u, {0, 1}, order);
            ops_v = opSet(obj.m_v, {0, 1}, order);

            I_u = speye(obj.m_u);
            I_v = speye(obj.m_v);

            obj.D1_u = ops_u.D1;
            obj.D2_u = ops_u.D2;

            H_u =  ops_u.H;
            Hi_u = ops_u.HI;
            e_l_u = ops_u.e_l;
            e_r_u = ops_u.e_r;
            d1_l_u = ops_u.d1_l;
            d1_r_u = ops_u.d1_r;

            obj.D1_v = ops_v.D1;
            obj.D2_v = ops_v.D2;
            H_v =  ops_v.H;
            Hi_v = ops_v.HI;
            e_l_v = ops_v.e_l;
            e_r_v = ops_v.e_r;
            d1_l_v = ops_v.d1_l;
            d1_r_v = ops_v.d1_r;

            obj.Du = kr(obj.D1_u,I_v);
            obj.Dv = kr(I_u,obj.D1_v);
            
            obj.H = kr(H_u,H_v);
            obj.Hi = kr(Hi_u,Hi_v);
            obj.Hu  = kr(H_u,I_v);
            obj.Hv  = kr(I_u,H_v);
            obj.Hiu = kr(Hi_u,I_v);
            obj.Hiv = kr(I_u,Hi_v);
            
            obj.e_w  = kr(e_l_u,I_v);
            obj.e_e  = kr(e_r_u,I_v);
            obj.e_s  = kr(I_u,e_l_v);
            obj.e_n  = kr(I_u,e_r_v);
            obj.du_w = kr(d1_l_u,I_v);
            obj.dv_w = (obj.e_w'*obj.Dv)';
            obj.du_e = kr(d1_r_u,I_v);   
            obj.dv_e = (obj.e_e'*obj.Dv)';
            obj.du_s = (obj.e_s'*obj.Du)';
            obj.dv_s = kr(I_u,d1_l_v);
            obj.du_n = (obj.e_n'*obj.Du)';
            obj.dv_n = kr(I_u,d1_r_v);
            
            obj.DUU = sparse(obj.m_tot);
            obj.DVV = sparse(obj.m_tot);
            obj.ind = grid.funcToMatrix(obj.grid, 1:obj.m_tot);
            
            obj.m = m;
            obj.h = [h_u h_v];
            obj.order = order;
            obj.D = @(t)obj.d_fun(t);
            obj.Hamiltonian = @(t)obj.h_fun(t);
            obj.variable_update(0);
        end
        
        function [D,n] = d_fun(obj,t)
            %         obj.update_vairables(t); In driscretization?
            %  D = obj.Ji*(-1/2*(obj.b1*obj.Du-obj.b1_u+obj.Du*obj.b1) -
            %  1/2*(obj.b2*obj.Dv - obj.b2_v +obj.Dv*obj.b2) +
            %  1i*obj.c^2*(obj.DUU + obj.DUV + obj.DVU + obj.DVV)); (ols
            %  not skew sym disc
            
            D = sqrt(obj.Ji)*(-1/2*(obj.b1*obj.Du + obj.Du*obj.b1) - 1/2*(obj.b2*obj.Dv + obj.Dv*obj.b2) + 1i*obj.c^2*(obj.DUU + obj.DUV + obj.DVU + obj.DVV))*sqrt(obj.Ji); 
        end
        
         function [Hamiltonian,n] = h_fun(obj,t)
              Hamiltonian =  obj.c^2*(obj.DUU + obj.DUV + obj.DVU + obj.DVV)*sqrt(obj.Ji);
         end    
        
        function [D ]= variable_update(obj,t)
            % Metric derivatives
            if(obj.t_up == t)
                return
            else
                ti = parametrization.Ti.points(obj.p{1}(t),obj.p{2}(t),obj.p{3}(t),obj.p{4}(t));
                ti_tau = parametrization.Ti.points(obj.p{5}(t),obj.p{6}(t),obj.p{7}(t),obj.p{8}(t));
                lcoords=points(obj.grid);
                [obj.xm,obj.ym]= ti.map(lcoords(1:obj.m_v:end,1),lcoords(1:obj.m_v,2));
                [x_tau,y_tau]= ti_tau.map(lcoords(1:obj.m_v:end,1),lcoords(1:obj.m_v,2));
                x = reshape(obj.xm,obj.m_tot,1);
                y = reshape(obj.ym,obj.m_tot,1);
                obj.x = x;
                obj.y = y;
                
                x_tau = reshape(x_tau,obj.m_tot,1);
                y_tau = reshape(y_tau,obj.m_tot,1);
                
                x_u = obj.Du*x;
                x_v = obj.Dv*x;
                y_u = obj.Du*y;
                y_v = obj.Dv*y;
                
                J = (x_u.*y_v - x_v.*y_u);
                obj.J = spdiags(J, 0, obj.m_tot, obj.m_tot);
                
                Ji = spdiags(1./J, 0, obj.m_tot, obj.m_tot);
                obj.Ji = Ji;
                
                a11 =  Ji* (x_v.^2  + y_v.^2);
                a12 = -Ji* (x_u.*x_v + y_u.*y_v);
                a22 =  Ji* (x_u.^2  + y_u.^2);
                
                obj.a11 = a11;
                obj.a12 = a12;
                obj.a22 = a22;
                
                % Assemble full operators
                L_12 = spdiags(a12, 0, obj.m_tot, obj.m_tot);
                obj.DUV = obj.Du*L_12*obj.Dv;
                obj.DVU = obj.Dv*L_12*obj.Du;
                
                
                for i = 1:obj.m_v
                    D = obj.D2_u(a11(obj.ind(:,i)));
                    p = obj.ind(:,i);
                    obj.DUU(p,p) = D;
                end
                
                for i = 1:obj.m_u
                    D = obj.D2_v(a22(obj.ind(i,:)));
                    p = obj.ind(i,:);
                    obj.DVV(p,p) = D;
                end
                
                obj.g_1 = x_v.*y_tau-x_tau.*y_v;
                obj.g_2 = x_tau.*y_u - y_tau.*x_u;
                
                obj.b1 = spdiags(obj.g_1, 0, obj.m_tot, obj.m_tot);
                obj.b2 = spdiags(obj.g_2, 0, obj.m_tot, obj.m_tot);
                
                obj.b1_u = spdiags(obj.Du*obj.g_1, 0, obj.m_tot, obj.m_tot);
                obj.b2_v = spdiags(obj.Dv*obj.g_2, 0, obj.m_tot, obj.m_tot);
                obj.t_up=t;
            end
        end

        % Closure functions return the opertors applied to the own doamin to close the boundary
            
        % Penalty functions return the opertors 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.
        %       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,~)
                    [e, d_n, d_t, coeff_t, coeff_n s, halfnorm_inv_n, halfnorm_inv_t, halfnorm_t,g] = obj.get_boundary_ops(boundary);
                 
                    a_t =  @(t) spdiag(coeff_t(t));
                    a_n = @(t) spdiag(coeff_n(t));
                 
                    
                    F = @(t)(s * a_n(t)*d_n' + s * a_t(t)  *d_t')';
                    tau1  = 1;       
                    a  = @(t)spdiag(g(t));
                    tau2 = @(t) (1*s*a(t))/2;

                    penalty_parameter_1 = @(t) 1*1i*halfnorm_inv_n*halfnorm_inv_t*F(t)*e'*halfnorm_t*e;
                    penalty_parameter_2 = @(t) halfnorm_inv_n*e*tau2(t);

                    closure = @(t)  sqrt(obj.Ji)*(obj.c^2 * penalty_parameter_1(t)*e' +  penalty_parameter_2(t)*e')*sqrt(obj.Ji);
                    penalty = @(t) -sqrt(obj.Ji)*(obj.c^2 * penalty_parameter_1(t)*e' +  penalty_parameter_2(t)*e')*sqrt(obj.Ji);
                
        end
        
        function [closure, penalty] = interface(obj,boundary,neighbour_scheme,neighbour_boundary)
            % u denotes the solution in the own domain
            % v denotes the solution in the neighbour domain
            [e_u, d_n_u, d_t_u,  coeff_t_u, coeff_n_u,s_u, halfnorm_inv_u_n, halfnorm_inv_u_t, halfnorm_u_t,gamm_u,Ji_u] = obj.get_boundary_ops(boundary);
            [e_v, d_n_v, d_t_v,  coeff_t_v, coeff_n_v s_v, halfnorm_inv_v_n, halfnorm_inv_v_t, halfnorm_v_t,gamm_v,Ji_v] = neighbour_scheme.get_boundary_ops(neighbour_boundary);

            a_n_u = @(t) spdiag(coeff_n_u(t));
            a_t_u = @(t) spdiag(coeff_t_u(t));
            a_n_v = @(t) spdiag(coeff_n_v(t));
            a_t_v = @(t) spdiag(coeff_t_v(t));
           

            F_u = @(t)(a_n_u(t)*d_n_u' + a_t_u(t)*d_t_u')';
            F_v = @(t)(a_n_v(t)*d_n_v' + a_t_v(t)*d_t_v')';
            
            a = @(t)spdiag(gamm_u(t));
       
            tau =  s_u*1*1i/2;            
            sig =  -s_u*1*1i/2;
            gamm = @(t) (-s_u*a(t))/2;
            
            penalty_parameter_1 = @(t) halfnorm_inv_u_n*(tau*halfnorm_inv_u_t*F_u(t)*e_u'*halfnorm_u_t*e_u);
            penalty_parameter_2 = @(t) halfnorm_inv_u_n * e_u  * (sig );
            penalty_parameter_3 = @(t) halfnorm_inv_u_n * e_u  * (gamm(t) );


            closure =@(t) sqrt(Ji_u)*obj.c^2 * ( penalty_parameter_1(t)*e_u' + penalty_parameter_2(t)*F_u(t)' + penalty_parameter_3(t)*e_u')*sqrt(Ji_u);
            penalty =@(t) sqrt(Ji_u)*obj.c^2 * ( -penalty_parameter_1(t)*e_v' - penalty_parameter_2(t)*F_v(t)' - penalty_parameter_3(t)*e_v')*sqrt(Ji_v);
        end


        function [e, d_n, d_t, coeff_t,coeff_n, s,  halfnorm_inv_n, halfnorm_inv_t, halfnorm_t,g,Ji] = get_boundary_ops(obj, boundary)

            % gridMatrix = zeros(obj.m(2),obj.m(1));
            % gridMatrix(:) = 1:numel(gridMatrix);

            ind = grid.funcToMatrix(obj.grid, 1:prod(obj.m));

            switch boundary
                case 'w'
                    e = obj.e_w;
                    d_n = obj.du_w;
                    d_t = obj.dv_w;
                    s = -1;

                    I = ind(1,:);
                    coeff_t = @(t)obj.a12(I);
                    coeff_n =@(t) obj.a11(I);
                    g = @(t)obj.g_1(I);
                case 'e'
                    e = obj.e_e;
                    d_n = obj.du_e;
                    d_t = obj.dv_e;
                    s = 1;

                    I = ind(end,:);
                    coeff_t = @(t)obj.a12(I);
                    coeff_n = @(t)obj.a11(I);
                    g = @(t)obj.g_1(I);
                case 's'
                    e = obj.e_s;
                    d_n = obj.dv_s;
                    d_t = obj.du_s;
                    s = -1;

                    I = ind(:,1)';
                    coeff_t = @(t)obj.a12(I);
                    coeff_n = @(t)obj.a22(I);
                    g = @(t)obj.g_2(I);
                case 'n'
                    e = obj.e_n;
                    d_n = obj.dv_n;
                    d_t = obj.du_n;
                    s = 1;

                    I = ind(:,end)';
                    coeff_t = @(t)obj.a12(I);
                    coeff_n = @(t)obj.a22(I);
                    g = @(t)obj.g_2(I);
                otherwise
                    error('No such boundary: boundary = %s',boundary);
            end

            switch boundary
                case {'w','e'}
                    halfnorm_inv_n = obj.Hiu;
                    halfnorm_inv_t = obj.Hiv;
                    halfnorm_t = obj.Hv;
                   
                case {'s','n'}
                    halfnorm_inv_n = obj.Hiv;
                    halfnorm_inv_t = obj.Hiu;
                    halfnorm_t = obj.Hu;
            end
             Ji = obj.Ji;
        end

        function N = size(obj)
            N = prod(obj.m);
        end
    end
    methods(Static)
        % Calculates the matrcis need for the inteface coupling between boundary bound_u of scheme schm_u
        % and bound_v of scheme schm_v.
        %   [uu, uv, vv, vu] = inteface_couplong(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