classdef Matrix < sympy.Base

    properties(Access=private)
        m
        n
    end

    methods 
        function o = Matrix(m, n, x)
            o@sympy.Base(x);
            o.m = m; o.n = n;
        end

        function o = subsref(m, i, j)
            obj = getObj(m);
            o = sympy.Expr(obj{i,j});
        end

        function n = numel(o)
            n = o.m * o.n;
        end

        function varargout = size(o, dim)
            if nargin == 1
                if nargout <= 1
                    varargout{1} = [o.m, o.n];
                else
                    varargout{1} = o.m;
                    varargout{2} = o.n;
                end
            else
                if dim == 1
                    varargout{1} = o.m;
                elseif dim == 2
                    varargout{1} = o.n;
                else
                    varargout{1} = 1;
                end
            end
        end

        function r = horzcat(varargin)
            sp = sympy__sp;
            nArgs = nargin;
            if nArgs == 0
                error("输入参数不能为空");
            end

            blocks = cell(1, nArgs);
            rows = -1;
            totalCols = 0;
            for i = 1:nArgs
                cur = varargin{i};
                if isa(cur, 'sympy.Matrix')
                    [m, n] = size(cur);
                    blocks{i} = struct('obj', getObj(cur), 'm', m, 'n', n);
                elseif isa(cur, 'sympy.Expr')
                    pobj = getObj(cur);
                    blocks{i} = struct('obj', sp.Matrix(1, 1, pylist(pobj)), 'm', 1, 'n', 1);
                else
                    error("horzcat 仅支持 sympy.Matrix 或 sympy.Expr");
                end

                if rows < 0
                    rows = blocks{i}.m;
                elseif rows ~= blocks{i}.m
                    error("水平拼接时行数必须一致");
                end
                totalCols = totalCols + blocks{i}.n;
            end

            data = pylist;
            for rIdx = 1:rows
                for i = 1:nArgs
                    blk = blocks{i};
                    for cIdx = 1:blk.n
                        data.append(blk.obj{rIdx - 1, cIdx - 1});
                    end
                end
            end
            r = sympy.Matrix(rows, totalCols, sp.Matrix(rows, totalCols, data));
        end

        function r = vertcat(varargin)
            sp = sympy__sp;
            nArgs = nargin;
            if nArgs == 0
                error("输入参数不能为空");
            end

            blocks = cell(1, nArgs);
            cols = -1;
            totalRows = 0;
            for i = 1:nArgs
                cur = varargin{i};
                if isa(cur, 'sympy.Matrix')
                    [m, n] = size(cur);
                    blocks{i} = struct('obj', getObj(cur), 'm', m, 'n', n);
                elseif isa(cur, 'sympy.Expr')
                    pobj = getObj(cur);
                    blocks{i} = struct('obj', sp.Matrix(1, 1, pylist(pobj)), 'm', 1, 'n', 1);
                else
                    error("vertcat 仅支持 sympy.Matrix 或 sympy.Expr");
                end

                if cols < 0
                    cols = blocks{i}.n;
                elseif cols ~= blocks{i}.n
                    error("垂直拼接时列数必须一致");
                end
                totalRows = totalRows + blocks{i}.m;
            end

            data = pylist;
            for i = 1:nArgs
                blk = blocks{i};
                for rIdx = 1:blk.m
                    for cIdx = 1:cols
                        data.append(blk.obj{rIdx - 1, cIdx - 1});
                    end
                end
            end
            r = sympy.Matrix(totalRows, cols, sp.Matrix(totalRows, cols, data));
        end

        function [V, D] = eig(A)
            % 使矩阵对角化。
            % D = symeig(A); 返回与A相似的对角阵。
            % [V, D] = symeig(A); A = V*D*V^-1，其中D是对角阵。
            A = getObj(A);
            r = A.diagonalize();
            if nargout == 1
                V = sympy.Matrix(r{1});
            elseif nargout == 2
                V = sympy.Matrix(r{0});
                D = sympy.Matrix(r{1});
            end
        end
    end

    methods (Static)
        function R = cumprod(A, dim)
            narginchk(1, 2);

            if ~isa(A, 'sympy.Matrix')
                R = builtin('cumprod', A);
                return;
            end

            [m, n] = size(A);
            if nargin < 2
                if m == 1 && n > 1
                    dim = 2;
                else
                    dim = 1;
                end
            end

            if ~isscalar(dim) || ~(dim == 1 || dim == 2)
                error("dim 只能取 1 或 2");
            end

            sp = sympy__sp;
            obj = getObj(A);
            data = pylist;

            if dim == 1
                colAcc = cell(1, n);
                for i = 1:m
                    for j = 1:n
                        cur = obj{i - 1, j - 1};
                        if i == 1
                            colAcc{j} = cur;
                        else
                            colAcc{j} = colAcc{j} * cur;
                        end
                        data.append(colAcc{j});
                    end
                end
            else
                for i = 1:m
                    rowAcc = [];
                    for j = 1:n
                        cur = obj{i - 1, j - 1};
                        if j == 1
                            rowAcc = cur;
                        else
                            rowAcc = rowAcc * cur;
                        end
                        data.append(rowAcc);
                    end
                end
            end

            R = sympy.Matrix(m, n, sp.Matrix(m, n, data));
        end

        function R = cumsum(A, dim)
            narginchk(1, 2);

            if ~isa(A, 'sympy.Matrix')
                R = builtin('cumsum', A);
                return;
            end

            [m, n] = size(A);
            if nargin < 2
                if m == 1 && n > 1
                    dim = 2;
                else
                    dim = 1;
                end
            end

            if ~isscalar(dim) || ~(dim == 1 || dim == 2)
                error("dim 只能取 1 或 2");
            end

            sp = sympy__sp;
            obj = getObj(A);
            data = pylist;

            if dim == 1
                colAcc = cell(1, n);
                for i = 1:m
                    for j = 1:n
                        cur = obj{i - 1, j - 1};
                        if i == 1
                            colAcc{j} = cur;
                        else
                            colAcc{j} = colAcc{j} + cur;
                        end
                        data.append(colAcc{j});
                    end
                end
            else
                for i = 1:m
                    rowAcc = [];
                    for j = 1:n
                        cur = obj{i - 1, j - 1};
                        if j == 1
                            rowAcc = cur;
                        else
                            rowAcc = rowAcc + cur;
                        end
                        data.append(rowAcc);
                    end
                end
            end

            R = sympy.Matrix(m, n, sp.Matrix(m, n, data));
        end
    end

end