#include <map>
#include <iostream>
#include <vector>
#include <chrono>
#include <thread>
#include <atomic>
#include <mutex>
#include <cmath>
#include <random>
#include <algorithm>

#include "../util/grid/grid3eigen.hpp"
#include "../util/output/bmpwriter.hpp"
#include "../util/output/frame.hpp"
#include "../util/noise/pnoise2.hpp"
#include "../util/noise/pnoise.cpp"
#include "../util/output/aviwriter.hpp"

#include "../imgui/imgui.h"
#include "../imgui/backends/imgui_impl_glfw.h"
#include "../imgui/backends/imgui_impl_opengl3.h"
#include <GLFW/glfw3.h>
#include "../stb/stb_image.h"

struct fluidParticle {
    Eigen::Matrix<float, 3, 1> pos;
    Eigen::Matrix<float, 3, 1> velocity;
    Eigen::Matrix<float, 3, 1> acceleration;
    Eigen::Matrix<float, 3, 1> forceAccumulator;
    float density;
    float pressure;
    float viscosity;
    float restitution;
    float mass;
};

struct gridConfig {
    float gridSizeCube = 1024;
    const float SMOOTHING_RADIUS = 32.0f;
    const float REST_DENSITY = 0.5f;
    const float TIMESTEP = 0.016f;
    const float G_ATTRACTION = 50.0f;
};

Eigen::Matrix<float, 3, 1> posGen() {
    static std::random_device rd;
    static std::mt19937 gen(rd());
    static std::normal_distribution<float> dist(0.0f, 1024.0f);
    
    return Eigen::Matrix<float, 3, 1>(dist(gen),dist(gen),dist(gen));
}

// Math Helpers for SPH Kernels
float poly6Kernel(float r, float h) {
    if (r < 0 || r > h) return 0.0f;
    float h2 = h * h;
    float r2 = r * r;
    float diff = h2 - r2;
    float coef = 315.0f / (64.0f * M_PI * pow(h, 9));
    return coef * diff * diff * diff;
}

Eigen::Vector3f spikyGradientKernel(Eigen::Vector3f r_vec, float r, float h) {
    if (r <= 0 || r > h) return Eigen::Vector3f::Zero();
    float diff = h - r;
    float coef = -45.0f / (M_PI * pow(h, 6));
    return r_vec.normalized() * coef * diff * diff;
}

Eigen::Vector3f buildGradient(float value, const std::map<float, Eigen::Vector3f>& gradientKeys) {
    auto exactMatch = gradientKeys.find(value);
    if (exactMatch != gradientKeys.end()) {
        return exactMatch->second;
    }
    auto lower = gradientKeys.lower_bound(value);
    if (lower == gradientKeys.begin()) {
        return gradientKeys.begin()->second;
    }
    if (lower == gradientKeys.end()) {
        return gradientKeys.rbegin()->second;
    }
    auto upper = lower;
    lower = std::prev(lower);
    float key1 = lower->first;
    float key2 = upper->first;
    const Eigen::Vector3f& color1 = lower->second;
    const Eigen::Vector3f& color2 = upper->second;
    float t = (value - key1) / (key2 - key1);
    t = std::clamp(t, 0.0f, 1.0f);
    return color1 + t * (color2 - color1);
}

class fluidSim {
private:
    std::unordered_map<size_t, Eigen::Matrix<float, 3, 1>> idposMap;
    float newMass = 1000;
    gridConfig config;
    int nextObjectId = 0;
    std::map<float, Eigen::Vector3f> gradientmap;
public:
    Octree<fluidParticle> grid;

    fluidSim() : grid({-config.gridSizeCube, -config.gridSizeCube, -config.gridSizeCube}, {config.gridSizeCube, config.gridSizeCube, config.gridSizeCube}) {
        grid.setBackgroundColor({0.1f, 0.1f, 0.2f});
        gradientmap.emplace(0.0, Eigen::Vector3f(1, 0, 0));
        gradientmap.emplace(0.5, Eigen::Vector3f(0, 1, 0));
        gradientmap.emplace(1.0, Eigen::Vector3f(0, 0, 1));

    }

    void spawnParticles(fluidParticle toSpawn, int count) {
        TIME_FUNCTION;
        for (int i = 0; i < count; i++) {
            Eigen::Matrix<float, 3, 1> pos = posGen();
            int id = nextObjectId++;
            
            Eigen::Vector3f color = buildGradient(toSpawn.mass / 1000, gradientmap);
            grid.set(toSpawn, pos, true, color, 1, true, id, false);
            idposMap[id] = pos;
        }
        newMass -= newMass / 10;
    }

    void applyPressureDensity() {
        #pragma omp parallel for
        for (auto& point : idposMap) {
            auto node = grid.find(point.second);
            if (!node) continue;
            std::vector<std::shared_ptr<Octree<fluidParticle>::NodeData>> neighbors = grid.findInRadius(point.second, config.SMOOTHING_RADIUS * node->data.mass);
            float density = 0.0f;
            for (const auto& neighbor : neighbors) {
                if (node == neighbor) continue;
                float dist = (point.second-neighbor->position).norm();
                density += neighbor->data.mass * poly6Kernel(dist, config.SMOOTHING_RADIUS * node->data.mass);
            }
            node->data.density = density;
            float pressure = node->data.restitution * (density - config.REST_DENSITY);
            node->data.pressure = std::max(pressure, 0.0f);
        }
    }
    
    void applyForce() {
        #pragma omp parallel for
        for (auto& point : idposMap) {
            auto node = grid.find(point.second);
            if (!node) continue;
            Eigen::Vector3f totalForce = -point.second.normalized() * 10.0f * node->data.mass;
            std::vector<std::shared_ptr<Octree<fluidParticle>::NodeData>> neighbors = grid.findInRadius(point.second, config.SMOOTHING_RADIUS * node->data.mass);

            for (const auto& neighbor : neighbors) {
                if (node == neighbor) continue;
                Eigen::Vector3f diff = point.second - neighbor->position;
                float dist = diff.norm();
                if (dist < EPSILON) continue;
                float densj = neighbor->data.density;
                if (densj < EPSILON) densj = EPSILON;
                float pressj = neighbor->data.pressure;

                Eigen::Vector3f pressureForceVec = -neighbor->data.mass * ((node->data.pressure + pressj) / (2.0f * densj)) * spikyGradientKernel(diff, dist, config.SMOOTHING_RADIUS);
                totalForce += pressureForceVec;
                Eigen::Vector3f velDiff = neighbor->data.velocity - node->data.velocity;
                float viscositCoef = node->data.viscosity * neighbor->data.mass * (1.0f / densj) * poly6Kernel(dist, config.SMOOTHING_RADIUS);
                totalForce += velDiff * viscositCoef;

                float clampDist = std::max(dist, 5.0f);
                Eigen::Vector3f dirToNeighbor = (neighbor->position - node->position).normalized();
                float attract = (config.G_ATTRACTION * node->data.mass * neighbor->data.mass) / (clampDist * clampDist);
                totalForce += dirToNeighbor * attract;
            }
            node->data.forceAccumulator = totalForce;
        }
    }

    void replaceLost() {
        int toReplace = 0;
        #pragma omp parallel for
        for (auto& point : idposMap) { 
            if (!grid.inGrid(point.second)) {
                grid.remove(point.second);
                toReplace++;
            }
        }
        fluidParticle newParticles;
        newParticles.mass = newMass;
        spawnParticles(newParticles, toReplace);
    }

    void applyPhysics() {
        TIME_FUNCTION;
            ///TODO: 
            // apply velocity to have particle move
            // use mass of neighboring particles to change direction towards them
            // ressure pushes particles away if they are too close
            // resitution controls how much pressure builds up based on how close particles are
            // have to find a way to have a clump use their combined mass instead of individual mass.
            //if a particle leaves the grid, destroy it and spawn a new one
    }
};