// main.cpp
// A minimal physically based renderer with ImGui integration,
// normal mapping (with per-vertex tangents),
// YAML-based material loading and material saving via yaml-cpp.
// Uses OpenGL, GLFW, GLEW, GLM, ImGui (and IM_PI), stb_image, and yaml-cpp.
// Compile (e.g. on Linux) with: 
//   g++ main.cpp -lglfw -lGLEW -lGL -ldl -limgui -lyaml-cpp -o pbr_renderer
// (Adjust include paths and linker flags as needed.)

#include <GL/glew.h>
#include <GLFW/glfw3.h>

#include <glm/glm.hpp>
#include <glm/gtc/matrix_transform.hpp>
#include <glm/gtc/type_ptr.hpp>

#include <fstream>
#include <iostream>
#include <vector>
#include <cmath>
#include <cstring>
#include <string>

// ImGui includes.
#include "imgui.h"
#include "imgui_internal.h"   // for IM_PI (ImGui's internal constant)
#include "imgui_impl_glfw.h"
#include "imgui_impl_opengl3.h"

// stb_image for texture loading.
#define STB_IMAGE_IMPLEMENTATION
#include "stb_image.h"

// yaml-cpp for YAML material loading and saving.
#include <yaml-cpp/yaml.h>

// --------------------------------------------------------------------------------------
// Shader sources
// --------------------------------------------------------------------------------------

// --- Vertex Shader ---
// Now accepts a tangent attribute (location = 3) so we can compute a TBN matrix.
const char* vertexShaderSource = R"(
#version 330 core
layout (location = 0) in vec3 aPos;
layout (location = 1) in vec3 aNormal;
layout (location = 2) in vec2 aTexCoords;
layout (location = 3) in vec3 aTangent;

out vec3 WorldPos;
out vec3 Normal;
out vec2 TexCoords;
out vec3 Tangent;

uniform mat4 model;
uniform mat4 view;
uniform mat4 projection;

void main(){
    WorldPos = vec3(model * vec4(aPos, 1.0));
    Normal = mat3(transpose(inverse(model))) * aNormal;
    // For tangent, we simply transform with the model matrix (ignoring potential scale issues).
    Tangent = mat3(model) * aTangent;
    TexCoords = aTexCoords;
    gl_Position = projection * view * vec4(WorldPos, 1.0);
}
)";

// --- Fragment Shader ---
// Uses a uniform for PI (set from ImGui's internal IM_PI) and supports texture maps.
// It now also supports a normal map: if enabled, it uses TBN to compute a perturbed normal.
const char* fragmentShaderSource = R"(
#version 330 core
out vec4 FragColor;

in vec3 WorldPos;
in vec3 Normal;
in vec2 TexCoords;
in vec3 Tangent;

uniform vec3 camPos;

// Base material parameters.
uniform vec3 albedo;
uniform float metallic;
uniform float roughness;
uniform float ao;

// Texture map usage flags and samplers.
uniform bool useAlbedoMap;
uniform bool useMetallicMap;
uniform bool useRoughnessMap;
uniform bool useAOMap;
uniform bool useNormalMap;

uniform sampler2D albedoMap;
uniform sampler2D metallicMap;
uniform sampler2D roughnessMap;
uniform sampler2D aoMap;
uniform sampler2D normalMap;

// Use ImGui's internal PI constant.
uniform float PI;

// Light parameters.
uniform vec3 lightPos;
uniform vec3 lightColor;

vec3 fresnelSchlick(float cosTheta, vec3 F0) {
    return F0 + (1.0 - F0) * pow(1.0 - cosTheta, 5.0);
}

float DistributionGGX(vec3 N, vec3 H, float roughness) {
    float a      = roughness * roughness;
    float a2     = a * a;
    float NdotH  = max(dot(N, H), 0.0);
    float NdotH2 = NdotH * NdotH;
    float denom  = (NdotH2 * (a2 - 1.0) + 1.0);
    denom        = PI * denom * denom;
    return a2 / max(denom, 0.001);
}

float GeometrySchlickGGX(float NdotV, float roughness) {
    float r = roughness + 1.0;
    float k = (r * r) / 8.0;
    return NdotV / (NdotV * (1.0 - k) + k);
}

float GeometrySmith(vec3 N, vec3 V, vec3 L, float roughness) {
    float NdotV = max(dot(N, V), 0.0);
    float NdotL = max(dot(N, L), 0.0);
    float ggx1  = GeometrySchlickGGX(NdotV, roughness);
    float ggx2  = GeometrySchlickGGX(NdotL, roughness);
    return ggx1 * ggx2;
}

void main(){
    // Start with the geometry normal.
    vec3 N = normalize(Normal);
    
    // If a normal map is used, sample the normal map and compute a perturbed normal.
    if (useNormalMap) {
        vec3 tangentNormal = texture(normalMap, TexCoords).rgb;
        tangentNormal = tangentNormal * 2.0 - 1.0; // Remap from [0,1] to [-1,1]
        vec3 T = normalize(Tangent);
        vec3 B = normalize(cross(N, T));
        mat3 TBN = mat3(T, B, N);
        N = normalize(TBN * tangentNormal);
    }
    
    vec3 V = normalize(camPos - WorldPos);

    // Get material parameter values from textures (if enabled) or use fallback values.
    vec3 albedoValue = albedo;
    if (useAlbedoMap)
        albedoValue = texture(albedoMap, TexCoords).rgb;
    
    float metallicValue = metallic;
    if (useMetallicMap)
        metallicValue = texture(metallicMap, TexCoords).r;
    
    float roughnessValue = roughness;
    if (useRoughnessMap)
        roughnessValue = texture(roughnessMap, TexCoords).r;
    
    float aoValue = ao;
    if (useAOMap)
        aoValue = texture(aoMap, TexCoords).r;

    // Calculate reflectance at normal incidence.
    vec3 F0 = vec3(0.04);
    F0 = mix(F0, albedoValue, metallicValue);

    vec3 L = normalize(lightPos - WorldPos);
    vec3 H = normalize(V + L);

    float NDF = DistributionGGX(N, H, roughnessValue);
    float G   = GeometrySmith(N, V, L, roughnessValue);
    vec3 F    = fresnelSchlick(max(dot(H, V), 0.0), F0);

    vec3 numerator    = NDF * G * F;
    float denominator = 4.0 * max(dot(N, V), 0.0) * max(dot(N, L), 0.0) + 0.001;
    vec3 specular     = numerator / denominator;

    vec3 kS = F;
    vec3 kD = vec3(1.0) - kS;
    kD *= 1.0 - metallicValue;

    float NdotL = max(dot(N, L), 0.0);
    vec3 irradiance = lightColor * NdotL;
    vec3 diffuse = (albedoValue / PI);

    vec3 color = (kD * diffuse + specular) * irradiance;
    vec3 ambient = vec3(0.03) * albedoValue * aoValue;
    color += ambient;

    // Simple HDR tonemapping & gamma correction.
    color = color / (color + vec3(1.0));
    color = pow(color, vec3(1.0/2.2));

    FragColor = vec4(color, 1.0);
}
)";

// --------------------------------------------------------------------------------------
// Helper functions: Shader compilation and linking.
// --------------------------------------------------------------------------------------
GLuint compileShader(GLenum shaderType, const char* source) {
    GLuint shader = glCreateShader(shaderType);
    glShaderSource(shader, 1, &source, nullptr);
    glCompileShader(shader);
    
    GLint success;
    glGetShaderiv(shader, GL_COMPILE_STATUS, &success);
    if (!success){
        char infoLog[512];
        glGetShaderInfoLog(shader, 512, nullptr, infoLog);
        std::cerr << "Shader compilation failed:\n" << infoLog << std::endl;
    }
    return shader;
}

GLuint createProgram(GLuint vertexShader, GLuint fragmentShader) {
    GLuint program = glCreateProgram();
    glAttachShader(program, vertexShader);
    glAttachShader(program, fragmentShader);
    glLinkProgram(program);
    
    GLint success;
    glGetProgramiv(program, GL_LINK_STATUS, &success);
    if (!success){
        char infoLog[512];
        glGetProgramInfoLog(program, 512, nullptr, infoLog);
        std::cerr << "Program linking failed:\n" << infoLog << std::endl;
    }
    return program;
}

// --------------------------------------------------------------------------------------
// Texture loading helper using stb_image.
// --------------------------------------------------------------------------------------
GLuint LoadTexture(const char* path) {
    int width, height, nrChannels;
    stbi_set_flip_vertically_on_load(true);
    unsigned char* data = stbi_load(path, &width, &height, &nrChannels, 0);
    if (!data) {
        std::cout << "Failed to load texture: " << path << std::endl;
        return 0;
    }
    GLenum format;
    if(nrChannels == 1)
        format = GL_RED;
    else if(nrChannels == 3)
        format = GL_RGB;
    else if(nrChannels == 4)
        format = GL_RGBA;
        
    GLuint textureID;
    glGenTextures(1, &textureID);
    glBindTexture(GL_TEXTURE_2D, textureID);
    glTexImage2D(GL_TEXTURE_2D, 0, format, width, height, 0, format, GL_UNSIGNED_BYTE, data);
    glGenerateMipmap(GL_TEXTURE_2D);
    // Texture parameters.
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_REPEAT);	
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_REPEAT);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR_MIPMAP_LINEAR);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);
    
    stbi_image_free(data);
    return textureID;
}

// --------------------------------------------------------------------------------------
// Sphere mesh generation with tangents.
// Produces per-vertex data: position (3), normal (3), texcoords (2), tangent (3) → 11 floats/vertex.
// --------------------------------------------------------------------------------------
void generateSphere(std::vector<float>& vertices, std::vector<unsigned int>& indices,
                    unsigned int X_SEGMENTS = 64, unsigned int Y_SEGMENTS = 64) {
    vertices.clear();
    indices.clear();
    
    for (unsigned int y = 0; y <= Y_SEGMENTS; ++y) {
        for (unsigned int x = 0; x <= X_SEGMENTS; ++x) {
            float xSegment = static_cast<float>(x) / X_SEGMENTS;
            float ySegment = static_cast<float>(y) / Y_SEGMENTS;
            // Using ImGui's IM_PI for consistency.
            float xPos = std::cos(xSegment * 2.0f * IM_PI) * std::sin(ySegment * IM_PI);
            float yPos = std::cos(ySegment * IM_PI);
            float zPos = std::sin(xSegment * 2.0f * IM_PI) * std::sin(ySegment * IM_PI);
            // Position.
            vertices.push_back(xPos);
            vertices.push_back(yPos);
            vertices.push_back(zPos);
            // Normal (unit sphere: same as position).
            vertices.push_back(xPos);
            vertices.push_back(yPos);
            vertices.push_back(zPos);
            // Texture coordinates.
            vertices.push_back(xSegment);
            vertices.push_back(ySegment);
            // Tangent: derivative with respect to u.
            float tangentX = -std::sin(xSegment * 2.0f * IM_PI);
            float tangentY = 0.0f;
            float tangentZ =  std::cos(xSegment * 2.0f * IM_PI);
            vertices.push_back(tangentX);
            vertices.push_back(tangentY);
            vertices.push_back(tangentZ);
        }
    }
    
    for (unsigned int y = 0; y < Y_SEGMENTS; ++y) {
        for (unsigned int x = 0; x < X_SEGMENTS; ++x) {
            unsigned int i0 = y * (X_SEGMENTS + 1) + x;
            unsigned int i1 = i0 + 1;
            unsigned int i2 = i0 + (X_SEGMENTS + 1);
            unsigned int i3 = i2 + 1;
            indices.push_back(i0);
            indices.push_back(i2);
            indices.push_back(i1);
            indices.push_back(i1);
            indices.push_back(i2);
            indices.push_back(i3);
        }
    }
}

// --------------------------------------------------------------------------------------
// Material and Sphere Instance structures.
// The Material now includes a normal map.
// --------------------------------------------------------------------------------------
struct Material {
    glm::vec3 albedo;
    float metallic;
    float roughness;
    float ao;
    
    // Texture IDs.
    GLuint albedoTex;
    GLuint metallicTex;
    GLuint roughnessTex;
    GLuint aoTex;
    GLuint normalTex;
    
    // Flags indicating if a texture map is used.
    bool useAlbedoMap;
    bool useMetallicMap;
    bool useRoughnessMap;
    bool useAOMap;
    bool useNormalMap;
    
    // File path buffers for manual texture loading via ImGui.
    char albedoPath[256];
    char metallicPath[256];
    char roughnessPath[256];
    char aoPath[256];
    char normalPath[256];
    
    Material() : albedo(0.5f, 0.0f, 0.0f), metallic(0.0f), roughness(0.5f), ao(1.0f),
                 albedoTex(0), metallicTex(0), roughnessTex(0), aoTex(0), normalTex(0),
                 useAlbedoMap(false), useMetallicMap(false), useRoughnessMap(false),
                 useAOMap(false), useNormalMap(false)
    {
        strcpy(albedoPath, "");
        strcpy(metallicPath, "");
        strcpy(roughnessPath, "");
        strcpy(aoPath, "");
        strcpy(normalPath, "");
    }
};

struct SphereInstance {
    glm::vec3 position;
    float rotation; // Rotation angle (in radians).
    Material material;
};

// --------------------------------------------------------------------------------------
// YAML Material Loading
// --------------------------------------------------------------------------------------
// Expects a YAML file with keys: albedo (list of 3), metallic, roughness, ao,
// plus optional texture paths: albedo_texture, metallic_texture, roughness_texture, ao_texture, normal_texture.
Material loadMaterialFromYAML(const std::string& filename) {
    Material mat;
    try {
        YAML::Node config = YAML::LoadFile(filename);
        if (config["albedo"]) {
            mat.albedo = glm::vec3(config["albedo"][0].as<float>(),
                                   config["albedo"][1].as<float>(),
                                   config["albedo"][2].as<float>());
        }
        if (config["metallic"])
            mat.metallic = config["metallic"].as<float>();
        if (config["roughness"])
            mat.roughness = config["roughness"].as<float>();
        if (config["ao"])
            mat.ao = config["ao"].as<float>();
        
        if (config["albedo_texture"]) {
            std::string path = config["albedo_texture"].as<std::string>();
            strncpy(mat.albedoPath, path.c_str(), sizeof(mat.albedoPath));
            GLuint tex = LoadTexture(mat.albedoPath);
            if (tex != 0) {
                mat.albedoTex = tex;
                mat.useAlbedoMap = true;
            }
        }
        if (config["metallic_texture"]) {
            std::string path = config["metallic_texture"].as<std::string>();
            strncpy(mat.metallicPath, path.c_str(), sizeof(mat.metallicPath));
            GLuint tex = LoadTexture(mat.metallicPath);
            if (tex != 0) {
                mat.metallicTex = tex;
                mat.useMetallicMap = true;
            }
        }
        if (config["roughness_texture"]) {
            std::string path = config["roughness_texture"].as<std::string>();
            strncpy(mat.roughnessPath, path.c_str(), sizeof(mat.roughnessPath));
            GLuint tex = LoadTexture(mat.roughnessPath);
            if (tex != 0) {
                mat.roughnessTex = tex;
                mat.useRoughnessMap = true;
            }
        }
        if (config["ao_texture"]) {
            std::string path = config["ao_texture"].as<std::string>();
            strncpy(mat.aoPath, path.c_str(), sizeof(mat.aoPath));
            GLuint tex = LoadTexture(mat.aoPath);
            if (tex != 0) {
                mat.aoTex = tex;
                mat.useAOMap = true;
            }
        }
        if (config["normal_texture"]) {
            std::string path = config["normal_texture"].as<std::string>();
            strncpy(mat.normalPath, path.c_str(), sizeof(mat.normalPath));
            GLuint tex = LoadTexture(mat.normalPath);
            if (tex != 0) {
                mat.normalTex = tex;
                mat.useNormalMap = true;
            }
        }
    } catch(const std::exception& e) {
        std::cerr << "Error loading YAML material from file " << filename << ": " << e.what() << std::endl;
    }
    return mat;
}

// --------------------------------------------------------------------------------------
// YAML Material Saving
// --------------------------------------------------------------------------------------
// Saves a material to a YAML file with keys similar to the loader.
bool saveMaterialToYAML(const std::string& filename, const Material& mat) {
    YAML::Emitter out;
    out << YAML::BeginMap;
    out << YAML::Key << "albedo" << YAML::Value << YAML::Flow << std::vector<float>{mat.albedo.r, mat.albedo.g, mat.albedo.b};
    out << YAML::Key << "metallic" << YAML::Value << mat.metallic;
    out << YAML::Key << "roughness" << YAML::Value << mat.roughness;
    out << YAML::Key << "ao" << YAML::Value << mat.ao;
    if (mat.useAlbedoMap)
        out << YAML::Key << "albedo_texture" << YAML::Value << std::string(mat.albedoPath);
    if (mat.useMetallicMap)
        out << YAML::Key << "metallic_texture" << YAML::Value << std::string(mat.metallicPath);
    if (mat.useRoughnessMap)
        out << YAML::Key << "roughness_texture" << YAML::Value << std::string(mat.roughnessPath);
    if (mat.useAOMap)
        out << YAML::Key << "ao_texture" << YAML::Value << std::string(mat.aoPath);
    if (mat.useNormalMap)
        out << YAML::Key << "normal_texture" << YAML::Value << std::string(mat.normalPath);
    out << YAML::EndMap;
    
    std::ofstream fout(filename);
    if (!fout.is_open()){
        std::cerr << "Failed to open file for saving material: " << filename << std::endl;
        return false;
    }
    fout << out.c_str();
    fout.close();
    return true;
}

// --------------------------------------------------------------------------------------
// Main
// --------------------------------------------------------------------------------------
int main(){
    // Initialize GLFW.
    if (!glfwInit()){
        std::cerr << "Failed to initialize GLFW!" << std::endl;
        return -1;
    }
    glfwWindowHint(GLFW_CONTEXT_VERSION_MAJOR, 3);
    glfwWindowHint(GLFW_CONTEXT_VERSION_MINOR, 3);
    glfwWindowHint(GLFW_OPENGL_PROFILE, GLFW_OPENGL_CORE_PROFILE);
    
    GLFWwindow* window = glfwCreateWindow(800, 600, "PBR Renderer with ImGui, Normal Mapping & YAML Materials", nullptr, nullptr);
    if (!window){
        std::cerr << "Failed to create GLFW window!" << std::endl;
        glfwTerminate();
        return -1;
    }
    glfwMakeContextCurrent(window);
    
    // Initialize GLEW.
    glewExperimental = GL_TRUE;
    if (glewInit() != GLEW_OK){
        std::cerr << "Failed to initialize GLEW!" << std::endl;
        return -1;
    }
    glEnable(GL_DEPTH_TEST);
    
    // ----------------------------------------------------------------------------------
    // Setup ImGui context.
    // ----------------------------------------------------------------------------------
    IMGUI_CHECKVERSION();
    ImGui::CreateContext();
    ImGuiIO& io = ImGui::GetIO(); (void)io;
    ImGui::StyleColorsDark();
    ImGui_ImplGlfw_InitForOpenGL(window, true);
    ImGui_ImplOpenGL3_Init("#version 330");
    
    // ----------------------------------------------------------------------------------
    // Compile shaders and create shader program.
    // ----------------------------------------------------------------------------------
    GLuint vertexShader = compileShader(GL_VERTEX_SHADER, vertexShaderSource);
    GLuint fragmentShader = compileShader(GL_FRAGMENT_SHADER, fragmentShaderSource);
    GLuint shaderProgram = createProgram(vertexShader, fragmentShader);
    glDeleteShader(vertexShader);
    glDeleteShader(fragmentShader);
    
    // ----------------------------------------------------------------------------------
    // Generate the sphere mesh (with position, normal, texcoords, tangent).
    // ----------------------------------------------------------------------------------
    std::vector<float> sphereVertices;
    std::vector<unsigned int> sphereIndices;
    generateSphere(sphereVertices, sphereIndices, 64, 64);
    unsigned int indexCount = sphereIndices.size();
    
    GLuint VAO, VBO, EBO;
    glGenVertexArrays(1, &VAO);
    glGenBuffers(1, &VBO);
    glGenBuffers(1, &EBO);
    
    glBindVertexArray(VAO);
    glBindBuffer(GL_ARRAY_BUFFER, VBO);
    glBufferData(GL_ARRAY_BUFFER, sphereVertices.size() * sizeof(float), &sphereVertices[0], GL_STATIC_DRAW);
    
    glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, EBO);
    glBufferData(GL_ELEMENT_ARRAY_BUFFER, sphereIndices.size() * sizeof(unsigned int), &sphereIndices[0], GL_STATIC_DRAW);
    
    // Vertex attributes: 0 = position (3), 1 = normal (3), 2 = texcoords (2), 3 = tangent (3).
    glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, 11 * sizeof(float), (void*)0);
    glEnableVertexAttribArray(0);
    
    glVertexAttribPointer(1, 3, GL_FLOAT, GL_FALSE, 11 * sizeof(float), (void*)(3 * sizeof(float)));
    glEnableVertexAttribArray(1);
    
    glVertexAttribPointer(2, 2, GL_FLOAT, GL_FALSE, 11 * sizeof(float), (void*)(6 * sizeof(float)));
    glEnableVertexAttribArray(2);
    
    glVertexAttribPointer(3, 3, GL_FLOAT, GL_FALSE, 11 * sizeof(float), (void*)(8 * sizeof(float)));
    glEnableVertexAttribArray(3);
    glBindVertexArray(0);
    
    // ----------------------------------------------------------------------------------
    // Prepare scene data: a list of spheres.
    // ----------------------------------------------------------------------------------
    std::vector<SphereInstance> spheres;
    // Add one default sphere.
    spheres.push_back({ glm::vec3(0.0f), 0.0f, Material() });
    
    // Global scene/light parameters.
    glm::vec3 camPos(0.0f, 0.0f, 5.0f);
    glm::vec3 lightPos(0.0f, 0.0f, 10.0f);
    glm::vec3 lightColor(300.0f, 300.0f, 300.0f);
    
    // For YAML material loading via ImGui.
    static char yamlPathBuffer[256] = "";
    
    // ----------------------------------------------------------------------------------
    // Main render loop.
    // ----------------------------------------------------------------------------------
    while (!glfwWindowShouldClose(window)) {
        glfwPollEvents();
        
        // Start new ImGui frame.
        ImGui_ImplOpenGL3_NewFrame();
        ImGui_ImplGlfw_NewFrame();
        ImGui::NewFrame();
        
        // ----------------------------------------------------------------------------------
        // ImGui: Global Scene Controls.
        // ----------------------------------------------------------------------------------
        {
            ImGui::Begin("Scene Controls");
            
            // Global light controls.
            ImGui::Text("Light Controls");
            ImGui::DragFloat3("Light Position", glm::value_ptr(lightPos), 0.1f);
            ImGui::ColorEdit3("Light Color", glm::value_ptr(lightColor));
            ImGui::Separator();
            
            // Button to add a new sphere.
            if (ImGui::Button("Add Sphere")) {
                spheres.push_back({ glm::vec3(0.0f), 0.0f, Material() });
            }
            ImGui::Separator();
            
            // Global YAML material loader: set file path and add a sphere from YAML.
            ImGui::InputText("YAML Material Path", yamlPathBuffer, sizeof(yamlPathBuffer));
            if (ImGui::Button("Add Sphere from YAML Material")) {
                Material matFromYAML = loadMaterialFromYAML(yamlPathBuffer);
                spheres.push_back({ glm::vec3(0.0f), 0.0f, matFromYAML });
            }
            ImGui::Separator();
            
            // Per-sphere controls.
            for (size_t i = 0; i < spheres.size(); i++) {
                std::string header = "Sphere " + std::to_string(i);
                if (ImGui::CollapsingHeader(header.c_str())) {
                    // Transform controls.
                    ImGui::DragFloat3(("Position##" + std::to_string(i)).c_str(), glm::value_ptr(spheres[i].position), 0.1f);
                    ImGui::DragFloat(("Rotation (radians)##" + std::to_string(i)).c_str(), &spheres[i].rotation, 0.01f);
                    
                    // Material properties.
                    ImGui::ColorEdit3(("Albedo##" + std::to_string(i)).c_str(), glm::value_ptr(spheres[i].material.albedo));
                    ImGui::SliderFloat(("Metallic##" + std::to_string(i)).c_str(), &spheres[i].material.metallic, 0.0f, 1.0f);
                    ImGui::SliderFloat(("Roughness##" + std::to_string(i)).c_str(), &spheres[i].material.roughness, 0.05f, 1.0f);
                    ImGui::SliderFloat(("AO##" + std::to_string(i)).c_str(), &spheres[i].material.ao, 0.0f, 1.0f);
                    
                    // Texture map controls.
                    ImGui::InputText(("Albedo Texture Path##" + std::to_string(i)).c_str(),
                                     spheres[i].material.albedoPath, sizeof(spheres[i].material.albedoPath));
                    if (ImGui::Button(("Load Albedo Texture##" + std::to_string(i)).c_str())) {
                        GLuint tex = LoadTexture(spheres[i].material.albedoPath);
                        if (tex != 0) {
                            spheres[i].material.albedoTex = tex;
                            spheres[i].material.useAlbedoMap = true;
                        }
                    }
                    ImGui::InputText(("Metallic Texture Path##" + std::to_string(i)).c_str(),
                                     spheres[i].material.metallicPath, sizeof(spheres[i].material.metallicPath));
                    if (ImGui::Button(("Load Metallic Texture##" + std::to_string(i)).c_str())) {
                        GLuint tex = LoadTexture(spheres[i].material.metallicPath);
                        if (tex != 0) {
                            spheres[i].material.metallicTex = tex;
                            spheres[i].material.useMetallicMap = true;
                        }
                    }
                    ImGui::InputText(("Roughness Texture Path##" + std::to_string(i)).c_str(),
                                     spheres[i].material.roughnessPath, sizeof(spheres[i].material.roughnessPath));
                    if (ImGui::Button(("Load Roughness Texture##" + std::to_string(i)).c_str())) {
                        GLuint tex = LoadTexture(spheres[i].material.roughnessPath);
                        if (tex != 0) {
                            spheres[i].material.roughnessTex = tex;
                            spheres[i].material.useRoughnessMap = true;
                        }
                    }
                    ImGui::InputText(("AO Texture Path##" + std::to_string(i)).c_str(),
                                     spheres[i].material.aoPath, sizeof(spheres[i].material.aoPath));
                    if (ImGui::Button(("Load AO Texture##" + std::to_string(i)).c_str())) {
                        GLuint tex = LoadTexture(spheres[i].material.aoPath);
                        if (tex != 0) {
                            spheres[i].material.aoTex = tex;
                            spheres[i].material.useAOMap = true;
                        }
                    }
                    ImGui::InputText(("Normal Texture Path##" + std::to_string(i)).c_str(),
                                     spheres[i].material.normalPath, sizeof(spheres[i].material.normalPath));
                    if (ImGui::Button(("Load Normal Texture##" + std::to_string(i)).c_str())) {
                        GLuint tex = LoadTexture(spheres[i].material.normalPath);
                        if (tex != 0) {
                            spheres[i].material.normalTex = tex;
                            spheres[i].material.useNormalMap = true;
                        }
                    }
                    
                    // YAML material reloading.
                    static char sphereYAMLPath[256] = "";
                    ImGui::InputText(("YAML Material Path##" + std::to_string(i)).c_str(), sphereYAMLPath, sizeof(sphereYAMLPath));
                    if (ImGui::Button(("Load Material from YAML##" + std::to_string(i)).c_str())) {
                        Material newMat = loadMaterialFromYAML(sphereYAMLPath);
                        spheres[i].material = newMat;
                    }
                    
                    // ---------------------------
                    // Material Saving UI.
                    // ---------------------------
                    // We use a simple static buffer per sphere to specify a save file path.
                    static char saveYAMLPath[256] = "";
                    ImGui::InputText(("Save YAML Path##" + std::to_string(i)).c_str(), saveYAMLPath, sizeof(saveYAMLPath));
                    if (ImGui::Button(("Save Material to YAML##" + std::to_string(i)).c_str())) {
                        if (saveMaterialToYAML(std::string(saveYAMLPath), spheres[i].material)) {
                            std::cout << "Saved material for sphere " << i << " to " << saveYAMLPath << std::endl;
                        }
                    }
                    
                    if (ImGui::Button(("Remove Sphere##" + std::to_string(i)).c_str())) {
                        spheres.erase(spheres.begin() + i);
                        break;
                    }
                }
            }
            ImGui::End();
        }
        
        // ----------------------------------------------------------------------------------
        // Rendering.
        // ----------------------------------------------------------------------------------
        int width, height;
        glfwGetFramebufferSize(window, &width, &height);
        float aspect = width / static_cast<float>(height);
        glViewport(0, 0, width, height);
        glClearColor(0.1f, 0.1f, 0.1f, 1.0f);
        glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);
        
        // Setup camera matrices.
        glm::mat4 projection = glm::perspective(glm::radians(45.0f), aspect, 0.1f, 100.0f);
        glm::mat4 view = glm::lookAt(camPos, glm::vec3(0.0f), glm::vec3(0.0, 1.0, 0.0));
        
        glUseProgram(shaderProgram);
        glUniformMatrix4fv(glGetUniformLocation(shaderProgram, "projection"), 1, GL_FALSE, glm::value_ptr(projection));
        glUniformMatrix4fv(glGetUniformLocation(shaderProgram, "view"), 1, GL_FALSE, glm::value_ptr(view));
        glUniform3fv(glGetUniformLocation(shaderProgram, "camPos"), 1, glm::value_ptr(camPos));
        // Set light uniforms.
        glUniform3fv(glGetUniformLocation(shaderProgram, "lightPos"), 1, glm::value_ptr(lightPos));
        glUniform3fv(glGetUniformLocation(shaderProgram, "lightColor"), 1, glm::value_ptr(lightColor));
        // Pass ImGui's internal PI constant.
        glUniform1f(glGetUniformLocation(shaderProgram, "PI"), IM_PI);
        
        glBindVertexArray(VAO);
        for (const auto& sphere : spheres) {
            // Compute model matrix.
            glm::mat4 model = glm::translate(glm::mat4(1.0f), sphere.position);
            model = glm::rotate(model, sphere.rotation, glm::vec3(0.0f, 1.0f, 0.0f));
            glUniformMatrix4fv(glGetUniformLocation(shaderProgram, "model"), 1, GL_FALSE, glm::value_ptr(model));
            
            // Set fallback material uniforms.
            glUniform3fv(glGetUniformLocation(shaderProgram, "albedo"), 1, glm::value_ptr(sphere.material.albedo));
            glUniform1f(glGetUniformLocation(shaderProgram, "metallic"), sphere.material.metallic);
            glUniform1f(glGetUniformLocation(shaderProgram, "roughness"), sphere.material.roughness);
            glUniform1f(glGetUniformLocation(shaderProgram, "ao"), sphere.material.ao);
            
            // Set texture usage flags.
            glUniform1i(glGetUniformLocation(shaderProgram, "useAlbedoMap"), sphere.material.useAlbedoMap ? 1 : 0);
            glUniform1i(glGetUniformLocation(shaderProgram, "useMetallicMap"), sphere.material.useMetallicMap ? 1 : 0);
            glUniform1i(glGetUniformLocation(shaderProgram, "useRoughnessMap"), sphere.material.useRoughnessMap ? 1 : 0);
            glUniform1i(glGetUniformLocation(shaderProgram, "useAOMap"), sphere.material.useAOMap ? 1 : 0);
            glUniform1i(glGetUniformLocation(shaderProgram, "useNormalMap"), sphere.material.useNormalMap ? 1 : 0);
            
            // Bind textures if enabled.
            if (sphere.material.useAlbedoMap) {
                glActiveTexture(GL_TEXTURE0);
                glBindTexture(GL_TEXTURE_2D, sphere.material.albedoTex);
                glUniform1i(glGetUniformLocation(shaderProgram, "albedoMap"), 0);
            }
            if (sphere.material.useMetallicMap) {
                glActiveTexture(GL_TEXTURE1);
                glBindTexture(GL_TEXTURE_2D, sphere.material.metallicTex);
                glUniform1i(glGetUniformLocation(shaderProgram, "metallicMap"), 1);
            }
            if (sphere.material.useRoughnessMap) {
                glActiveTexture(GL_TEXTURE2);
                glBindTexture(GL_TEXTURE_2D, sphere.material.roughnessTex);
                glUniform1i(glGetUniformLocation(shaderProgram, "roughnessMap"), 2);
            }
            if (sphere.material.useAOMap) {
                glActiveTexture(GL_TEXTURE3);
                glBindTexture(GL_TEXTURE_2D, sphere.material.aoTex);
                glUniform1i(glGetUniformLocation(shaderProgram, "aoMap"), 3);
            }
            if (sphere.material.useNormalMap) {
                glActiveTexture(GL_TEXTURE4);
                glBindTexture(GL_TEXTURE_2D, sphere.material.normalTex);
                glUniform1i(glGetUniformLocation(shaderProgram, "normalMap"), 4);
            }
            
            glDrawElements(GL_TRIANGLES, indexCount, GL_UNSIGNED_INT, 0);
        }
        glBindVertexArray(0);
        
        // Render ImGui over the scene.
        ImGui::Render();
        ImGui_ImplOpenGL3_RenderDrawData(ImGui::GetDrawData());
        
        glfwSwapBuffers(window);
    }
    
    // ----------------------------------------------------------------------------------
    // Cleanup.
    // ----------------------------------------------------------------------------------
    glDeleteVertexArrays(1, &VAO);
    glDeleteBuffers(1, &VBO);
    glDeleteBuffers(1, &EBO);
    glDeleteProgram(shaderProgram);
    
    ImGui_ImplOpenGL3_Shutdown();
    ImGui_ImplGlfw_Shutdown();
    ImGui::DestroyContext();
    
    glfwTerminate();
    return 0;
}