#include <math.h>
#include "potential.h"
#include "scene.h"

#include "Common.h"
#include "PatchProbabilityAssigner.h"



PatchProbabilityAssigner::PatchProbabilityAssigner()
{
  numpatches = 0;
  sum = 0.0;
  pa = NULL;
  pb = NULL;
}



PatchProbabilityAssigner::~PatchProbabilityAssigner()
{
  freeProbabilities();
}


int PatchProbabilityAssigner::assignProbabilities(PATCH *patch)
{
  // Alloc arrays ...
  if (allocProbabilities())
    {
      // Use derived classes computations
      recomputeProbabilities(patch);

      // Normalise so that numbers in pb[] becomes trues probabilities.
      // sum = 1, and always between 0 and 1. 
      normalizeProbabilities();
      
      return 1;
    }

  return 0;
}


double PatchProbabilityAssigner::getProbability(const int pid)
{ return pb[pid]; }


double PatchProbabilityAssigner::getProbability(const PATCH *patch)
{ return pb[patch->id]; }


int PatchProbabilityAssigner::choosePatchId()
{
  int    m;
  double rn;

  // Choose a patch according to probability array.
  rn = drand48();
  m = 1;
  while (m < numpatches)
    {
      if (rn <= pb[m])
	break;
      rn -= pb[m];
      m++;
    }
  
  // Numeric errors are possible...
  if (m==numpatches)
    m--;

  return m;
}


PATCH* PatchProbabilityAssigner::choosePatch()
{ return pa[choosePatchId()]; }



void PatchProbabilityAssigner::normalizeProbabilities()
{
  int i;

  // If everything is zero, fall back to simple area sampling ...
  if (sum==0.0)
    {
      printf("Probability assignment failed !\nFalling back to simple pure area sampling.\n");
      for(i=1;i<numpatches;i++)
	{
	  pb[i] = pa[i]->area;
	  sum += pb[i];
	}
    } 
  
  // Multiplication is faster than division on most computers
  sum = 1.0 / sum;

  // Normalize
  for(i=1;i<numpatches;i++)
    pb[i] *= sum;

  sum = 1.0;
}



int PatchProbabilityAssigner::allocProbabilities()
{
  PATCHLIST *pl;
  int       idx;
  int       np;

  // Get the number of patches
  np = PatchListCount(Patches)+1; // +1 for patch 0...
  if (np != numpatches)
    {
      // Free old probabilities...
      freeProbabilities();
      numpatches = np;
    }

  // Alloc arrays    
  if (pa==NULL)
    pa = new PATCH * [numpatches];
  
  if (pb==NULL)
    pb = new double [numpatches];

  // Check memory allocation
  if ((pa==NULL)||(pb==NULL))
    {
      // free what's been allocated.
      freeProbabilities();
      return 0;
    }

  // Fill array
  idx = 1;
  pl = Patches;
  ForAllPatches(cp,pl) {
    pa[idx] = cp;
    pb[idx] = 0.0;
    idx++;
  } EndForAll;

  // Others
  pb[0] = sum = 0.0;
  pa[0] = (PATCH *)NULL;

  return 1;
}


void PatchProbabilityAssigner::freeProbabilities()
{
  // Free all arrays
  if (pb!=NULL)
    {
      delete pb;
      pb = NULL;
    }

  if (pa!=NULL)
    {
      delete pa;
      pa = NULL;
    }

  // No more patches
  numpatches = 0;
  sum = 0.0;
}

void PatchProbabilityAssigner::clearProbabilities()
{
  int i;

  for (i=0;i<numpatches;i++)
    pb[i] = 0.0;
  sum = 0.0;
}




// ----------------------------------------------- AreaPatchProbabilityAssigner implementation

void AreaPatchProbabilityAssigner::recomputeProbabilities(PATCH*)
{
  int i;

  sum = 0.0;
  for (i=1;i<numpatches;i++)
    {
      pb[i] = pa[i]->area;
      sum += pb[i];
    }
}



// ----------------------------------------------- DirectPotentialPatchProbabilityAssigner implementation

void DirectPotentialPatchProbabilityAssigner::recomputeProbabilities(PATCH*)
{
  int i;

  // Other modules of renderpark are doing it...
  UpdateDirectPotential();

  sum = 0.0;
  for (i=1;i<numpatches;i++)
    {
      pb[i] = pa[i]->direct_potential;
      sum += pb[i];
    }
}
       


// ----------------------------------------------- EquiWinPatchProbabilityAssigner implementation

void EquiWinPatchProbabilityAssigner::recomputeProbabilities(PATCH*)
{
  double hmin,htmp;
  int    nhit,i;

  // Find minimum window size
  hmin = -1.0;
  for (i=1;i<numpatches;i++)
    {
      nhit = deState.isp->getNumberOfImpacts(i-1);
     
      // If the patch receive no particle htmp is infinite, so we skip this one
      if (nhit>0)
	{
	  htmp = sqrt((pa[i]->area*deState.c1)/((double)nhit));
	  if ((htmp<hmin)||(hmin<0.0f))
	    hmin = htmp;
	}
    }
  
  // This may happen in some weird cases -- fall back to AreaPatch
  if (hmin<0.0)
    hmin = 1.0;

  // Precalculate
  htmp = deState.c1 / (hmin*hmin);

  // Probability is proportionnal to the number of particles needed to get constant window size
  sum = 0.0;
  for (i=1;i<numpatches;i++)
    {
      pb[i] = (htmp*pa[i]->area) - deState.isp->getNumberOfImpacts(i-1);
      sum += pb[i];
    }
}


// ----------------------------------------------- CombinePatchProbabilityAssigner implementation


CombinePatchProbabilityAssigner::CombinePatchProbabilityAssigner(PatchProbabilityAssigner *nppa1,PatchProbabilityAssigner *nppa2)
{
  // Save
  ppa1 = nppa1;
  ppa2 = nppa2;
}


CombinePatchProbabilityAssigner::~CombinePatchProbabilityAssigner()
{
  delete ppa1;
  delete ppa2;
}


void CombinePatchProbabilityAssigner::recomputeProbabilities(PATCH* p)
{
  int i;

  // Compute ...
  ppa1->assignProbabilities(p);
  ppa2->assignProbabilities(p);
  
  // Combine ...
  sum = 0.0;
  for(i=1;i<numpatches;i++)
    {
      pb[i] = ppa1->getProbability(i) * ppa2->getProbability(i);
      sum += pb[i];
    }
}