pocketpy/src/linalg.cpp

#include "pocketpy/linalg.h"

namespace pkpy{

#define BIND_VEC_VEC_OP(D, name, op)                                                    \
        vm->bind##name(type->as<Type>(), [](VM* vm, PyVar _0, PyVar _1){          \
            Vec##D self = _CAST(Vec##D, _0);                                            \
            Vec##D other = CAST(Vec##D, _1);                                            \
            return VAR(self op other);                                                  \
        });

#define BIND_VEC_FLOAT_OP(D, name, op)  \
        vm->bind##name(type->as<Type>(), [](VM* vm, PyVar _0, PyVar _1){          \
            Vec##D self = _CAST(Vec##D, _0);                                            \
            f64 other = CAST(f64, _1);                                                  \
            return VAR(self op other);                                                  \
        });

#define BIND_VEC_FUNCTION_0(T, name)                                        \
        vm->bind_func(type, #name, 1, [](VM* vm, ArgsView args){            \
            T self = _CAST(T, args[0]);                                     \
            return VAR(self.name());                                        \
        });

#define BIND_VEC_FUNCTION_1(T, name)                                        \
        vm->bind_func(type, #name, 2, [](VM* vm, ArgsView args){            \
            T self = _CAST(T, args[0]);                                     \
            T other = CAST(T, args[1]);                                     \
            return VAR(self.name(other));                                   \
        });

#define BIND_VEC_MUL_OP(D)                                                                  \
        vm->bind__mul__(type->as<Type>(), [](VM* vm, PyVar _0, PyVar _1){             \
            Vec##D self = _CAST(Vec##D, _0);                                                \
            if(vm->is_user_type<Vec##D>(_1)){                                               \
                Vec##D other = _CAST(Vec##D, _1);                                           \
                return VAR(self * other);                                                   \
            }                                                                               \
            f64 other = CAST(f64, _1);                                                      \
            return VAR(self * other);                                                       \
        });                                                                                 \
        vm->bind_func(type, "__rmul__", 2, [](VM* vm, ArgsView args){                       \
            Vec##D self = _CAST(Vec##D, args[0]);                                           \
            f64 other = CAST(f64, args[1]);                                                 \
            return VAR(self * other);                                                       \
        });                                                                                 \
        vm->bind__truediv__(type->as<Type>(), [](VM* vm, PyVar _0, PyVar _1){ \
            Vec##D self = _CAST(Vec##D, _0);                                          \
            f64 other = CAST(f64, _1);                                                      \
            return VAR(self / other);                                                       \
        });

#define BIND_VEC_GETITEM(D) \
        vm->bind__getitem__(type->as<Type>(), [](VM* vm, PyVar obj, PyVar index){ \
            Vec##D self = _CAST(Vec##D, obj); \
            i64 i = CAST(i64, index); \
            if(i < 0 || i >= D) vm->IndexError("index out of range"); \
            return VAR(self[i]); \
        });

#define BIND_SSO_VEC_COMMON(D)  \
        vm->bind__eq__(type->as<Type>(), [](VM* vm, PyVar _0, PyVar _1){  \
            Vec##D self = _CAST(Vec##D, _0);                                    \
            if(!vm->is_user_type<Vec##D>(_1)) return vm->NotImplemented;        \
            Vec##D other = _CAST(Vec##D, _1);                                   \
            return VAR(self == other);                                          \
        });                                                                     \
        vm->bind_func(type, "__getnewargs__", 1, [](VM* vm, ArgsView args){     \
            Vec##D self = _CAST(Vec##D, args[0]);                               \
            Tuple t(D);                                                         \
            for(int i=0; i<D; i++) t[i] = VAR(self[i]);                         \
            return VAR(std::move(t));                                           \
        });

// https://github.com/Unity-Technologies/UnityCsReference/blob/master/Runtime/Export/Math/Vector2.cs#L289
static Vec2 SmoothDamp(Vec2 current, Vec2 target, Vec2& currentVelocity, float smoothTime, float maxSpeed, float deltaTime)
{
    // Based on Game Programming Gems 4 Chapter 1.10
    smoothTime = std::max(0.0001F, smoothTime);
    float omega = 2.0F / smoothTime;

    float x = omega * deltaTime;
    float exp = 1.0F / (1.0F + x + 0.48F * x * x + 0.235F * x * x * x);

    float change_x = current.x - target.x;
    float change_y = current.y - target.y;
    Vec2 originalTo = target;

    // Clamp maximum speed
    float maxChange = maxSpeed * smoothTime;

    float maxChangeSq = maxChange * maxChange;
    float sqDist = change_x * change_x + change_y * change_y;
    if (sqDist > maxChangeSq)
    {
        float mag = std::sqrt(sqDist);
        change_x = change_x / mag * maxChange;
        change_y = change_y / mag * maxChange;
    }

    target.x = current.x - change_x;
    target.y = current.y - change_y;

    float temp_x = (currentVelocity.x + omega * change_x) * deltaTime;
    float temp_y = (currentVelocity.y + omega * change_y) * deltaTime;

    currentVelocity.x = (currentVelocity.x - omega * temp_x) * exp;
    currentVelocity.y = (currentVelocity.y - omega * temp_y) * exp;

    float output_x = target.x + (change_x + temp_x) * exp;
    float output_y = target.y + (change_y + temp_y) * exp;

    // Prevent overshooting
    float origMinusCurrent_x = originalTo.x - current.x;
    float origMinusCurrent_y = originalTo.y - current.y;
    float outMinusOrig_x = output_x - originalTo.x;
    float outMinusOrig_y = output_y - originalTo.y;

    if (origMinusCurrent_x * outMinusOrig_x + origMinusCurrent_y * outMinusOrig_y > 0)
    {
        output_x = originalTo.x;
        output_y = originalTo.y;

        currentVelocity.x = (output_x - originalTo.x) / deltaTime;
        currentVelocity.y = (output_y - originalTo.y) / deltaTime;
    }
    return Vec2(output_x, output_y);
}

    void Vec2::_register(VM* vm, PyObject* mod, PyObject* type){
        type->attr().set("ZERO", vm->new_user_object<Vec2>(0, 0));
        type->attr().set("ONE", vm->new_user_object<Vec2>(1, 1));

        vm->bind_func(type, __new__, 3, [](VM* vm, ArgsView args){
            float x = CAST_F(args[1]);
            float y = CAST_F(args[2]);
            return vm->new_object<Vec2>(PK_OBJ_GET(Type, args[0]), x, y);
        });

        // @staticmethod
        vm->bind(type, "smooth_damp(current: vec2, target: vec2, current_velocity_: vec2, smooth_time: float, max_speed: float, delta_time: float) -> vec2", [](VM* vm, ArgsView args){
            Vec2 current = CAST(Vec2, args[0]);
            Vec2 target = CAST(Vec2, args[1]);
            Vec2 current_velocity_ = CAST(Vec2, args[2]);
            float smooth_time = CAST_F(args[3]);
            float max_speed = CAST_F(args[4]);
            float delta_time = CAST_F(args[5]);
            Vec2 ret = SmoothDamp(current, target, current_velocity_, smooth_time, max_speed, delta_time);
            return VAR(Tuple(VAR(ret), VAR(current_velocity_)));
        }, {}, BindType::STATICMETHOD);

        // @staticmethod
        vm->bind(type, "angle(__from: vec2, __to: vec2) -> float", [](VM* vm, ArgsView args){
            Vec2 __from = CAST(Vec2, args[0]);
            Vec2 __to = CAST(Vec2, args[1]);
            float val = atan2f(__to.y, __to.x) - atan2f(__from.y, __from.x);
            const float PI = 3.1415926535897932384f;
            if(val > PI) val -= 2*PI;
            if(val < -PI) val += 2*PI;
            return VAR(val);
        }, {}, BindType::STATICMETHOD);

        vm->bind__repr__(type->as<Type>(), [](VM* vm, PyVar obj) -> Str{
            Vec2 self = _CAST(Vec2, obj);
            SStream ss;
            ss.setprecision(3);
            ss << "vec2(" << self.x << ", " << self.y << ")";
            return ss.str();
        });

        vm->bind_func(type, "rotate", 2, [](VM* vm, ArgsView args){
            Vec2 self = _CAST(Vec2, args[0]);
            float radian = CAST(f64, args[1]);
            return vm->new_user_object<Vec2>(self.rotate(radian));
        });

        PY_READONLY_FIELD(Vec2, "x", x)
        PY_READONLY_FIELD(Vec2, "y", y)

        BIND_VEC_VEC_OP(2, __add__, +)
        BIND_VEC_VEC_OP(2, __sub__, -)
        BIND_VEC_MUL_OP(2)
        BIND_VEC_FLOAT_OP(2, __truediv__, /)
        BIND_VEC_FUNCTION_1(Vec2, dot)
        BIND_VEC_FUNCTION_1(Vec2, cross)
        BIND_VEC_FUNCTION_0(Vec2, length)
        BIND_VEC_FUNCTION_0(Vec2, length_squared)
        BIND_VEC_FUNCTION_0(Vec2, normalize)
        BIND_VEC_GETITEM(2)
        BIND_SSO_VEC_COMMON(2)
    }

    void Vec3::_register(VM* vm, PyObject* mod, PyObject* type){
        type->attr().set("ZERO", vm->new_user_object<Vec3>(0, 0, 0));
        type->attr().set("ONE", vm->new_user_object<Vec3>(1, 1, 1));

        vm->bind_func(type, __new__, 4, [](VM* vm, ArgsView args){
            float x = CAST_F(args[1]);
            float y = CAST_F(args[2]);
            float z = CAST_F(args[3]);
            return vm->new_object<Vec3>(PK_OBJ_GET(Type, args[0]), x, y, z);
        });

        vm->bind__repr__(type->as<Type>(), [](VM* vm, PyVar obj) -> Str{
            Vec3 self = _CAST(Vec3, obj);
            SStream ss;
            ss.setprecision(3);
            ss << "vec3(" << self.x << ", " << self.y << ", " << self.z << ")";
            return ss.str();
        });

        PY_READONLY_FIELD(Vec3, "x", x)
        PY_READONLY_FIELD(Vec3, "y", y)
        PY_READONLY_FIELD(Vec3, "z", z)

        BIND_VEC_VEC_OP(3, __add__, +)
        BIND_VEC_VEC_OP(3, __sub__, -)
        BIND_VEC_MUL_OP(3)
        BIND_VEC_FUNCTION_1(Vec3, dot)
        BIND_VEC_FUNCTION_1(Vec3, cross)
        BIND_VEC_FUNCTION_0(Vec3, length)
        BIND_VEC_FUNCTION_0(Vec3, length_squared)
        BIND_VEC_FUNCTION_0(Vec3, normalize)
        BIND_VEC_GETITEM(3)
        BIND_SSO_VEC_COMMON(3)
    }

    void Vec4::_register(VM* vm, PyObject* mod, PyObject* type){
        PY_STRUCT_LIKE(Vec4)

        type->attr().set("ZERO", vm->new_user_object<Vec4>(0, 0, 0, 0));
        type->attr().set("ONE", vm->new_user_object<Vec4>(1, 1, 1, 1));

        vm->bind_func(type, __new__, 5, [](VM* vm, ArgsView args){
            float x = CAST_F(args[1]);
            float y = CAST_F(args[2]);
            float z = CAST_F(args[3]);
            float w = CAST_F(args[4]);
            return vm->new_object<Vec4>(PK_OBJ_GET(Type, args[0]), x, y, z, w);
        });

        vm->bind__repr__(type->as<Type>(), [](VM* vm, PyVar obj) -> Str{
            Vec4 self = _CAST(Vec4&, obj);
            SStream ss;
            ss.setprecision(3);
            ss << "vec4(" << self.x << ", " << self.y << ", " << self.z << ", " << self.w << ")";
            return ss.str();
        });

        PY_FIELD(Vec4, "x", x)
        PY_FIELD(Vec4, "y", y)
        PY_FIELD(Vec4, "z", z)
        PY_FIELD(Vec4, "w", w)

        BIND_VEC_VEC_OP(4, __add__, +)
        BIND_VEC_VEC_OP(4, __sub__, -)
        BIND_VEC_MUL_OP(4)
        BIND_VEC_FUNCTION_1(Vec4&, dot)
        BIND_VEC_FUNCTION_1(Vec4&, copy_)
        BIND_VEC_FUNCTION_0(Vec4&, length)
        BIND_VEC_FUNCTION_0(Vec4&, length_squared)
        BIND_VEC_FUNCTION_0(Vec4&, normalize)
        BIND_VEC_FUNCTION_0(Vec4&, normalize_)
        BIND_VEC_GETITEM(4)
    }

#undef BIND_VEC_VEC_OP
#undef BIND_VEC_MUL_OP
#undef BIND_VEC_FUNCTION_0
#undef BIND_VEC_FUNCTION_1
#undef BIND_VEC_GETITEM

    void Mat3x3::_register(VM* vm, PyObject* mod, PyObject* type){
        PY_STRUCT_LIKE(Mat3x3)

        vm->bind_func(type, __new__, -1, [](VM* vm, ArgsView args){
            if(args.size() == 1+0) return vm->new_object<Mat3x3>(PK_OBJ_GET(Type, args[0]), Mat3x3::zeros());
            if(args.size() == 1+1){
                const List& list = CAST(List&, args[1]);
                if(list.size() != 9) vm->TypeError("Mat3x3.__new__ takes a list of 9 floats");
                Mat3x3 mat;
                for(int i=0; i<9; i++) mat.v[i] = CAST_F(list[i]);
                return vm->new_object<Mat3x3>(PK_OBJ_GET(Type, args[0]), mat);
            }
            if(args.size() == 1+9){
                Mat3x3 mat;
                for(int i=0; i<9; i++) mat.v[i] = CAST_F(args[1+i]);
                return vm->new_object<Mat3x3>(PK_OBJ_GET(Type, args[0]), mat);
            }
            vm->TypeError(_S("Mat3x3.__new__ takes 0 or 1 or 9 arguments, got ", args.size()-1));
            return vm->None;
        });

        vm->bind_func(type, "copy_", 2, [](VM* vm, ArgsView args){
            Mat3x3& self = _CAST(Mat3x3&, args[0]);
            const Mat3x3& other = CAST(Mat3x3&, args[1]);
            self = other;
            return vm->None;
        });

        vm->bind__repr__(type->as<Type>(), [](VM* vm, PyVar obj) -> Str{
            const Mat3x3& self = _CAST(Mat3x3&, obj);
            SStream ss;
            ss.setprecision(3);
            ss << "mat3x3([" << self._11 << ", " << self._12 << ", " << self._13 << ",\n";
            ss << "        " << self._21 << ", " << self._22 << ", " << self._23 << ",\n";
            ss << "        " << self._31 << ", " << self._32 << ", " << self._33 << "])";
            return ss.str();
        });

        vm->bind__getitem__(type->as<Type>(), [](VM* vm, PyVar obj, PyVar index){
            Mat3x3& self = _CAST(Mat3x3&, obj);
            Tuple& t = CAST(Tuple&, index);
            if(t.size() != 2){
                vm->TypeError("Mat3x3.__getitem__ takes a tuple of 2 integers");
            }
            i64 i = CAST(i64, t[0]);
            i64 j = CAST(i64, t[1]);
            if(i < 0 || i >= 3 || j < 0 || j >= 3){
                vm->IndexError("index out of range");
            }
            return VAR(self.m[i][j]);
        });

        vm->bind__setitem__(type->as<Type>(), [](VM* vm, PyVar obj, PyVar index, PyVar value){
            Mat3x3& self = _CAST(Mat3x3&, obj);
            const Tuple& t = CAST(Tuple&, index);
            if(t.size() != 2){
                vm->TypeError("Mat3x3.__setitem__ takes a tuple of 2 integers");
            }
            i64 i = CAST(i64, t[0]);
            i64 j = CAST(i64, t[1]);
            if(i < 0 || i >= 3 || j < 0 || j >= 3){
                vm->IndexError("index out of range");
            }
            self.m[i][j] = CAST_F(value);
        });

        vm->bind_field(type, "_11", &Mat3x3::_11);
        vm->bind_field(type, "_12", &Mat3x3::_12);
        vm->bind_field(type, "_13", &Mat3x3::_13);
        vm->bind_field(type, "_21", &Mat3x3::_21);
        vm->bind_field(type, "_22", &Mat3x3::_22);
        vm->bind_field(type, "_23", &Mat3x3::_23);
        vm->bind_field(type, "_31", &Mat3x3::_31);
        vm->bind_field(type, "_32", &Mat3x3::_32);
        vm->bind_field(type, "_33", &Mat3x3::_33);

        vm->bind__add__(type->as<Type>(), [](VM* vm, PyVar _0, PyVar _1){
            Mat3x3& self = _CAST(Mat3x3&, _0);
            Mat3x3& other = CAST(Mat3x3&, _1);
            return vm->new_user_object<Mat3x3>(self + other);
        });

        vm->bind__sub__(type->as<Type>(), [](VM* vm, PyVar _0, PyVar _1){
            Mat3x3& self = _CAST(Mat3x3&, _0);
            Mat3x3& other = CAST(Mat3x3&, _1);
            return vm->new_user_object<Mat3x3>(self - other);
        });

        vm->bind__mul__(type->as<Type>(), [](VM* vm, PyVar _0, PyVar _1){
            Mat3x3& self = _CAST(Mat3x3&, _0);
            f64 other = CAST_F(_1);
            return vm->new_user_object<Mat3x3>(self * other);
        });

        vm->bind_func(type, "__rmul__", 2, [](VM* vm, ArgsView args){
            Mat3x3& self = _CAST(Mat3x3&, args[0]);
            f64 other = CAST_F(args[1]);
            return vm->new_user_object<Mat3x3>(self * other);
        });

        vm->bind__truediv__(type->as<Type>(), [](VM* vm, PyVar _0, PyVar _1){
            Mat3x3& self = _CAST(Mat3x3&, _0);
            f64 other = CAST_F(_1);
            return vm->new_user_object<Mat3x3>(self / other);
        });

        vm->bind__matmul__(type->as<Type>(), [](VM* vm, PyVar _0, PyVar _1){
            Mat3x3& self = _CAST(Mat3x3&, _0);
            if(vm->is_user_type<Mat3x3>(_1)){
                const Mat3x3& other = _CAST(Mat3x3&, _1);
                return vm->new_user_object<Mat3x3>(self.matmul(other));
            }
            if(vm->is_user_type<Vec3>(_1)){
                const Vec3 other = _CAST(Vec3, _1);
                return vm->new_user_object<Vec3>(self.matmul(other));
            }
            return vm->NotImplemented;
        });

        vm->bind(type, "matmul(self, other: mat3x3, out: mat3x3 = None)", [](VM* vm, ArgsView args){
            const Mat3x3& self = _CAST(Mat3x3&, args[0]);
            const Mat3x3& other = CAST(Mat3x3&, args[1]);
            if(args[2] == vm->None){
                return vm->new_user_object<Mat3x3>(self.matmul(other));
            }else{
                Mat3x3& out = CAST(Mat3x3&, args[2]);
                out = self.matmul(other);
                return vm->None;
            }
        });

        vm->bind_func(type, "determinant", 1, [](VM* vm, ArgsView args){
            Mat3x3& self = _CAST(Mat3x3&, args[0]);
            return VAR(self.determinant());
        });

        vm->bind_func(type, "transpose", 1, [](VM* vm, ArgsView args){
            Mat3x3& self = _CAST(Mat3x3&, args[0]);
            return vm->new_user_object<Mat3x3>(self.transpose());
        });

        vm->bind__invert__(type->as<Type>(), [](VM* vm, PyVar obj){
            Mat3x3& self = _CAST(Mat3x3&, obj);
            Mat3x3 ret;
            if(!self.inverse(ret)) vm->ValueError("matrix is not invertible");
            return vm->new_user_object<Mat3x3>(ret);
        });

        vm->bind_func(type, "inverse", 1, [](VM* vm, ArgsView args){
            Mat3x3& self = _CAST(Mat3x3&, args[0]);
            Mat3x3 ret;
            if(!self.inverse(ret)) vm->ValueError("matrix is not invertible");
            return vm->new_user_object<Mat3x3>(ret);
        });

        vm->bind_func(type, "inverse_", 1, [](VM* vm, ArgsView args){
            Mat3x3& self = _CAST(Mat3x3&, args[0]);
            Mat3x3 ret;
            if(!self.inverse(ret)) vm->ValueError("matrix is not invertible");
            self = ret;
            return vm->None;
        });

        vm->bind_func(type, "transpose_", 1, [](VM* vm, ArgsView args){
            Mat3x3& self = _CAST(Mat3x3&, args[0]);
            self = self.transpose();
            return vm->None;
        });

        // @staticmethod
        vm->bind_func(type, "zeros", 0, [](VM* vm, ArgsView args){
            return vm->new_user_object<Mat3x3>(Mat3x3::zeros());
        }, {}, BindType::STATICMETHOD);

        // @staticmethod
        vm->bind_func(type, "ones", 0, [](VM* vm, ArgsView args){
            return vm->new_user_object<Mat3x3>(Mat3x3::ones());
        }, {}, BindType::STATICMETHOD);

        // @staticmethod
        vm->bind_func(type, "identity", 0, [](VM* vm, ArgsView args){
            return vm->new_user_object<Mat3x3>(Mat3x3::identity());
        }, {}, BindType::STATICMETHOD);

        /*************** affine transformations ***************/
        // @staticmethod
        vm->bind(type, "trs(t: vec2, r: float, s: vec2)", [](VM* vm, ArgsView args){
            Vec2 t = CAST(Vec2, args[0]);
            f64 r = CAST_F(args[1]);
            Vec2 s = CAST(Vec2, args[2]);
            return vm->new_user_object<Mat3x3>(Mat3x3::trs(t, r, s));
        }, {}, BindType::STATICMETHOD);

        vm->bind(type, "copy_trs_(self, t: vec2, r: float, s: vec2)", [](VM* vm, ArgsView args){
            Mat3x3& self = _CAST(Mat3x3&, args[0]);
            Vec2 t = CAST(Vec2, args[1]);
            f64 r = CAST_F(args[2]);
            Vec2 s = CAST(Vec2, args[3]);
            self = Mat3x3::trs(t, r, s);
            return vm->None;
        });

        vm->bind(type, "copy_t_(self, t: vec2)", [](VM* vm, ArgsView args){
            Mat3x3& self = _CAST(Mat3x3&, args[0]);
            Vec2 t = CAST(Vec2, args[1]);
            self = Mat3x3::trs(t, self._r(), self._s());
            return vm->None;
        });

        vm->bind(type, "copy_r_(self, r: float)", [](VM* vm, ArgsView args){
            Mat3x3& self = _CAST(Mat3x3&, args[0]);
            f64 r = CAST_F(args[1]);
            self = Mat3x3::trs(self._t(), r, self._s());
            return vm->None;
        });

        vm->bind(type, "copy_s_(self, s: vec2)", [](VM* vm, ArgsView args){
            Mat3x3& self = _CAST(Mat3x3&, args[0]);
            Vec2 s = CAST(Vec2, args[1]);
            self = Mat3x3::trs(self._t(), self._r(), s);
            return vm->None;
        });

        vm->bind_func(type, "is_affine", 1, [](VM* vm, ArgsView args){
            Mat3x3& self = _CAST(Mat3x3&, args[0]);
            return VAR(self.is_affine());
        });

        vm->bind_func(type, "_t", 1, [](VM* vm, ArgsView args){
            Mat3x3& self = _CAST(Mat3x3&, args[0]);
            return vm->new_user_object<Vec2>(self._t());
        });

        vm->bind_func(type, "_r", 1, [](VM* vm, ArgsView args){
            Mat3x3& self = _CAST(Mat3x3&, args[0]);
            return VAR(self._r());
        });

        vm->bind_func(type, "_s", 1, [](VM* vm, ArgsView args){
            Mat3x3& self = _CAST(Mat3x3&, args[0]);
            return vm->new_user_object<Vec2>(self._s());
        });

        vm->bind_func(type, "transform_point", 2, [](VM* vm, ArgsView args){
            const Mat3x3& self = _CAST(Mat3x3&, args[0]);
            Vec2 v = CAST(Vec2, args[1]);
            Vec2 res(self._11 * v.x + self._12 * v.y + self._13, self._21 * v.x + self._22 * v.y + self._23);
            return vm->new_user_object<Vec2>(res);
        });

        vm->bind_func(type, "inverse_transform_point", 2, [](VM* vm, ArgsView args){
            const Mat3x3& self = _CAST(Mat3x3&, args[0]);
            Vec2 v = CAST(Vec2, args[1]);
            Mat3x3 inv;
            if(!self.inverse(inv)) vm->ValueError("matrix is not invertible");
            Vec2 res(inv._11 * v.x + inv._12 * v.y + inv._13, inv._21 * v.x + inv._22 * v.y + inv._23);
            return vm->new_user_object<Vec2>(res);
        });

        vm->bind_func(type, "transform_vector", 2, [](VM* vm, ArgsView args){
            const Mat3x3& self = _CAST(Mat3x3&, args[0]);
            Vec2 v = CAST(Vec2, args[1]);
            Vec2 res(self._11 * v.x + self._12 * v.y, self._21 * v.x + self._22 * v.y);
            return vm->new_user_object<Vec2>(res);
        });

        vm->bind_func(type, "inverse_transform_vector", 2, [](VM* vm, ArgsView args){
            const Mat3x3& self = _CAST(Mat3x3&, args[0]);
            Vec2 v = CAST(Vec2, args[1]);
            Mat3x3 inv;
            if(!self.inverse(inv)) vm->ValueError("matrix is not invertible");
            Vec2 res(inv._11 * v.x + inv._12 * v.y, inv._21 * v.x + inv._22 * v.y);
            return vm->new_user_object<Vec2>(res);
        });
    }


void add_module_linalg(VM* vm){
    PyObject* linalg = vm->new_module("linalg");

    vm->register_user_class<Vec2>(linalg, "vec2", VM::tp_object);
    vm->register_user_class<Vec3>(linalg, "vec3", VM::tp_object);
    vm->register_user_class<Vec4>(linalg, "vec4", VM::tp_object, true);
    vm->register_user_class<Mat3x3>(linalg, "mat3x3", VM::tp_object, true);

    PyVar float_p = vm->_modules["c"]->attr("float_p");
    linalg->attr().set("vec4_p", float_p);
    linalg->attr().set("mat3x3_p", float_p);
}


    /////////////// mat3x3 ///////////////
    Mat3x3::Mat3x3() {}
    Mat3x3::Mat3x3(float _11, float _12, float _13,
           float _21, float _22, float _23,
           float _31, float _32, float _33)
        : _11(_11), _12(_12), _13(_13)
        , _21(_21), _22(_22), _23(_23)
        , _31(_31), _32(_32), _33(_33) {}

    Mat3x3 Mat3x3::zeros(){
        return Mat3x3(0, 0, 0, 0, 0, 0, 0, 0, 0);
    }

    Mat3x3 Mat3x3::ones(){
        return Mat3x3(1, 1, 1, 1, 1, 1, 1, 1, 1);
    }

    Mat3x3 Mat3x3::identity(){
        return Mat3x3(1, 0, 0, 0, 1, 0, 0, 0, 1);
    }

    Mat3x3 Mat3x3::operator+(const Mat3x3& other) const{ 
        Mat3x3 ret;
        for (int i=0; i<9; ++i) ret.v[i] = v[i] + other.v[i];
        return ret;
    }

    Mat3x3 Mat3x3::operator-(const Mat3x3& other) const{ 
        Mat3x3 ret;
        for (int i=0; i<9; ++i) ret.v[i] = v[i] - other.v[i];
        return ret;
    }

    Mat3x3 Mat3x3::operator*(float scalar) const{ 
        Mat3x3 ret;
        for (int i=0; i<9; ++i) ret.v[i] = v[i] * scalar;
        return ret;
    }

    Mat3x3 Mat3x3::operator/(float scalar) const{ 
        Mat3x3 ret;
        for (int i=0; i<9; ++i) ret.v[i] = v[i] / scalar;
        return ret;
    }

    bool Mat3x3::operator==(const Mat3x3& other) const{
        for (int i=0; i<9; ++i){
            if (!isclose(v[i], other.v[i])) return false;
        }
        return true;
    }

    bool Mat3x3::operator!=(const Mat3x3& other) const{
        for (int i=0; i<9; ++i){
            if (!isclose(v[i], other.v[i])) return true;
        }
        return false;
    }

    Mat3x3 Mat3x3::matmul(const Mat3x3& other) const{
        Mat3x3 out;
        out._11 = _11 * other._11 + _12 * other._21 + _13 * other._31;
        out._12 = _11 * other._12 + _12 * other._22 + _13 * other._32;
        out._13 = _11 * other._13 + _12 * other._23 + _13 * other._33;
        out._21 = _21 * other._11 + _22 * other._21 + _23 * other._31;
        out._22 = _21 * other._12 + _22 * other._22 + _23 * other._32;
        out._23 = _21 * other._13 + _22 * other._23 + _23 * other._33;
        out._31 = _31 * other._11 + _32 * other._21 + _33 * other._31;
        out._32 = _31 * other._12 + _32 * other._22 + _33 * other._32;
        out._33 = _31 * other._13 + _32 * other._23 + _33 * other._33;
        return out;
    }

    Vec3 Mat3x3::matmul(const Vec3& other) const{
        Vec3 out;
        out.x = _11 * other.x + _12 * other.y + _13 * other.z;
        out.y = _21 * other.x + _22 * other.y + _23 * other.z;
        out.z = _31 * other.x + _32 * other.y + _33 * other.z;
        return out;
    }

    float Mat3x3::determinant() const{
        return _11 * _22 * _33 + _12 * _23 * _31 + _13 * _21 * _32
             - _11 * _23 * _32 - _12 * _21 * _33 - _13 * _22 * _31;
    }

    Mat3x3 Mat3x3::transpose() const{
        Mat3x3 ret;
        ret._11 = _11;  ret._12 = _21;  ret._13 = _31;
        ret._21 = _12;  ret._22 = _22;  ret._23 = _32;
        ret._31 = _13;  ret._32 = _23;  ret._33 = _33;
        return ret;
    }

    bool Mat3x3::inverse(Mat3x3& out) const{
        float det = determinant();
        if (isclose(det, 0)) return false;
        float inv_det = 1.0f / det;
        out._11 = (_22 * _33 - _23 * _32) * inv_det;
        out._12 = (_13 * _32 - _12 * _33) * inv_det;
        out._13 = (_12 * _23 - _13 * _22) * inv_det;
        out._21 = (_23 * _31 - _21 * _33) * inv_det;
        out._22 = (_11 * _33 - _13 * _31) * inv_det;
        out._23 = (_13 * _21 - _11 * _23) * inv_det;
        out._31 = (_21 * _32 - _22 * _31) * inv_det;
        out._32 = (_12 * _31 - _11 * _32) * inv_det;
        out._33 = (_11 * _22 - _12 * _21) * inv_det;
        return true;
    }

    Mat3x3 Mat3x3::trs(Vec2 t, float radian, Vec2 s){
        float cr = cosf(radian);
        float sr = sinf(radian);
        return Mat3x3(s.x * cr,   -s.y * sr,  t.x,
                      s.x * sr,   s.y * cr,   t.y,
                      0.0f,       0.0f,       1.0f);
    }

    bool Mat3x3::is_affine() const{
        float det = _11 * _22 - _12 * _21;
        if(isclose(det, 0)) return false;
        return _31 == 0.0f && _32 == 0.0f && _33 == 1.0f;
    }

    Vec2 Mat3x3::_t() const { return Vec2(_13, _23); }
    float Mat3x3::_r() const { return atan2f(_21, _11); }
    Vec2 Mat3x3::_s() const {
        return Vec2(
            sqrtf(_11 * _11 + _21 * _21),
            sqrtf(_12 * _12 + _22 * _22)
        );
    }

}   // namespace pkpy