// ================================================================================================
// 
// 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.
// 
// ================================================================================================

// (c) 2017 Jakob Weiss <jakob.weiss@tum.de>

#include "tools/transferfunction.frag"
#include "tools/texture3d.frag"

// The number of rays in each direction
#ifndef NUM_AO_RAYS
#define NUM_AO_RAYS 12
#endif

// The number of steps in each direction
#ifndef NUM_AO_RAY_STEPS
#define NUM_AO_RAY_STEPS 12
#endif

// The gamma factor that controls the visual strength of the AO term
#ifndef AO_GAMMA
#define AO_GAMMA 1.0
#endif

// factor to control the strength of emissive materials.
#ifndef AO_EMISSIVE_SCALE
#define AO_EMISSIVE_SCALE 5.0
#endif

// factor to control the transparency used when computing the AO effect
#ifndef AO_OPACITY_SCALE
#define AO_OPACITY_SCALE 1.0
#endif

const float AO_SAMPLING_BASE_INTERVAL_RCP = 200.0;

vec3 nthSphereSample(float R, int n, int N) {
    float golden_angle = 3.1415926535 * (3. - sqrt(5.));
    float theta = golden_angle * float(n);
    float t = float(n) / float(N - 1);
    float z = (1. - t) * (1 - 1.0 / float(N)) + t * (1.0 / float(N) - 1); // TODO use mix() here
    float radius = sqrt(1 - z * z);

    return R * vec3(radius * cos(theta), radius * sin(theta), z);
}


vec4 computeLAORay(vec3 pos, vec3 dir, float aoStepSize, sampler3D tex, sampler1D volumeTF, TFParameters1D volumeTFParams, sampler1D aoEmissiveTF, TFParameters1D aoEmissiveTFParams)
{
    vec3 aoStep = dir;
    vec3 aoRayPos = pos + aoStep; // offset
    vec4 aoColor = vec4(vec3(0), 1.0);

    for(int i=0; i < NUM_AO_RAY_STEPS; ++i) {
        // sample AO ray
        float aoIntensity = texture(tex, aoRayPos).r;
        vec4 aoSampleColor = lookupTF(volumeTF, volumeTFParams, aoIntensity);
        float aoSampleOpacity = aoSampleColor.a * AO_OPACITY_SCALE; // Volume tf sample is used for it's opacity value
        vec4 aoSampleEmissive = lookupTF(aoEmissiveTF, aoEmissiveTFParams, aoIntensity); // AO sample is used for the material's emissive color

        // compensate step size
        aoSampleOpacity = 1.0 - pow(1.0 - aoSampleOpacity, AO_GAMMA * aoStepSize * AO_SAMPLING_BASE_INTERVAL_RCP);
        aoSampleEmissive.rgb *= aoStepSize * AO_SAMPLING_BASE_INTERVAL_RCP  * aoSampleEmissive.a; // ao emissive color is multiplied by alpha for controlling the intensity

        // accumulate
        aoColor.rgb += (vec3(1) + AO_EMISSIVE_SCALE * aoSampleEmissive.rgb) * aoColor.a;
        aoColor.a *= (1 - aoSampleOpacity);

        aoRayPos += aoStep;
    }

    return aoColor;
}

void initLAODirs(inout vec4 aoRayDirs[NUM_AO_RAYS], float aoSphereRadius, TextureParameters3D volumeParameters)
{
    //Normalize for anisotropic voxel sizes
    vec3 anisotropicVoxel = volumeParameters._sizeRCP * volumeParameters._voxelSizeRCP;
    vec3 anisotropicRayStep = anisotropicVoxel / NUM_AO_RAY_STEPS;
    float voxelSize = length(volumeParameters._sizeRCP) / sqrt(3);
    float scale = aoSphereRadius * voxelSize / NUM_AO_RAY_STEPS;

    for(int i=0; i < NUM_AO_RAYS; ++i) {
        vec4 dir = vec4( nthSphereSample(1.0, i, NUM_AO_RAYS), 0.0);
        aoRayDirs[i].xyz = scale * normalize((volumeParameters._worldToTextureMatrixInvTransp * dir).xyz);

        //this is somehow slower to compute...
        //aoRayDirs[i].xyz = anisotropicRayStep * nthSphereSample(aoSphereRadius, i, NUM_AO_RAYS);

        aoRayDirs[i].w = length(aoRayDirs[i].xyz);
    }
}

vec3 computeLAO(vec3 position, in vec4 aoRayDirs[NUM_AO_RAYS], sampler3D tex, sampler1D volumeTF, TFParameters1D volumeTFParams, sampler1D aoEmissiveTF, TFParameters1D aoEmissiveTFParams, inout vec3 minOcclusionDir)
{
    minOcclusionDir = vec3(0);

    if(NUM_AO_RAYS*NUM_AO_RAY_STEPS > 0) {
        vec3 ambient = vec3(0);
        for(int i=0; i < NUM_AO_RAYS; ++i) {
            //vec3 dir = nthSphereSample(1.0, i, NUM_AO_RAYS);
            vec4 dir = aoRayDirs[i];
            vec4 aoRayColor = computeLAORay(position, dir.xyz, dir.w, tex, volumeTF, volumeTFParams, aoEmissiveTF, aoEmissiveTFParams);
            ambient += aoRayColor.rgb;

            minOcclusionDir += (1.- aoRayColor.a) * dir.xyz;
        }

        minOcclusionDir = normalize(minOcclusionDir);
        // normalize against number of samples
        return ambient / (NUM_AO_RAYS * NUM_AO_RAY_STEPS);
   }
   else
        return vec3(1);
}

vec3 computeLAO(vec3 position, in vec4 aoRayDirs[NUM_AO_RAYS], sampler3D tex, sampler1D volumeTF, TFParameters1D volumeTFParams, sampler1D aoEmissiveTF, TFParameters1D aoEmissiveTFParams)
{
    vec3 minOcclusionDir;
    return computeLAO(position, aoRayDirs, tex, volumeTF, volumeTFParams, aoEmissiveTF, aoEmissiveTFParams, minOcclusionDir);
}