// ===================================================================
// $Id$
//
// simplified_cook_torrance.cpp
//    Implementation of simplified Cook-Torrance BRDF.
//
// Class: CSimpleCookTorrance
//
// Initial coding by Hana Truskova & Michael Smutny, 2009.
// ===================================================================

#include "config.h"
#include "simplified_cook_torrance.h"
#include "functions.h"
#include <iostream>

__BEGIN_BRDFLIB_SOURCE


/**********************************************************************
 *
 *             CSimpleCookTorrance methods
 *
 **********************************************************************/
/*
  Albedo is the integral of BRDF multiplied by cos(Eye,N) over the hemisphere.
  Integral of simplified Cook-Torrance BRDF can be aproximated.
 */
//========================================================================
bool
CSimpleCookTorrance::GetAlbedo(CValues &albedo,
                          const CHitPointInfo &hitInfo,
                          EDistributionType   distType) const
{
  CValues c(CValues::EE_Weigth3); // result - zero by default
  CValues d(CValues::EE_Weigth3);
  CVector3D N       = hitInfo.GetNormal();
  float incEyeCos   = -(hitInfo.GetDotProd());
  CVector3D dirEye  = (hitInfo.GetRay())->GetDir();
  CVector3D  H  =  Normalize(dirEye - (hitInfo.GetRay()->GetDir()));
  float cos_HV  =  DotProd(H, -dirEye);
  bool retval = false;
  CValues white(CValues::EE_Weigth3);
  white.Set(0, 1); white.Set(1, 1);  white.Set(2, 1);

  // Specular part of albedo
  CValues a_i;
  if (((distType & _microfacetReflectionType()) && _specular) && _slope) {
    hitInfo.GetTextureValue(d, *_specular, CMappingBase::EE_KS);
    CValues rho_f = CValues(d + pow(1-cos_HV,5)*(white-d));
    a_i = CValues(incEyeCos/cos_HV*rho_f);
    c += a_i;
    retval = true;
  }

  // Diffuse part of albedo
  if ((distType & EE_REFLECTED_DIFFUSE) && _diffuse) {
    hitInfo.GetTextureValue(d, *_diffuse, CMappingBase::EE_KD);
    c += d;//white - a_i;
    retval = true;
  }

//#define _DEBUG_PHYSICAL_PLAUSIBILITY
#ifdef _DEBUG_PHYSICAL_PLAUSIBILITY
  // test if the albedo is smaller than 1.0 -- physically plausible
  if (c.Max() > 1.0) {
    if (_diffuse) {
      hitInfo.GetTextureValue(d, *_diffuse, CMappingBase::EE_KD);
      DEBUG << "diffuse = " << d << "\n";
    }
    if (_specular) {
      hitInfo.GetTextureValue(d, *_specular, CMappingBase::EE_KS);
      DEBUG << "specular = " << d << "\n";
    }
    cout << "Current parameter settings for Cook-Torrance model make it\n";
    cout << "physically unplausible\n";
    exit(0);
  }
#endif // _DEBUG_PHYSICAL_PLAUSIBILITY

  if (retval) {
    albedo = c;
    return true;
  }

  // unable to compute the value for some reason
#ifdef _DEBUG_PHYSICAL_PLAUSIBILITY
  if ((!_diffuse) && (!_specular)) {
    WARNING << "This is the black body - albedo cannot be computed" << endl;
  }
#endif // _DEBUG

  return false;
}

// Direct lighting computation
//=============================================================================
bool
CSimpleCookTorrance::EvaluateBSDF(CValues &result,
			 const CHitPointInfo &hitInfo,
			 const CVector3D     &dirLit,
			 EDistributionType   distType) const
{
  bool retval 	    = false;
  float d           = Magnitude(dirLit);
  CVector3D N       = hitInfo.GetNormal();
  CVector3D dirEye  = (hitInfo.GetRay())->GetDir();
  float lightCos    = DotProd(dirLit, N);
  float incEyeCos   = -(hitInfo.GetDotProd());

  CVector3D  H  =  Normalize(dirLit - (hitInfo.GetRay()->GetDir()));
  float cos_HV  =  DotProd(H, -dirEye);
  float cos_HN  =  DotProd(H, N);
  CValues white(CValues::EE_Weigth3);
  white.Set(0, 1); white.Set(1, 1);  white.Set(2, 1);

  // Provide 'zero' result if Light source or viewer is under the surface
  if(lightCos >= -1e-06 && incEyeCos >= -1e-06)
  {
    if(lightCos < WTINY) lightCos       = WTINY;
    if(incEyeCos < WTINY) incEyeCos  = WTINY;

    if ( (d < 1.0-1e-06) || (d > 1.0+1e-06) )
      Normalize(dirLit);

    // Diffuse part
    //~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    if(_diffuse && (distType & EE_REFLECTED_DIFFUSE)) {
      CValues rho_d;
      hitInfo.GetTextureValue(rho_d, *_diffuse, CMappingBase::EE_KS);
      // in the initial cook_torance implementation there is following matte component
      result = rho_d * M_1_PI * lightCos;
      /*CValues rho_f = CValues(rho_d + pow(1-cos_HV,5)*(white-rho_d));
      CValues a_avg = CValues(lightCos*lightCos/cos_HV*rho_f);
      CValues a_i = CValues(lightCos/cos_HV*rho_f);
      CValues a_r = CValues(incEyeCos/cos_HV*rho_f);
      CValues temp = CValues(M_PI*(white-a_avg));
      CValues help = CValues(( white- a_i)*(white-a_r)*lightCos);
      result = CValues(help.Get(0)/temp.Get(0),help.Get(1)/temp.Get(1), help.Get(2)/temp.Get(2) );*/
      retval = true;
    }

   // Specular part
   //~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
   if((_specular && (distType & EE_REFLECTED_GLOSSY)) && _slope) {
      int slopes_size  = _slope->GetSlopesSize();
      int weights_size = _slope->GetWeightsSize();

      // Compute specular part only if slope structure is defined
      if(slopes_size > 0 && (slopes_size == weights_size))
      {
        float *a2      = new float[slopes_size];
        float *weights = new float[slopes_size];

        // Get values for slopes and their weights
        for(int i = 0; i < slopes_size; i++)
        {
                float weight = 1.0;
          CTexture *slope = _slope->GetSlope(i, weight);
          assert(slope);
                CValues s_tmp(CValues::EE_Weigth3);
          hitInfo.GetTextureValue(s_tmp, *slope, CMappingBase::EE_KSI);

          a2[i]      = (s_tmp.Get(0)) * (s_tmp.Get(0));
                if(a2[i] < WTINY) a2[i] = WTINY;

          weights[i] = weight;
        }

        CValues rho_s;
        bool use_fresnel = _specular->IsFresnel();

        // Specular texture for Cook-Torrance model must be defined
        // as Fresnel texture
        if(use_fresnel)
          hitInfo.GetTextureValue(rho_s, *_specular, CMappingBase::EE_KFM);
        else
		{
			WARNING_INFO("CSimpleCookTorrance::EvaluateBSDF() : Specular texture for Cook-Torrance model must be defined as Fresnel texture!");
			return false;
		}
          

        assert(cos_HN >= 0);

        if(cos_HN < WTINY)  cos_HN = WTINY;
        float cos_HN3 = pow(cos_HN, 3.0f);

        float tan2_HN =  1.0f/(cos_HN * cos_HN) - 1.0;
        if(tan2_HN < 0) tan2_HN = 0;

        // compute term Distribution
        //--------------------------------------------------------------
        float distribution = 0;
        for(int i = 0; i < slopes_size; i++)
        {
          float memb    = exp( -tan2_HN/a2[i] )/(M_PI * a2[i] * cos_HN3);
          distribution += weights[i] * memb;
        }

        float k = distribution * lightCos ;

        // Compute geometrical attenuation factor (accounting for
        // shadowing and masking)
        //--------------------------------------------------------------
        float G = 1/(4*cos_HV*cos_HV);//1/(1+DotProd(dirEye, dirLit));

        // Add all together
        // calculate Fresnel - (rho_s + pow(1-DotProd(-dirEye,N),5)*(white-rho_s))
        result += rho_s * k * G;
        retval = true;

        // delete temporary structures
        delete [] a2;
        delete [] weights;
      }
    }
    /* returned result = BRDF * lightCos
       But BRDF model contains lightCos in denominator => result
       should not be multiplied on it */
  }
  else
  {
    retval = false;
  }
  return retval;
}

/* Specular part for stochastic sampling which is implemented in a similar
   to importance sampling scheme for Ward model way. (see paper Kelemen'2001)

   Distribution function contains several members (in other words, surface
   consists of several types of microfacets). In this way, to provide the
   importance sampling only one type of microfacet should be choosen.

   For vizualisation not only the diffuse part should be added to
   the computed sample weight, but the `rest` of the specular part represented
   by unchoosen distribution members. */
//========================================================================
int
CSimpleCookTorrance::SpawnStochasticSpecularRays(
  const CValues       &ks_local,
  CVector2D           &slopeXY,
  CVector2D           &slopeXY_sq,
  const CValues       &kd_local,
  CContributionArray  &carray,
  const CHitPointInfo &hitInfo,
  CSamplingInfo       &samplingInfo,
  int                 maxRays) const
{
  int num = 0;		// the number of samples generated

  float incEyeCos  = hitInfo.GetDotProd();
  CVector3D N 	   = hitInfo.GetNormal();
  CVector3D dirEye = (hitInfo.GetRay())->GetDir();

  if (incEyeCos > 0.0)
    N = -N;

  // Choose one of the members of the distribution fucntion for
  // importance sampling scheme
  //----------------------------------------------------------------------
  float rnd_dist    = specular_rnd[0];
  int   choose_dist = 0;
  float chWeight    = 1.0;
  int   slopes_size = _slope->GetSlopesSize();
  int  weights_size = _slope->GetWeightsSize();
  if(slopes_size == 0 || (slopes_size != weights_size))
  {
    WARNING_INFO("For Cook-Torrance BRDF model non-empty slopes structure\n\
with non-zero weights must be specified.");
    return 0;
  }
  CTexture *chSlope = NULL;

  float alpha    = 0;
  float *a2      = new float[slopes_size];
  float *weights = new float[slopes_size];
  bool choosen   = false;

  // choose the distribution member according to its weight
  for(int i = 0; i < slopes_size; i++)
  {
    float weight = 1.0;
    CTexture *slope = _slope->GetSlope(i, weight);
    assert(slope);
    CValues s_tmp(CValues::EE_Weigth3);
    hitInfo.GetTextureValue(s_tmp, *slope, CMappingBase::EE_KSI);

    a2[i] = (s_tmp.Get(0)) * (s_tmp.Get(0));
    if(a2[i] < WTINY) a2[i] = WTINY;

    weights[i] = weight;

    if(!choosen)
    {
      alpha += weight;
      if(rnd_dist <= alpha)
      {
        choose_dist = i;
        chSlope  = slope;
        chWeight = weight;
        choosen = true;
      }
    }
  }
  //----------------------------------------------------------------------

  // Get value of the choosen distribution member
  CValues slopeC;
  hitInfo.GetTextureValue(slopeC, *chSlope, CMappingBase::EE_KSI);
  float slopeC_a = (slopeC.Get(0));

  // Make new basis with R as Z axis
  CVector3D U = ArbitraryNormal(N); // U is perpendicular to R - X-axis
  CVector3D V = CrossProd(N, U);    // V is perpendicular to R and V - Y-axis

  // TODO diffuseness of the bounce
  float diffuseness = slopeXY.xx;

  // Compute directions for maxRays outgoing rays
  //----------------------------------------------------------------------
  for(int i = maxRays; i--; num++) {
    // the ray is sampled in this direction
    CContributionItem *contrib = carray.AllocateNext();
    contrib->weigth.SetRepresentation(CValues::EE_Weigth3);

    // set the flags concerning ray type
    contrib->SetFlags(CContributionItem::EE_ReflectedRay);

    float outCos = -1;
    CVector3D dirRes(0,0,1);
    // If slope is not very small
    //-------------------------------------------------------------------------
    if(slopeC_a > 1e-06)
    {
      float cos_HI = -1;
      CVector3D hvector(0,0,1);
      float cos_HN = 0, cos2_HN = 0, sin_HN = 1, tan2_HN = 0;
      bool rejected = false;

      // rejection sampling
      //~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
      while(outCos < 0) {
        CVector2D h;
        if( !rejected ) {
          float alpha, phi;
          alpha = atan(slopeC_a*sqrt(-log(1-specular_rnd[2*num])));
          phi = 2 * M_PI * specular_rnd[2*num + 1];
         // h.SetValue(specular_rnd[2*num], specular_rnd[2*num+1]);
          float one = 1.0;
          PolarToCartesian(N, phi, alpha, one, hvector);
        }
        else {
          h.SetValue(samplingInfo.Sample(0), samplingInfo.Sample(1));

          tan2_HN  = slopeC_a * (-log(1.0-h.xx));
          assert(tan2_HN >= 0);
          cos2_HN  = (1.0/(1.0 + tan2_HN));
          cos_HN   = sqrt(cos2_HN);
          sin_HN   = sqrt(1.0 - cos2_HN);

          assert(sin_HN >= 0 && cos_HN >= 0);

          float phi_h   = 2.0 * M_PI * h.yy;

          // Outgoing direction (half-vector H) in local coordinate system
          CVector3D localDir(sin_HN * cos(phi_h), sin_HN * sin(phi_h), cos_HN);

          // Change the basis from the VUN coordinate system (local coordinate
          // system where normal R is Z-axis) to the World coordinate system
          hvector = localDir.x * U + localDir.y * V + localDir.z * N;
        }

        assert(Magnitude(hvector) >= 1.0-1e-06 &&
	       Magnitude(hvector) <= 1.0+1e-06);

	cos_HI = DotProd(hvector,-dirEye);

	if(cos_HI < 0 && fabs(cos_HI) < 1e-06) cos_HI = WTINY;
        if(cos_HI >= 0)
        {
          // Compute Light ray direction
          dirRes = (2.0 * cos_HI) * hvector + dirEye;

          assert(Magnitude(dirRes) >= 1.0-1e-06 &&
	         Magnitude(dirRes) <= 1.0+1e-06);
          outCos = DotProd(N, dirRes);  // (angle btw the light ray and normal)

	  // reject unappropriate samples
       if(outCos < 0.0){
          rejected = true;
       }
        }
        else //(outCos < 0)
        {
          rejected = true;
        }
      }
      //~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
      float cos3_HN = pow(cos_HN, 3.0f);

      // set PDF value
      contrib->prob = exp(-tan2_HN/(slopeC_a*slopeC_a))/(slopeC_a*slopeC_a * M_PI * cos3_HN) / (4*cos_HI);

      // compute distribution function for the rest distribution members
      float distribution = 0;
      if(slopes_size > 1)
      {
        float cos4_HN = cos3_HN * cos_HN;
        for(int i = 0; i < slopes_size; i++)
        {
          if(i != choose_dist)
	  {
	    float memb    = exp( -tan2_HN/a2[i] )/(M_PI * a2[i] * cos3_HN);
 	    distribution += weights[i] * memb;
	  }
        }
	distribution /= cos4_HN;
      }
      incEyeCos = fabs(incEyeCos);
      if(incEyeCos < WTINY) incEyeCos = WTINY;

      // Compute geometrical attenuation factor (accounting for
      // shadowing and masking)
      float G = 1/(4*cos_HI*cos_HI);



      // Set the value of BSDF weight for choosen distribution member
      //~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
      // weight = BRDF * cos(Inc,N)/ probability_density_of_generated
      //        = F(ks_local,cosBeta) * cosTheta_in / cosBeta

      //~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
      contrib->weigth.SetRepresentation(CValues::EE_Weigth3);
      // just avoid additional multiplication on ks_local * m1
      contrib->weigth = DotProd(N,dirRes)/ cos_HI * ks_local;

       //// set the reflected ray, direction is normalized here
      contrib->ray.Init(hitInfo.GetPoint(), dirRes, REFLECTED_RAY);
      // type of interaction
      contrib->interaction = EE_REFLECTED_SPECULAR;
      // diffuseness of the bounce
      contrib->diffuseness = diffuseness;
    }
    else
    { // If slope is very small, surface behaves like a mirror. However,
      // we should keep normalization parameter 1/(4*Pi*alpha^2) and
      // only approximate outgoing direction as mirror-reflected
      // direction and omit the exponent:
      /*
          alpha != 0, slopeXY -> 0  => exp(-tan(alpha)/slopeXY)/slopeXY -> 0
	  alpha -> 0, slopeXY -> 0  => exp(-tan(alpha)/slopeXY)/slopeXY ->
	                               1/slopeXY

	  In other words, exp(-tan(alpha)/slopeXY) behaves like delta function
      */
      // set the reflected ray, direction is normalized here
      dirRes = hitInfo.GetReflectedDir();
      outCos = DotProd(dirRes, N);
      contrib->ray.Init(hitInfo.GetPoint(), dirRes, REFLECTED_RAY);
      // type of interaction
      contrib->interaction = EE_REFLECTED_MIRROR;
      // diffuseness of the bounce
      contrib->diffuseness = 0.0;

      // exp(-tan(0)/slopeXY) = 1, cos(H,N) = 1 =>
      contrib->prob = 1.0 / (slopeC_a * slopeC_a * M_PI);

      // For mirror cosine of incident equal to the cosine of
      // reflected ray => sqrt(outCos/incEyeCos) =>
      contrib->weigth = ks_local * 0.25;
    }
    //-------------------------------------------------------------------------

    // Additional (diffuse) components in order to provide nice and simple
    // vizualisation
    //------------------------------------------------------------------------
    if(kd_local.Luminance() > 1e-06)
       contrib->weigth += kd_local * (outCos/(M_PI*contrib->prob));

#ifdef __TRAVERSAL_STATISTICS
      world->_statistics->IncReflectedRays();
#endif
  } // for i
  // delete temporary arrays
  delete [] a2;
  delete [] weights;

 // return number of generated samples
  assert(num == maxRays);
  return num;
}
//========================================================================

const char* CSimpleCookTorrance::GetHelp() const
{
	return BrdfStrings::KSimpleCookTorranceDesc;
}
//========================================================================

__END_BRDFLIB_SOURCE

