#include "scene.h"
#include "tiny_obj_loader.h"
#include "sh.h"
#include "shproject.h"
#include <OpenImageIO/imageio.h>

#include <algorithm>
#include <iterator>

using namespace FireRays;

static int GetTextureForemat(OIIO_NAMESPACE::ImageSpec const& spec)
{
    OIIO_NAMESPACE_USING

    if (spec.format.basetype == TypeDesc::UINT8)
        return Scene::RGBA8;
    if (spec.format.basetype == TypeDesc::HALF)
        return Scene::RGBA16;
    else
        return Scene::RGBA32;
}


static void LoadTexture(std::string const& filename, Scene::Texture& texture, std::vector<std::unique_ptr<char[]> >& data)
{
    OIIO_NAMESPACE_USING

        ImageInput* input = ImageInput::open(filename);

    if (!input)
    {
        throw std::runtime_error("Can't load " + filename + " image");
    }

    ImageSpec const& spec = input->spec();

    texture.w = spec.width;
    texture.h = spec.height;
    texture.d = spec.depth;
    texture.fmt = GetTextureForemat(spec);

    // Save old size for reading offset
    texture.dataoffset = (int)data.size();

    if (texture.fmt == Scene::RGBA8)
    {
        texture.size = spec.width * spec.height * spec.depth * 4;

        // Resize storage
        std::unique_ptr<char[]> texturedata(new char[spec.width * spec.height * spec.depth * 4]);

        // Read data to storage
        input->read_image(TypeDesc::UINT8, texturedata.get(), sizeof(char)* 4);

        // Close handle
        input->close();

        // Add to texture pool
        data.push_back(std::move(texturedata));
    }
    else if (texture.fmt == Scene::RGBA16)
    {
        texture.size = spec.width * spec.height * spec.depth * sizeof(float)* 2;

        // Resize storage
        std::unique_ptr<char[]> texturedata(new char[spec.width * spec.height * spec.depth * sizeof(float)* 2]);

        // Read data to storage
        input->read_image(TypeDesc::HALF, texturedata.get(), sizeof(float)* 2);

        // Close handle
        input->close();

        // Add to texture pool
        data.push_back(std::move(texturedata));
    }
    else
    {
        texture.size = spec.width * spec.height * spec.depth * sizeof(float3);

        // Resize storage
        std::unique_ptr<char[]> texturedata(new char[spec.width * spec.height * spec.depth * sizeof(float3)]);

        // Read data to storage
        input->read_image(TypeDesc::FLOAT, texturedata.get(), sizeof(float3));

        // Close handle
        input->close();

        // Add to texture pool
        data.push_back(std::move(texturedata));
    }

    // Cleanup
    delete input;
}

Scene* Scene::LoadFromObj(std::string const& filename, std::string const& basepath)
{
    using namespace tinyobj;

    // Loader data
    std::vector<shape_t> objshapes;
    std::vector<material_t> objmaterials;

    // Try loading file
    std::string res = LoadObj(objshapes, objmaterials, filename.c_str(), basepath.c_str());

    // Allocate scene
    Scene* scene(new Scene);

    // Texture map
    std::map<std::string, int> textures;

    // Enumerate and translate materials
    for (int i = 0; i < (int)objmaterials.size(); ++i)
    {
        Material material;

        material.kd = float3(objmaterials[i].diffuse[0], objmaterials[i].diffuse[1], objmaterials[i].diffuse[2]);
        float3 ks = float3(objmaterials[i].specular[0], objmaterials[i].specular[1], objmaterials[i].specular[2]);
        material.ke = float3(objmaterials[i].emission[0], objmaterials[i].emission[1], objmaterials[i].emission[2]);
        material.kd.w = objmaterials[i].ior;
        material.roughness = 1.f - clamp(objmaterials[i].shininess / 64.f, 0.f, 1.f);
        material.metallic = material.kd.sqnorm() == 0.f ? 1.f : (ks.sqnorm() / material.kd.sqnorm());

        /// TODO: dont forget to remove
        //if (i == 1)
        /*{
            //CARPAINT:
            material.kd = float3(0.66f, 0.66f, 0.66f);
            material.kd.w = 6.f;
            material.metallic = 0.1f;
            material.specular = 0.9f;

            material.roughness = 0.1f;
            material.clearcoat = 0.99f;
            material.clearcoatgloss = 0.99f;
            }*/

        // Load diffuse texture if needed
        if (!objmaterials[i].diffuse_texname.empty())
        {
            auto iter = textures.find(objmaterials[i].diffuse_texname);
            if (iter != textures.end())
            {
                material.kdmapidx = iter->second;
            }
            else
            {
                Texture texture;

                // Load texture
                LoadTexture(basepath + "/" + objmaterials[i].diffuse_texname, texture, scene->texturedata_);

                // Add texture desc
                material.kdmapidx = (int)scene->textures_.size();
                scene->textures_.push_back(texture);

                // Save in the map
                textures[objmaterials[i].diffuse_texname] = material.kdmapidx;
            }
        }

        // Load specular texture
        /*if (!objmaterials[i].specular_texname.empty())
        {
        auto iter = textures.find(objmaterials[i].specular_texname);
        if (iter != textures.end())
        {
        material.ksmapidx = iter->second;
        }
        else
        {
        Texture texture;

        // Load texture
        LoadTexture(basepath + "/" + objmaterials[i].specular_texname, texture, scene->texturedata_);

        // Add texture desc
        material.ksmapidx = (int)scene->textures_.size();
        scene->textures_.push_back(texture);

        // Save in the map
        textures[objmaterials[i].specular_texname] = material.ksmapidx;
        }
        }*/

        // Load normal map
        if (!objmaterials[i].normal_texname.empty())
        {
            auto iter = textures.find(objmaterials[i].normal_texname);
            if (iter != textures.end())
            {
                material.nmapidx = iter->second;
            }
            else
            {
                Texture texture;

                // Load texture
                LoadTexture(basepath + "/" + objmaterials[i].normal_texname, texture, scene->texturedata_);

                // Add texture desc
                material.nmapidx = (int)scene->textures_.size();
                scene->textures_.push_back(texture);

                // Save in the map
                textures[objmaterials[i].normal_texname] = material.nmapidx;
            }
        }


        scene->materials_.push_back(material);
    }

    // Enumerate all shapes in the scene
    for (int s = 0; s < (int)objshapes.size(); ++s)
    {
        // Prepate shape
        Shape shape;
        shape.startidx = (int)scene->indices_.size();
        shape.numprims = (int)objshapes[s].mesh.indices.size() / 3;
        shape.startvtx = (int)scene->vertices_.size();
        shape.numvertices = (int)objshapes[s].mesh.positions.size() / 3;
        shape.m = matrix();
        shape.linearvelocity = float3(0.0f, 0.f, 0.f);
        shape.angularvelocity = quaternion(0.f, 0.f, 0.f, 1.f);
        // Save last index to add to this shape indices
        // int baseidx = (int)scene->vertices_.size();

        // Enumerate and copy vertex data
        for (int i = 0; i < (int)objshapes[s].mesh.positions.size() / 3; ++i)
        {
            scene->vertices_.push_back(float3(objshapes[s].mesh.positions[3 * i], objshapes[s].mesh.positions[3 * i + 1], objshapes[s].mesh.positions[3 * i + 2]));
        }

        for (int i = 0; i < (int)objshapes[s].mesh.normals.size() / 3; ++i)
        {
            scene->normals_.push_back(float3(objshapes[s].mesh.normals[3 * i], objshapes[s].mesh.normals[3 * i + 1], objshapes[s].mesh.normals[3 * i + 2]));
        }

        for (int i = 0; i < (int)objshapes[s].mesh.texcoords.size() / 2; ++i)
        {
            scene->uvs_.push_back(float2(objshapes[s].mesh.texcoords[2 * i], objshapes[s].mesh.texcoords[2 * i + 1]));
        }


        // Enumerate and copy indices (accounting for base index) and material indices
        for (int i = 0; i<(int)objshapes[s].mesh.indices.size() / 3; ++i)
        {
            scene->indices_.push_back(objshapes[s].mesh.indices[3 * i]);
            scene->indices_.push_back(objshapes[s].mesh.indices[3 * i + 1]);
            scene->indices_.push_back(objshapes[s].mesh.indices[3 * i + 2]);

            int matidx = objshapes[s].mesh.material_ids[i];
            scene->materialids_.push_back(matidx);

            if (scene->materials_[matidx].ke.sqnorm() > 0)
            {
                Emissive emissive = { s, i };
                scene->emissives_.push_back(emissive);
            }
        }

        scene->shapes_.push_back(shape);
    }

    // Check if there is no UVs
    if (scene->uvs_.empty())
    {
        scene->uvs_.resize(scene->normals_.size());
        std::fill(scene->uvs_.begin(), scene->uvs_.end(), float2(0, 0));
    }

    std::cout << "Loading complete\n";
    std::cout << "Number of objects: " << scene->shapes_.size() << "\n";
    std::cout << "Number of textures: " << scene->textures_.size() << "\n";

    return scene;
}

void Scene::SetEnvironment(std::string const& filename, std::string const& basepath, float envmapmul)
{
    // Save multiplier
    envmapmul_ = envmapmul;

    Texture texture;

    // Load texture
    if (!basepath.empty())
    {
        LoadTexture(basepath + "/" + filename, texture, texturedata_);
    }
    else
    {
        LoadTexture(filename, texture, texturedata_);
    }

    // Save index
    envidx_ = (int)textures_.size();

    // Add texture desc
    textures_.push_back(texture);
}

void Scene::SetBackground(std::string const& filename, std::string const& basepath)
{
    Texture texture;

    // Load texture
    if (!basepath.empty())
    {
        LoadTexture(basepath + "/" + filename, texture, texturedata_);
    }
    else
    {
        LoadTexture(filename, texture, texturedata_);
    }

    // Save index
    bgimgidx_ = (int)textures_.size();

    // Add texture desc
    textures_.push_back(texture);
}
