// ================================================================================================
// 
// This file is part of the CAMPVis Software Framework.
// 
// If not explicitly stated otherwise: Copyright (C) 2012-2015, all rights reserved,
//      Christian Schulte zu Berge <christian.szb@in.tum.de>
//      Chair for Computer Aided Medical Procedures
//      Technische Universitaet Muenchen
//      Boltzmannstr. 3, 85748 Garching b. Muenchen, Germany
// 
// For a full list of authors and contributors, please refer to the file "AUTHORS.txt".
// 
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file 
// except in compliance with the License. You may obtain a copy of the License at
// 
// http://www.apache.org/licenses/LICENSE-2.0
// 
// Unless required by applicable law or agreed to in writing, software distributed under the 
// License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, 
// either express or implied. See the License for the specific language governing permissions 
// and limitations under the License.
// 
// ================================================================================================

#include "tfpreintegrator.h"

#include "cgt/buffer.h"
#include "cgt/imageunit.h"
#include "cgt/logmanager.h"
#include "cgt/shadermanager.h"
#include "cgt/textureunit.h"
#include "cgt/texture.h"

#include "core/datastructures/imagedata.h"
#include "core/datastructures/imagerepresentationgl.h"
#include "core/datastructures/renderdata.h"

#include "core/tools/cshelper.h"
#include "core/tools/quadrenderer.h"
#include "core/tools/stringutils.h"

#include "core/classification/geometry1dtransferfunction.h"

namespace campvis {

    const std::string TFPreIntegrator::loggerCat_ = "CAMPVis.modules.advancedraycasting";


    TFPreIntegrator::TFPreIntegrator()
        : AbstractProcessor()
        , p_inputImage("InputImage", "Input Image", "", DataNameProperty::READ)
        , p_outputImage("OutputImage", "Output Image", "PreIntegratedVolume", DataNameProperty::WRITE)
        , p_transferFunction("TransferFunction", "Transfer Function", new Geometry1DTransferFunction(256, cgt::vec2(0, 1)))
        , p_workgroupSize("WorkgroupSize", "Workgroup Size", cgt::ivec3(4, 4, 2), cgt::ivec3(1), cgt::ivec3(16))
        , _shader(nullptr)
    {
        addProperty(p_inputImage, INVALID_RESULT | INVALID_PROPERTIES);
        addProperty(p_outputImage);
        addProperty(p_workgroupSize);
        addProperty(p_transferFunction);
    }

    TFPreIntegrator::~TFPreIntegrator() {

    }

    void TFPreIntegrator::init() {
        AbstractProcessor::init();

        LGL_ERROR;
    }

    void TFPreIntegrator::deinit() {
        ShdrMgr.dispose(_shader);

        AbstractProcessor::deinit();
    }

    void TFPreIntegrator::updateResult(DataContainer& data) {
        ImageRepresentationGL::ScopedRepresentation img(data, p_inputImage.getValue());

        if (img != 0) {
            if (img->getParent()->getDimensionality() > 1) {

                cgt::ImageUnit outputUnit;
                generateShader(img, outputUnit);

                if (_shader) {
                    cgt::ivec3 size = img->getSize();
                    cgt::ivec3 wgSize = p_workgroupSize.getValue();
                    cgt::ivec3 groups(cgt::ceil(cgt::vec3(size) / cgt::vec3(wgSize)));

                    cgt::TextureUnit inputUnit;
                    inputUnit.activate();

                    cgt::TextureUnit tfUnit;

                    _shader->activate();

                    p_transferFunction.getTF()->bind(_shader, tfUnit);

                    // create texture for result
                    std::unique_ptr<cgt::Texture> resultTexture(new cgt::Texture(img->getTexture()->getType(), size, GL_RGBA8, cgt::Texture::LINEAR));
                    resultTexture->setWrapping(cgt::Texture::CLAMP_TO_EDGE);
                    img->bind(_shader, inputUnit);
                    LGL_ERROR;
                    resultTexture->bindImage(outputUnit, GL_WRITE_ONLY);
                    LGL_ERROR;
                    _shader->setUniform("_outputImage", outputUnit.getUnitNumber());
                    LGL_ERROR;

                    // make sure the textures have been uploaded completely
                    glMemoryBarrier(GL_TEXTURE_UPDATE_BARRIER_BIT);

                    glDispatchCompute(groups.x, groups.y, groups.z);

                    LGL_ERROR;


                    _shader->deactivate();

                    // Make sure the result is actually finished
                    glFinish();

                    // put resulting image into DataContainer
                    std::unique_ptr<ImageData> id(new ImageData(img->getParent()->getDimensionality(), size, 4));
                    ImageRepresentationGL::create(id.get(), resultTexture.release());
                    id->setMappingInformation(img->getParent()->getMappingInformation());
                    data.addData(p_outputImage.getValue(), id.release());

                    cgt::TextureUnit::setZeroUnit();
                    LGL_ERROR;
                }
            }
            else {
                LERROR("Supports only 2D and 3D Median Filtering.");
            }
        }
        else {
            LDEBUG("No suitable input image found.");
        }
    }

    void TFPreIntegrator::updateProperties(DataContainer& dc) {
        ScopedTypedData<ImageData> img(dc, p_inputImage.getValue());

        if(img)
            p_transferFunction.setImageHandle(img.getDataHandle());
    }

    void TFPreIntegrator::generateShader(const ImageRepresentationGL::ScopedRepresentation& img, const cgt::ImageUnit& imgUnit)
    {
        auto wg = p_workgroupSize.getValue();
        std::stringstream ss;
        ss << "#define WORK_GROUP_SIZE_X " << wg.x << std::endl;
        ss << "#define WORK_GROUP_SIZE_Y " << wg.y << std::endl;
        ss << "#define WORK_GROUP_SIZE_Z " << wg.z << std::endl;
        ss << "#define OUTPUT_TEXTURE_FORMAT " << cgt::Texture::calcMatchingWriteFormat(img->getTexture()->getInternalFormat()) << std::endl;
        ss << "#define TEXTURE_DIMENSIONALITY " << img->getDimensionality() << std::endl;
        ss << CSHelper::generateGLSLImageDefinition(*(img->getTexture()), "_outputImage", imgUnit) << std::endl;

        // very simple caching to eliminate the glsl compiler as a bottleneck
        static std::string headers;
        if (!_shader || headers != ss.str()) {
            headers = ss.str();
            //_shader->setHeaders(headers);
            //_shader->rebuild();
            _shader = ShdrMgr.loadCompute("modules/advancedraycasting/glsl/tfpreintegrator.comp", headers);
        }
    }
}