classdef Eqs < sympy.Base

    methods
        function o = Eqs(m, n, x)
            % Eqs 方程矩阵：obj 为 sympy.Matrix，元素为 sympy Expr（等价于 f==0 左侧）
            o@sympy.Base(x);
        end

        function n = numel(o)
            s = getObj(o).shape;
            n = double(s{0}) * double(s{1});
        end

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

        function S = solve(obj, var)
            py = pyroot;
            sp = sympy__sp;
            eqList = pylist;
            o = getObj(obj);
            [mm, nn] = size(obj);
            for i = 1:mm
                for j = 1:nn
                    eqList.append(o{i - 1, j - 1});
                end
            end
            if py.len(eqList) == 0
                error('Eqs.solve: 空方程组');
            end
            if py.len(eqList) == 1
                symeq_sp = eqList{0};
            else
                symeq_sp = eqList;
            end
            if iscell(var)
                var_sp = pylist;
                for i = 1:length(var)
                    if ~isa(var{i}, 'sympy.Expr') && ~isa(var{i}, 'sympy.Function')
                        error('求解变量必须是 sympy.Expr 或 sympy.Function');
                    end
                    var_sp.append(var{i}.getObj());
                end
            elseif isa(var, 'sympy.Matrix')
                var_sp = pylist;
                vo = getObj(var);
                [vm, vn] = size(var);
                for i = 1:vm
                    for j = 1:vn
                        var_sp.append(vo{i - 1, j - 1});
                    end
                end
            else
                var_sp = var.getObj();
            end
            S_py = sp.solve(symeq_sp, var_sp);
            S = sympyEqs_convSolveResult(py, sp, S_py);
        end

        function S = vpasolve(obj, vars)
            py = pyroot;
            sp = sympy__sp;
            o = getObj(obj);
            [mm, nn] = size(obj);
            exprs2 = pylist;
            for i = 1:mm
                for j = 1:nn
                    exprs2.append(o{i - 1, j - 1});
                end
            end
            varList = sympyEqs_varsToPyList(vars);
            nv = py.len(varList);
            if nv == 0
                error('Eqs.vpasolve: 需要至少一个求解变量');
            end
            guess = pylist;
            for k = 1:nv
                guess.append(py.float(1.0));
            end
            if nv == 1
                sol = sp.nsolve(exprs2{0}, varList{0}, py.float(0.5));
                S = sympy.Expr(sol);
                return;
            end
            sol = sp.nsolve(py.tuple(exprs2), py.tuple(varList), py.tuple(guess));
            names = sympyEqs_varNames(varList, nv);
            S = struct;
            solM = sp.Matrix(sol);
            sm = solM.shape;
            sm0 = double(sm{0});
            sm1 = double(sm{1});
            for k = 1:nv
                if sm1 == 1
                    entry = solM{k - 1, 0};
                else
                    entry = solM{0, k - 1};
                end
                S.(names{k}) = sympy.Expr(entry);
            end
        end

        function r = horzcat(varargin)
            sp = sympy__sp;
            if nargin == 0
                error('horzcat: 输入不能为空');
            end
            rows = -1;
            totalCols = 0;
            blocks = cell(1, nargin);
            for i = 1:nargin
                cur = varargin{i};
                if isa(cur, 'sympy.Eqs')
                    [m, n] = size(cur);
                    blocks{i} = struct('obj', getObj(cur), 'm', m, 'n', n);
                elseif isa(cur, 'sympy.Eq')
                    pobj = getObj(cur);
                    pex = (pobj.lhs - pobj.rhs).simplify();
                    blocks{i} = struct('obj', sp.Matrix(1, 1, pylist(pex)), 'm', 1, 'n', 1);
                else
                    error('Eqs.horzcat 仅支持 sympy.Eqs 或 sympy.Eq');
                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:nargin
                    blk = blocks{i};
                    for cIdx = 1:blk.n
                        data.append(blk.obj{rIdx - 1, cIdx - 1});
                    end
                end
            end
            r = sympy.Eqs(rows, totalCols, sp.Matrix(rows, totalCols, data));
        end

        function r = vertcat(varargin)
            sp = sympy__sp;
            if nargin == 0
                error('vertcat: 输入不能为空');
            end
            cols = -1;
            totalRows = 0;
            blocks = cell(1, nargin);
            for i = 1:nargin
                cur = varargin{i};
                if isa(cur, 'sympy.Eqs')
                    [m, n] = size(cur);
                    blocks{i} = struct('obj', getObj(cur), 'm', m, 'n', n);
                elseif isa(cur, 'sympy.Eq')
                    pobj = getObj(cur);
                    pex = (pobj.lhs - pobj.rhs).simplify();
                    blocks{i} = struct('obj', sp.Matrix(1, 1, pylist(pex)), 'm', 1, 'n', 1);
                else
                    error('Eqs.vertcat 仅支持 sympy.Eqs 或 sympy.Eq');
                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:nargin
                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.Eqs(totalRows, cols, sp.Matrix(totalRows, cols, data));
        end
    end
end

function L = sympyEqs_varsToPyList(vars)
    if iscell(vars)
        L = pylist;
        for i = 1:numel(vars)
            v = vars{i};
            if isa(v, 'sympy.Expr') || isa(v, 'sympy.Function')
                L.append(v.getObj());
            else
                error('vars 元胞内须为 sympy.Expr/Function');
            end
        end
    elseif isa(vars, 'sympy.Matrix')
        L = pylist;
        vo = vars.getObj();
        [vm, vn] = size(vars);
        for i = 1:vm
            for j = 1:vn
                L.append(vo{i - 1, j - 1});
            end
        end
    elseif isa(vars, 'sympy.Expr') || isa(vars, 'sympy.Function')
        L = pylist(vars.getObj());
    else
        error('vars 类型不支持');
    end
end

function names = sympyEqs_varNames(varList, nv)
    names = cell(1, double(nv));
    for k = 1:double(nv)
        sy = varList{k - 1};
        names{k} = char(sy.name);
    end
end

function S = sympyEqs_convSolveResult(py, ~, S_py)
    % 与 sympy.Eq.solve 相同的结果封装逻辑
    if py.isinstance(S_py, py.dict)
        S = struct;
        dict_keys = py.list(S_py.keys());
        dict_vals = py.list(S_py.values());
        for i = 1:py.len(dict_keys)
            key_sp = py.str(dict_keys{i - 1});
            val_sp = dict_vals{i - 1};
            S.(key_sp) = sympy.Expr(val_sp);
        end
    elseif py.isinstance(S_py, py.list) || py.isinstance(S_py, py.tuple)
        S_len = py.len(S_py);
        if S_len == 1
            S = sympy.Expr(S_py{0}.simplify());
            return;
        end
        S = {};
        for i = 1:py.len(S_py)
            sol_sp = S_py{i - 1}.simplify();
            S{end + 1} = sympy.Expr(sol_sp);
        end
    else
        S = sympy.Expr(S_py);
    end
end
