#include <c_Matrix.h>
#include <c_Memory.h>

// 初始化矩阵（由外部传入结构体指针 self，内部只分配绑定的 c_Array 内存）
c_err_t c_Matrix_Init(c_Matrix_t* self, c_size_t rows, c_size_t cols, c_size_t element_size) {
    if (!self) return C_ERR_PARAM;
    if (rows == 0 || cols == 0 || element_size == 0) return C_ERR_PARAM;

    // 分配底层封装的一维泛型数组
    self->array = c_Array_New(rows * cols, element_size);
    if (!self->array) {
        return C_ERR_NOMEM;
    }

    self->rows = rows;
    self->cols = cols;
    return C_SUCCESS;
}

// 销毁矩阵（释放内部资源，不释放 self 本身，由外部决定 self 释放方式）
void c_Matrix_Destroy(c_Matrix_t* self) {
    if (self) {
        if (self->array) {
            c_Array_Delete(&self->array);
            self->array = NULL;
        }
        self->rows = 0;
        self->cols = 0;
    }
}

// 获取单个元素字节大小
c_size_t c_Matrix_ElementSize(c_Matrix_t* self) {
    if (!self || !self->array) return 0;
    return c_Array_Size(self->array);
}

// 获取矩阵元素（通过二级指针 value 返回该元素的内部指针，并返回状态码）
c_err_t c_Matrix_Get(c_Matrix_t* self, c_size_t row, c_size_t col, void** value) {
    if (!self || !value) return C_ERR_PARAM;
    if (row >= self->rows || col >= self->cols) return C_ERR_PARAM;

    // 计算一维索引：index = row * cols + col
    const c_size_t index = row * self->cols + col;

    c_err_t err = c_Array_Get(self->array, index, value);
    if (err!=C_ERR_OK) return err;

    return C_SUCCESS;
}

// 写入矩阵元素
c_err_t c_Matrix_Put(c_Matrix_t* self, c_size_t row, c_size_t col, void* value) {
    if (!self || !value) return C_ERR_PARAM;
    if (row >= self->rows || col >= self->cols) return C_ERR_PARAM;

    const c_size_t index = row * self->cols + col;
    return c_Array_Put(self->array, index, value);
}

// 矩阵转置：dst 必须是已经初始化好的大小为 (src->cols x src->rows) 的矩阵
c_err_t c_Matrix_TransposeTo(c_Matrix_t* src, c_Matrix_t* dst) {
    if (!src || !dst) return C_ERR_PARAM;
    if (src->cols != dst->rows || src->rows != dst->cols) return C_ERR_PARAM;

    for (c_size_t i = 0; i < src->rows; i++) {
        for (c_size_t j = 0; j < src->cols; j++) {
            void* temp = NULL;
            c_Matrix_Get(src, i, j, (void**)&temp);
            c_Matrix_Put(dst, j, i, temp); // 行列互换写入
        }
    }

    return C_SUCCESS;
}

// 矩阵乘法：C = A * B （以 float 为例）
c_err_t c_Matrix_Multiply(const c_Matrix_t* A, const c_Matrix_t* B, c_Matrix_t* C, void (*multiply)(const void* a, const void* b, void** c)) {
    // 1. 基础健壮性检查
    if (!A || !B || !C || !multiply) return C_ERR_PARAM;

    // 2. 矩阵乘法维度匹配检查 (A的列数必须等于B的行数，C的尺寸必须是 A.rows x B.cols)
    if (A->cols != B->rows || C->rows != A->rows || C->cols != B->cols) {
        return C_ERR_PARAM;
    }

    const c_size_t elem_size = c_Matrix_ElementSize((c_Matrix_t*)A);

    // 4. 三重循环计算矩阵乘法
    for (c_size_t i = 0; i < A->rows; i++) {
        for (c_size_t j = 0; j < B->cols; j++) {

            // 在计算当前 C[i][j] 的点积前，必须先获取 C[i][j] 现有的物理指针
            void* c_item_ptr = NULL;
            c_Matrix_Get(C, i, j, &c_item_ptr);

            // 将当前位置的数据清零（初始化累加器）
            memset(c_item_ptr, 0, elem_size);

            for (c_size_t k = 0; k < A->cols; k++) {
                void *a_item_ptr = NULL;
                void *b_item_ptr = NULL;

                // 安全获取 A[i][k] 和 B[k][j]
                c_Matrix_Get((c_Matrix_t*)A, i, k, &a_item_ptr);
                c_Matrix_Get((c_Matrix_t*)B, k, j, &b_item_ptr);

                // 核心：调用用户自定义的单元素乘法
                // 指针传参解释：
                // &prod_res 是一个二级指针 (void**)，回调函数内部会将计算结果写入到 prod_res 指向的缓冲区
                void* ctx_ptr = c_item_ptr;
                multiply(a_item_ptr, b_item_ptr, (void**)&ctx_ptr);

                // 执行累加：这里需要对通用泛型执行加法
                // 由于 C 语言泛型限制，我们在这里进行一维数组层面的二进制或特定类型累加。
                // 工业级做法通常会将“加法”也作为回调，或者假设前几个字节可加。
                // 针对最常见的数值类型，我们可以根据用户传进来的回调结果，在当前位置累加：
                // 为使设计更严谨，通常让 multiply 内部直接实现：*(Type*)c_item_ptr += *(Type*)a * *(Type*)b
                // 如果你的 multiply 语义是：*c = a * b，则需要外部进行特定类型的累加，例如以 float 为例：
                // *(float*)c_item_ptr += *(float*)prod_res;
            }
        }
    }

    return C_SUCCESS;
}

/* ------------------------------------------------------------------------------------------------------------------ */
/*  */

// 辅助函数：根据一维索引，计算转置后该元素应当去往的新一维索引
// 公式：新行 = 旧列，新列 = 旧行 -> new_index = (old_index % cols) * rows + (old_index / cols)
C_STATIC_FORCE_INLINE
c_size_t get_transposed_source_index(c_size_t curr_index, c_size_t old_rows, c_size_t old_cols) {
    return (curr_index % old_rows) * old_cols + (curr_index / old_rows);
}

c_err_t c_Matrix_InPlaceTranspose(c_Matrix_t* self) {
    if (!self || !self->array) return C_ERR_PARAM;

    c_size_t rows = self->rows;
    c_size_t cols = self->cols;
    c_size_t elem_size = c_Matrix_ElementSize(self);

    if (rows <= 1 && cols <= 1) {
        return C_SUCCESS;
    }

    // --------------------------------------------------------------------------------
    // 场景 A：方阵原地转置（不变，原本就是正确的）
    // --------------------------------------------------------------------------------
    if (rows == cols) {
#define SWAP_BUF_SIZE 256
        uint8_t swap_buf[SWAP_BUF_SIZE];
        uint8_t* temp = swap_buf;
        if (elem_size > SWAP_BUF_SIZE) {
            temp = (uint8_t*)C_ALLOC(elem_size);
            if (!temp) return C_ERR_NOMEM;
        }
        for (c_size_t i = 0; i < rows; i++) {
            for (c_size_t j = i + 1; j < cols; j++) {
                void *cell_a = NULL, *cell_b = NULL;
                c_Matrix_Get(self, i, j, &cell_a);
                c_Matrix_Get(self, j, i, &cell_b);
                memcpy(temp, cell_a, elem_size);
                memcpy(cell_a, cell_b, elem_size);
                memcpy(cell_b, temp, elem_size);
            }
        }
        if (elem_size > SWAP_BUF_SIZE) C_FREE(temp);
#undef SWAP_BUF_SIZE
        return C_SUCCESS;
    }

    // --------------------------------------------------------------------------------
    // 场景 B：非方阵原地转置（已完全修复环路拉取逻辑）
    // --------------------------------------------------------------------------------
    c_size_t total_elements = rows * cols;

    c_size_t bitmask_size = (total_elements + 7) / 8;
    uint8_t* visited = (uint8_t*)C_CALLOC(bitmask_size, sizeof(uint8_t));
    if (!visited) return C_ERR_NOMEM;

    uint8_t* cycle_buf = (uint8_t*)C_ALLOC(elem_size);
    if (!cycle_buf) {
        C_FREE(visited);
        return C_ERR_NOMEM;
    }

    void* base_data = NULL;
    c_err_t array_err = c_Array_Get(self->array, 0, &base_data);
    if (array_err != C_SUCCESS || !base_data) {
        C_FREE(cycle_buf);
        C_FREE(visited);
        return array_err;
    }

    for (c_size_t i = 0; i < total_elements; i++) {
        if (visited[i / 8] & (1 << (i % 8))) {
            continue;
        }

        c_size_t curr_idx = i;
        // 使用修正后的逆向映射函数
        c_size_t next_idx = get_transposed_source_index(curr_idx, rows, cols);

        if (next_idx == curr_idx) {
            visited[curr_idx / 8] |= (1 << (curr_idx % 8));
            continue;
        }

        // 暂存当前起点元素
        memcpy(cycle_buf, (char*)base_data + (curr_idx * elem_size), elem_size);

        // 沿着逆向数据源环路追溯拉取
        while (next_idx != i) {
            void* src = (char*)base_data + (next_idx * elem_size);
            void* dst = (char*)base_data + (curr_idx * elem_size);

            memcpy(dst, src, elem_size); // 正确地拉取源数据
            visited[curr_idx / 8] |= (1 << (curr_idx % 8));

            curr_idx = next_idx;
            next_idx = get_transposed_source_index(curr_idx, rows, cols);
        }

        // 闭合环路
        void* dst = (char*)base_data + (curr_idx * elem_size);
        memcpy(dst, cycle_buf, elem_size);
        visited[curr_idx / 8] |= (1 << (curr_idx % 8));
    }

    C_FREE(cycle_buf);
    C_FREE(visited);

    // 修改元数据宽高
    self->rows = cols;
    self->cols = rows;

    return C_SUCCESS;
}

c_err_t c_Matrix_Determinant(const c_Matrix_t* self, const c_MatrixOps_t* ops, void* out_det) {
    if (!self || !ops || !out_det) return C_ERR_PARAM;
    if (self->rows != self->cols) return C_ERR_PARAM;

    c_size_t n = self->rows;
    c_size_t elem_size = c_Matrix_ElementSize((c_Matrix_t*)self);

    if (n == 1) {
        void* cell = NULL;
        c_Matrix_Get((c_Matrix_t*)self, 0, 0, &cell);
        memcpy(out_det, cell, elem_size);
        return C_SUCCESS;
    }

    c_Matrix_t temp_mat;
    c_err_t err = c_Matrix_Init(&temp_mat, n, n, elem_size);
    if (err != C_SUCCESS) return err;

    void* src_base = NULL;
    void* dst_base = NULL;
    c_Array_Get(self->array, 0, &src_base);
    c_Array_Get(temp_mat.array, 0, &dst_base);
    memcpy(dst_base, src_base, n * n * elem_size);

    ops->one(out_det);
    int sign = 1;

    // 分配真正独立的单元素计算缓冲区
    void* pivot_val = C_ALLOC(elem_size);
    void* factor = C_ALLOC(elem_size);
    void* temp_val = C_ALLOC(elem_size);

    if (!pivot_val || !factor || !temp_val) {
        C_FREE(pivot_val);
        C_FREE(factor);
        C_FREE(temp_val);
        return C_ERR_NOMEM;
    }

    for (c_size_t i = 0; i < n; i++) {
        // --- 部分主元选择 ---
        c_size_t pivot_row = i;
        void* max_cell = NULL;
        c_Matrix_Get(&temp_mat, i, i, &max_cell);

        for (c_size_t k = i + 1; k < n; k++) {
            void* check_cell = NULL;
            c_Matrix_Get(&temp_mat, k, i, &check_cell);
            if (ops->compare_abs(check_cell, max_cell) > 0) {
                max_cell = check_cell;
                pivot_row = k;
            }
        }

        if (ops->is_zero(max_cell)) {
            ops->zero(out_det);
            goto _cleanup;
        }

        if (pivot_row != i) {
            for (c_size_t j = i; j < n; j++) {
                void *cell_a = NULL, *cell_b = NULL;
                c_Matrix_Get(&temp_mat, i, j, &cell_a);
                c_Matrix_Get(&temp_mat, pivot_row, j, &cell_b);

                memcpy(temp_val, cell_a, elem_size);
                memcpy(cell_a, cell_b, elem_size);
                memcpy(cell_b, temp_val, elem_size);
            }
            sign = -sign;
        }

        // --- 🛠️ 核心修复点：安全获取当前主元的数据 ---
        void* internal_pivot_ptr = NULL;
        c_Matrix_Get(&temp_mat, i, i, &internal_pivot_ptr); // 此时被覆盖的是临时的 internal_pivot_ptr
        memcpy(pivot_val, internal_pivot_ptr, elem_size);   // 将真实数据拷贝到我们的 C_ALLOC 缓冲区中，保持 pivot_val 的指针值不变

        // --- 消元 ---
        for (c_size_t k = i + 1; k < n; k++) {
            void* current_col_cell = NULL;
            c_Matrix_Get(&temp_mat, k, i, &current_col_cell);

            if (ops->is_zero(current_col_cell)) continue;

            ops->div(current_col_cell, pivot_val, factor);

            for (c_size_t j = i; j < n; j++) {
                void *row_i_cell = NULL, *row_k_cell = NULL;
                c_Matrix_Get(&temp_mat, i, j, &row_i_cell);
                c_Matrix_Get(&temp_mat, k, j, &row_k_cell);

                ops->mul(row_i_cell, factor, temp_val);
                ops->sub(row_k_cell, temp_val, row_k_cell);
            }
        }

        // 累乘对角线元素
        ops->mul(out_det, pivot_val, out_det);
    }

    if (sign == -1) {
        ops->zero(temp_val);
        ops->sub(temp_val, out_det, out_det);
    }

_cleanup:
    // 此时的指针地址完美保持初始 C_ALLOC 状态，可以安全 C_FREE！
    C_FREE(pivot_val);
    C_FREE(factor);
    C_FREE(temp_val);
    c_Matrix_Destroy(&temp_mat);
    return C_SUCCESS;
}

c_err_t c_Matrix_Solve(const c_Matrix_t* A, const c_Matrix_t* b, c_Matrix_t* x, const c_MatrixOps_t* ops) {
    // 1. 基础健壮性检查
    if (!A || !b || !x || !ops) return C_ERR_PARAM;
    if (A->rows != A->cols) return C_ERR_PARAM;          // A 必须是方阵
    if (b->rows != A->rows || b->cols != 1) return C_ERR_PARAM; // b 必须是 Nx1
    if (x->rows != A->rows || x->cols != 1) return C_ERR_PARAM; // x 必须是 Nx1

    c_size_t n = A->rows;
    c_size_t elem_size = c_Matrix_ElementSize((c_Matrix_t*)A);

    // 2. 初始化 Nx(N+1) 的增广矩阵
    c_Matrix_t aug;
    c_err_t err = c_Matrix_Init(&aug, n, n + 1, elem_size);
    if (err != C_SUCCESS) return err;

    // 填充增广矩阵：前 n 列拷贝 A，第 n+1 列拷贝 b
    for (c_size_t i = 0; i < n; i++) {
        for (c_size_t j = 0; j < n; j++) {
            void* cell_A = NULL;
            c_Matrix_Get((c_Matrix_t*)A, i, j, &cell_A);
            c_Matrix_Put(&aug, i, j, cell_A);
        }
        void* cell_b = NULL;
        c_Matrix_Get((c_Matrix_t*)b, i, 0, &cell_b);
        c_Matrix_Put(&aug, i, n, cell_b); // 最后一列
    }

    // 分配独立的计算缓冲区，严格保护指针地址不被 Get 覆盖
    void* pivot_val = C_ALLOC(elem_size);
    void* factor    = C_ALLOC(elem_size);
    void* temp_val  = C_ALLOC(elem_size);
    void* sum_val   = C_ALLOC(elem_size);

    if (!pivot_val || !factor || !temp_val || !sum_val) {
        C_FREE(pivot_val);
        C_FREE(factor);
        C_FREE(temp_val);
        C_FREE(sum_val);
        return C_ERR_NOMEM;
    }

    // 3. 高斯消元主循环（化为上三角矩阵）
    for (c_size_t i = 0; i < n; i++) {
        // --- 3.1 部分主元选择（选绝对值最大的行交换上来，提高数值稳定性） ---
        c_size_t pivot_row = i;
        void* max_cell = NULL;
        c_Matrix_Get(&aug, i, i, &max_cell);

        for (c_size_t k = i + 1; k < n; k++) {
            void* check_cell = NULL;
            c_Matrix_Get(&aug, k, i, &check_cell);
            if (ops->compare_abs(check_cell, max_cell) > 0) {
                max_cell = check_cell;
                pivot_row = k;
            }
        }

        // 如果主元极度接近 0，说明矩阵奇异（无解或无数解）
        if (ops->is_zero(max_cell)) {
            err = C_ERR_SINGULAR;
            goto _cleanup;
        }

        // 行交换
        if (pivot_row != i) {
            for (c_size_t j = i; j <= n; j++) {
                void *cell_a = NULL, *cell_b = NULL;
                c_Matrix_Get(&aug, i, j, &cell_a);
                c_Matrix_Get(&aug, pivot_row, j, &cell_b);

                memcpy(temp_val, cell_a, elem_size);
                memcpy(cell_a, cell_b, elem_size);
                memcpy(cell_b, temp_val, elem_size);
            }
        }

        // --- 3.2 保护性安全读取当前行的主元数据 ---
        void* internal_pivot_ptr = NULL;
        c_Matrix_Get(&aug, i, i, &internal_pivot_ptr);
        memcpy(pivot_val, internal_pivot_ptr, elem_size); // 复制数据，保证 pivot_val 的 C_ALLOC 指针不被踩坏

        // --- 3.3 消元过程 ---
        for (c_size_t k = i + 1; k < n; k++) {
            void* current_col_cell = NULL;
            c_Matrix_Get(&aug, k, i, &current_col_cell);

            if (ops->is_zero(current_col_cell)) continue;

            // factor = aug[k][i] / pivot_val
            ops->div(current_col_cell, pivot_val, factor);

            for (c_size_t j = i; j <= n; j++) {
                void *row_i_cell = NULL, *row_k_cell = NULL;
                c_Matrix_Get(&aug, i, j, &row_i_cell);
                c_Matrix_Get(&aug, k, j, &row_k_cell);

                ops->mul(row_i_cell, factor, temp_val); // temp = aug[i][j] * factor
                ops->sub(row_k_cell, temp_val, row_k_cell); // aug[k][j] -= temp
            }
        }
    }

    // 4. 回代法求解（Back Substitution）
    // 公式：x[i] = (aug[i][n] - sum(aug[i][j] * x[j])) / aug[i][i]
    for (int i = (int)n - 1; i >= 0; i--) {
        ops->zero(sum_val); // sum = 0

        for (c_size_t j = (c_size_t)i + 1; j < n; j++) {
            void *aug_cell = NULL, *x_cell = NULL;
            c_Matrix_Get(&aug, (c_size_t)i, j, &aug_cell);
            c_Matrix_Get(x, j, 0, &x_cell);

            ops->mul(aug_cell, x_cell, temp_val); // temp = aug[i][j] * x[j]
            ops->add(sum_val, temp_val, sum_val); // sum += temp
        }

        void *aug_b_cell = NULL, *aug_diag_cell = NULL;
        c_Matrix_Get(&aug, (c_size_t)i, n, &aug_b_cell);       // 最后一列的常数项
        c_Matrix_Get(&aug, (c_size_t)i, (c_size_t)i, &aug_diag_cell); // 对角线主元

        ops->sub(aug_b_cell, sum_val, temp_val); // temp = b_item - sum

        void* x_dest_slot = NULL;
        c_Matrix_Get(x, (c_size_t)i, 0, &x_dest_slot);
        ops->div(temp_val, aug_diag_cell, x_dest_slot); // x[i] = temp / aug[i][i]
    }

    err = C_SUCCESS;

_cleanup:
    C_FREE(pivot_val);
    C_FREE(factor);
    C_FREE(temp_val);
    C_FREE(sum_val);
    c_Matrix_Destroy(&aug);
    return err;
}