#include "raylib.h"
#include "raymath.h"
#include <stdio.h> // 增加用於調試輸出的引用 (如果需要)
#include <stdlib.h> // 用於 rand(), malloc(), free()

#define ARENA_IMPLEMENTATION
#include "arena.h"

#define MAX_PARTICLES 30000 // 最大粒子數量
#define MAX_PARTICLES_PER_NODE 4 // 四叉樹每個葉節點能容納的最大粒子索引數量

// 粒子組件結構，包含位置和速度
typedef struct {
    Vector2 position;
    Vector2 velocity;
} ParticleComponent;

// 粒子系統結構，包含所有粒子和當前粒子數量
typedef struct {
    ParticleComponent particles[MAX_PARTICLES]; // 粒子數據陣列
    int count; // 當前粒子數量 (此範例中固定為 MAX_PARTICLES)
} ParticleSystem;

// 四叉樹節點結構
typedef struct QuadTreeNode {
    Rectangle bounds; // 此節點代表的空間範圍
    int particles[MAX_PARTICLES_PER_NODE]; // 儲存在此葉節點中的粒子 *索引* (僅葉節點有效)
    int count; // 儲存在此節點的粒子數量 (僅葉節點有效)
    struct QuadTreeNode* children[4]; // 子節點指標 (左上, 右上, 左下, 右下)
                                      // 如果 children[0] == NULL，表示這是葉節點
} QuadTreeNode;

// --- 函數原型宣告 ---
// 初始化粒子系統
void InitParticleSystem(ParticleSystem* system);
// 更新粒子系統狀態 (包含物理模擬和四叉樹交互)
void UpdateParticleSystem_NoQuadTree(ParticleSystem* system, Vector2 mousePos, float deltaTime);
void UpdateParticleSystem(ParticleSystem* system, Vector2 mousePos, float deltaTime);
// 繪製粒子系統
void DrawParticleSystem(ParticleSystem* system);
// 創建一個新的四叉樹節點
QuadTreeNode* CreateQuadTreeNode(Rectangle bounds);
// 分裂一個滿載的四叉樹葉節點
void SplitQuadTreeNode(QuadTreeNode* node, ParticleSystem* system);
// 將粒子插入到四叉樹節點 (會遞迴處理)
void InsertParticleToNode(QuadTreeNode* node, int particleIndex, ParticleSystem* system);
// 建立整個四叉樹
QuadTreeNode* BuildQuadTree(ParticleSystem* system, Rectangle bounds);
// 查詢四叉樹中與指定矩形範圍重疊的粒子
void QueryQuadTree(QuadTreeNode* node, Rectangle queryRect, int* result, int* count);
// 釋放整個四叉樹的記憶體
void FreeQuadTree(QuadTreeNode* node);

// 主函數
Arena arena = { 0 };

int main()
{
    // 螢幕尺寸
    const int screenWidth = 800;
    const int screenHeight = 600;
    arena_init(&arena);
    // 初始化 Raylib 窗口
    InitWindow(screenWidth, screenHeight, "粒子系統與四叉樹優化 (Particle System with QuadTree)");
    // 創建粒子系統實例
    ParticleSystem system;
    // 初始化粒子
    InitParticleSystem(&system);
    // 設定目標幀率
    SetTargetFPS(60);
    // 主遊戲迴圈
    while (!WindowShouldClose()) { // 持續執行直到使用者關閉視窗
        // 獲取幀時間差 (用於物理計算)
        float deltaTime = GetFrameTime();
        // 獲取滑鼠當前位置
        Vector2 mousePos = GetMousePosition();
        // 更新粒子系統狀態
        UpdateParticleSystem(&system, mousePos, deltaTime);
        // UpdateParticleSystem_NoQuadTree(&system, mousePos, deltaTime);
        //  開始繪圖
        BeginDrawing();
        // 清空背景為黑色
        ClearBackground(BLACK);
        // 繪製所有粒子
        DrawParticleSystem(&system);
        // 在左上角顯示當前 FPS
        DrawFPS(10, 10);
        // 結束繪圖
        EndDrawing();
    }
    arena_destroy(&arena);
    // 關閉 Raylib 窗口並釋放資源
    CloseWindow();
    // 程式正常結束
    return 0;
}

// 初始化粒子系統
void InitParticleSystem(ParticleSystem* system)
{
    system->count = MAX_PARTICLES; // 設定粒子數量
    // 遍歷所有粒子進行初始化
    for (int i = 0; i < MAX_PARTICLES; i++) {
        // 在螢幕範圍內隨機設定初始位置
        system->particles[i].position.x = rand() % 800;
        system->particles[i].position.y = rand() % 600;
        // 設定隨機初始速度 (-10 到 +10 之間)
        system->particles[i].velocity.x = (rand() % 100 - 50) / 5.0f;
        system->particles[i].velocity.y = (rand() % 100 - 50) / 5.0f;
    }
}
void UpdateParticleSystem_NoQuadTree(ParticleSystem* system, Vector2 mousePos, float deltaTime)
{
    // --- 直接處理所有粒子與滑鼠的互動 ---
    for (int i = 0; i < system->count; i++) {
        ParticleComponent* p = &system->particles[i];

        // 計算粒子到滑鼠的向量和距離 (與原碼相同)
        Vector2 toMouse = Vector2Subtract(mousePos, p->position);
        float distance = Vector2Length(toMouse);

        // 如果粒子在滑鼠 100 像素範圍內 (與原碼相同)
        if (distance > 1.0f && distance < 100.0f) {
            Vector2 direction = Vector2Normalize(toMouse);
            // 施加吸引力 (與原碼相同)
            p->velocity = Vector2Add(p->velocity, Vector2Scale(direction, 50.0f * (100.0f - distance) / 100.0f * deltaTime));
            // 增加阻尼 (與原碼相同)
            p->velocity = Vector2Scale(p->velocity, 1.0f - 0.1f * deltaTime); // 可以移到下面統一處理
        }

        // 可以在這裡統一應用阻尼，即使沒受到滑鼠影響
        // p->velocity = Vector2Scale(p->velocity, 1.0f - 0.1f * deltaTime);
    }

    // --- 更新所有粒子的位置 (與原碼相同) ---
    for (int i = 0; i < system->count; i++) {
        ParticleComponent* p = &system->particles[i];
        p->position = Vector2Add(p->position, Vector2Scale(p->velocity, deltaTime));

        // 處理邊界條件 (與原碼相同)
        if (p->position.x < 0)
            p->position.x = 800;
        if (p->position.x > 800)
            p->position.x = 0;
        if (p->position.y < 0)
            p->position.y = 600;
        if (p->position.y > 600)
            p->position.y = 0;
    }
}
// 更新粒子系統狀態
void UpdateParticleSystem(ParticleSystem* system, Vector2 mousePos, float deltaTime)
{
    // --- 效能瓶頸核心區域開始 ---
    // 定義滑鼠周圍的查詢矩形範圍 (200x200)
    Rectangle queryRect = { mousePos.x - 100, mousePos.y - 100, 200, 200 };

    // **關鍵效能問題**: 每幀重新建立完整的四叉樹
    // 理想情況下，應更新現有樹或使用記憶體池來減少分配開銷
    QuadTreeNode* root = BuildQuadTree(system, (Rectangle) { 0, 0, 800, 600 }); // 根節點範圍為整個螢幕

    // 用於儲存查詢結果的粒子索引陣列
    int result[MAX_PARTICLES]; // 注意: 這裡的大小是 MAX_PARTICLES，以防萬一，但通常遠小於此
    int count = 0; // 查詢結果的數量

    // 使用四叉樹查詢在滑鼠附近的粒子
    QueryQuadTree(root, queryRect, result, &count);

    // --- 僅處理靠近滑鼠的粒子 ---
    // 這個迴圈只遍歷被四叉樹查詢到的粒子，數量相對較少
    for (int i = 0; i < count; i++) {
        int index = result[i]; // 獲取粒子在 system->particles 中的索引
        ParticleComponent* p = &system->particles[index]; // 獲取該粒子的指針

        // 計算粒子到滑鼠的向量
        Vector2 toMouse = Vector2Subtract(mousePos, p->position);
        // 計算粒子到滑鼠的距離
        float distance = Vector2Length(toMouse);

        // 如果粒子在滑鼠 100 像素範圍內
        if (distance > 1.0f && distance < 100.0f) { // 避免除以零或距離過近時力過大
            // 計算從粒子指向滑鼠的單位向量 (方向)
            Vector2 direction = Vector2Normalize(toMouse);
            // 施加一個排斥力 (方向相反)，力的大小與距離成反比
            // 使用 deltaTime 來確保力的作用與幀率無關
            // p->velocity = Vector2Subtract(p->velocity, Vector2Scale(direction, 100.0f / distance * deltaTime)); // 原始排斥
            p->velocity = Vector2Add(p->velocity, Vector2Scale(direction, 50.0f * (100.0f - distance) / 100.0f * deltaTime)); // 吸引力示例
        }
        // 可以增加阻尼效果，讓粒子移動更平滑
        p->velocity = Vector2Scale(p->velocity, 1.0f - 0.1f * deltaTime);
    }

    // --- 更新所有粒子的位置 ---
    // 這個迴圈仍然需要遍歷所有粒子
    for (int i = 0; i < system->count; i++) {
        ParticleComponent* p = &system->particles[i];
        // 根據速度更新位置: position += velocity * deltaTime
        p->position = Vector2Add(p->position, Vector2Scale(p->velocity, deltaTime));

        // 處理邊界條件 (穿牆效果)
        if (p->position.x < 0)
            p->position.x = 800; // 從左邊出去，右邊回來
        if (p->position.x > 800)
            p->position.x = 0; // 從右邊出去，左邊回來
        if (p->position.y < 0)
            p->position.y = 600; // 從上面出去，下面回來
        if (p->position.y > 600)
            p->position.y = 0; // 從下面出去，上面回來
    }

    // **關鍵效能問題**: 每幀釋放整個四叉樹的記憶體
    FreeQuadTree(root);
    // --- 效能瓶頸核心區域結束 ---
}

// 繪製粒子系統
void DrawParticleSystem(ParticleSystem* system)
{
    // 遍歷所有粒子並繪製
    for (int i = 0; i < system->count; i++) {
        ParticleComponent* p = &system->particles[i];
        // 在粒子的位置繪製一個白色的圓點 (半徑 2 像素)
        //        DrawCircle(p->position.x, p->position.y, 2, WHITE);
        DrawPixel(p->position.x, p->position.y, WHITE);
    }
}

// 創建一個新的四叉樹節點
QuadTreeNode* CreateQuadTreeNode(Rectangle bounds)
{
    // 為新節點分配記憶體
    // 優化提示: 這裡可以使用記憶體池來代替 malloc
    //    QuadTreeNode* node = malloc(sizeof(QuadTreeNode));
    QuadTreeNode* node = arena_alloc(&arena, sizeof(QuadTreeNode));
    if (!node) {
        // 處理記憶體分配失敗的情況 (例如，打印錯誤訊息並退出)
        perror("Failed to allocate memory for QuadTreeNode");
        exit(EXIT_FAILURE);
    }
    node->bounds = bounds; // 設定節點的空間範圍
    node->count = 0; // 初始化粒子計數為 0 (僅葉節點有效)
    // 初始化所有子節點為 NULL (表示當前是葉節點)
    for (int i = 0; i < 4; i++) {
        node->children[i] = NULL;
    }
    // 注意: 'particles' 陣列不需要在此處初始化，因為它是節點結構的一部分
    return node;
}

// 分裂一個滿載的四叉樹葉節點
void SplitQuadTreeNode(QuadTreeNode* node, ParticleSystem* system)
{
    // 計算子節點的寬度和高度
    float subWidth = node->bounds.width / 2;
    float subHeight = node->bounds.height / 2;
    // 獲取父節點的左上角座標
    float x = node->bounds.x;
    float y = node->bounds.y;

    // 創建四個子節點，分別對應四個象限
    // 優化提示: 這裡也可以使用記憶體池
    node->children[0] = CreateQuadTreeNode((Rectangle) { x, y, subWidth, subHeight }); // 左上
    node->children[1] = CreateQuadTreeNode((Rectangle) { x + subWidth, y, subWidth, subHeight }); // 右上
    node->children[2] = CreateQuadTreeNode((Rectangle) { x, y + subHeight, subWidth, subHeight }); // 左下
    node->children[3] = CreateQuadTreeNode((Rectangle) { x + subWidth, y + subHeight, subWidth, subHeight }); // 右下

    // --- 將原來儲存在父節點(現在變為內部節點)的粒子重新分配到新的子節點中 ---
    // 臨時儲存父節點的粒子索引和數量
    int tempParticles[MAX_PARTICLES_PER_NODE];
    int tempCount = node->count;
    for (int i = 0; i < tempCount; ++i) {
        tempParticles[i] = node->particles[i];
    }

    // 清空父節點的粒子計數 (因為它不再是葉節點)
    node->count = 0;

    // 遍歷臨時儲存的粒子索引
    for (int i = 0; i < tempCount; i++) {
        int index = tempParticles[i]; // 獲取粒子索引
        Vector2 pos = system->particles[index].position; // 獲取粒子位置

        // 檢查粒子屬於哪個子節點，並將其插入
        for (int j = 0; j < 4; j++) {
            if (CheckCollisionPointRec(pos, node->children[j]->bounds)) {
                InsertParticleToNode(node->children[j], index, system);
                break; // 粒子只會屬於一個子節點，找到後跳出內層迴圈
            }
        }
    }
}

// 將粒子插入到四叉樹節點 (會遞迴處理)
void InsertParticleToNode(QuadTreeNode* node, int particleIndex, ParticleSystem* system)
{
    // 獲取要插入粒子的位置
    Vector2 pos = system->particles[particleIndex].position;
    // 檢查當前節點是否為葉節點 (沒有子節點)
    if (node->children[0] == NULL) { // 是葉節點
        // 檢查葉節點是否還有空間容納更多粒子
        if (node->count < MAX_PARTICLES_PER_NODE) {
            // --- 修改點: 直接存入靜態陣列 ---
            node->particles[node->count] = particleIndex; // 將粒子索引存入陣列
            node->count++; // 增加節點的粒子計數
        } else {
            // 葉節點已滿，需要分裂
            SplitQuadTreeNode(node, system);
            // 分裂後，將 *當前正要插入的* 粒子插入到對應的新子節點中
            for (int i = 0; i < 4; i++) {
                if (CheckCollisionPointRec(pos, node->children[i]->bounds)) {
                    InsertParticleToNode(node->children[i], particleIndex, system);
                    break;
                }
            }
        }
    } else { // 不是葉節點 (是內部節點)
        // 將粒子遞迴地插入到對應的子節點中
        for (int i = 0; i < 4; i++) {
            if (CheckCollisionPointRec(pos, node->children[i]->bounds)) {
                InsertParticleToNode(node->children[i], particleIndex, system);
                break;
            }
        }
    }
}

// 建立整個四叉樹
QuadTreeNode* BuildQuadTree(ParticleSystem* system, Rectangle bounds)
{
    // 創建根節點，範圍由參數指定 (通常是整個螢幕)
    QuadTreeNode* root = CreateQuadTreeNode(bounds);
    // 遍歷系統中的所有粒子
    for (int i = 0; i < system->count; i++) {
        // 將每個粒子插入到四叉樹中 (從根節點開始)
        InsertParticleToNode(root, i, system);
    }
    // 返回構建好的四叉樹的根節點
    return root;
}

// 查詢四叉樹中與指定矩形範圍重疊的粒子
void QueryQuadTree(QuadTreeNode* node, Rectangle queryRect, int* result, int* count)
{
    // 剪枝: 如果節點的範圍與查詢範圍完全沒有重疊，則無需繼續查詢此分支
    if (!CheckCollisionRecs(node->bounds, queryRect)) {
        return;
    }

    // 檢查是否為葉節點
    if (node->children[0] == NULL) { // 是葉節點
        // 遍歷此葉節點中儲存的所有粒子索引
        for (int i = 0; i < node->count; i++) {
            // 這裡可以選擇性地再次檢查粒子實際位置是否真的在 queryRect 內 (更精確)
            // if (CheckCollisionPointRec(system->particles[node->particles[i]].position, queryRect)) {
                // 將粒子索引添加到結果陣列中
                result[*count] = node->particles[i];
                (*count)++; // 增加結果計數
            // }
        }
    } else { // 是內部節點
        // 遞迴地查詢四個子節點
        for (int i = 0; i < 4; i++) {
            QueryQuadTree(node->children[i], queryRect, result, count);
        }
    }
}

// 遞迴地釋放整個四叉樹的記憶體
void FreeQuadTree(QuadTreeNode* node)
{
/*    
    // 檢查是否為內部節點
    if (node->children[0] != NULL) {
        // 遞迴地釋放所有子節點
        for (int i = 0; i < 4; i++) {
            FreeQuadTree(node->children[i]);
            node->children[i] = NULL; // 避免懸空指針 (好習慣)
        }
    }
    free(node);
*/
    arena_reset(&arena);    
}