#include "stdafx.h"
#include "analyze3.h"
#include "Blob.h"
#include "TrackedObject.h"
#include "BGModel.h"
#include "Analyzer.h"
#include "timer.h"

#ifdef _DEBUG
#undef THIS_FILE
static char THIS_FILE[]=__FILE__;
#define new DEBUG_NEW
#endif

//////////////////////////////////////////////////////////////////////

enum{
   vRaw = 0,
   vSeg,
   vNUM
} Vwnds;

char *vName[] = { "Raw Image", "Segmented (per pixel) Image" };

#define NUM_COLORS   1   /* 0  */
BYTE colors[NUM_COLORS][3] = 
{
    { 255, 0, 0 }    /*,
   { 0, 255, 0 },   
   { 255, 255, 0 },
   { 255, 130, 255 },
   { 0, 255, 255 },
   { 160, 70, 110 },
   { 110, 150, 220 },
   { 50, 120, 20 },
   { 210, 160, 40 },
   { 20, 220, 170 }    */
};

char *colorNames[] =
{
   "Red"   /*  , "Green", "Yellow", "Pink", "Cyan", "Dark Purple",
      "Sky Blue", "Dark Green", "Orange", "Teal"
      */
};

//////////////////////////////////////////////////////////////////////

Analyzer::Analyzer()
{
   bgm = NULL;
   pLabels = NULL;
   pFgMask = NULL;
   points = NULL;
   nBlobs = 0;
   nTobs = 0;
   int i;
   for( i = 0; i<MAX_NUM_BLOBS; i++) blobs[i] = NULL;
   for( i = 0; i<MAX_NUM_TOBS; i++) tobs[i] = NULL;
}

//////////////////////////////////////////////////////////////////////

Analyzer::~Analyzer()
{
   Cleanup();
}

//////////////////////////////////////////////////////////////////////

void Analyzer::OnStart(VIDEOINFOHEADER *pvih)
{   
   Cleanup();
   AbstractAnalyzer::OnStart(pvih);
   
   if (!CreateVwnds(vNUM)) return;

   CvSize size;
   size.width = vw;
   size.height = vh;
   pFgMask = cvCreateImage(size, 8, 1);
   iplSet(pFgMask, 0);
   pLabels = cvCreateImage(size, 8, 3);   
   iplSet(pLabels, 0);
   
   bgm = new BGModel(vw, vh);

   int i;
   for( i = 0; i<MAX_NUM_BLOBS; i++) blobs[i] = new Blob();
   nBlobs = 0;

   for( i = 0; i<MAX_NUM_TOBS; i++) tobs[i] = new TrackedObject();
   nTobs = 0;

   points = new POINTS[nPixels];   
      
   for( i = 0; i<vNUM; i++)
   {
      vwnds[i] = new CVideoWnd(vName[i], vw, vh);
      Call(vwnds[i], ShowWindow(SW_SHOWNORMAL));
   }   
   
}

//////////////////////////////////////////////////////////////////////

bool Analyzer::BuildLabelImageFromMask(IplImage *pMask, IplImage *pLabels)
{
   if (!pMask || !pLabels) return false;

   int *mask = (int*)pMask->imageData;
   BYTE *label = (BYTE*)pLabels->imageData;

   int i;
   for( i = 0; i<nPixels; i++)
   {
      if (*mask)
      {
         int id = *mask;
         *label = (BYTE)(id * 42);
         *(label+1) = (BYTE)(id * 8);
         *(label+2) = (BYTE)(255 - id * 6);
      }
      else
      {
         *label = 0;
         *(label+1) = 0;
         *(label+2) = 0;
      }

      mask++;
      label += 3;
   }

   return true;
}

//////////////////////////////////////////////////////////////////////

/*
 Finds connected components in binary image pMask (8x1) using a 2-pass
  method.  Stores resulting labels in pLabels (32x1) such that every pixel in
  a blob has the same label and each blob has a unique label.
 */
bool Analyzer::FindConnectedComponents(IplImage *pMask, IplImage *pLabels)
{
   if (!pMask || !pLabels) return false;
   
   BYTE *mask = (BYTE*)pMask->imageData;
   int *id = (int*)pLabels->imageData;
   int x, y, i, j, k, nextId;
   int ids[4];
   int idOfs[4] = { -vw-1, -vw, -vw+1, -1 };   

   // to start, every id value maps to itself
   nextId = 1;
   for( i = 0; i<MAX_CC_IDS; i++) map[i] = i;
   
   // scan first pixel as a special case
   if (*mask)
   {
       *id = nextId++;
   }
   else
   {
       *id = 0;
   }
   mask++;
   id++;

   // scan rest of first row as a special case   
   for(x=1; x<vw; x++)
   {
      if (*mask)
      {
         j = *(id - 1);
         if (j > 0)
         {
             *id = j;
         }
         else 
         {
             *id = nextId++;
         }
      }
      else 
      {
          *id = 0;
      }
      
      mask++;
      id++;
   }

   // scan rest of rows
   for(y=1; y<vh; y++)
   {
      // check first pixel of row as a special case
      if (*mask)
      {
         i = *(id - vw);
         j = *(id - vw + 1);
         
         if (j>i)
         {
             i = j;
         }
         if (i>0) 
         {
             *id = i;
         }
         else 
         {
             *id = nextId++;
         }
      }
      else 
      {
          *id = 0;
      }
      mask++;
      id++;

      // now check the 'middle' of the row
      for(x=1; x<vw-1; x++)
      {
         if (*mask)
         {
            j = 0;
            // find the max neighbor
            for( i = 0; i<4; i++)
            {
               k = *(id + idOfs[i]);
               ids[i] = map[k];
               if (ids[i] > j) j = ids[i];
            }

            if (j > 0)
            {
               for( i = 0; i<4; i++)
               {
                  if (ids[i]==0 || ids[i]==j) continue;
                  for(k=1; k<nextId; k++)
                  {
                     if (map[k]==ids[i]) map[k] = j;
                  }                  
               }
               *id = j;
            }
            else
            {
               *id = nextId++;
            }
         }
         else
         {
             *id = 0;
         }

         mask++;
         id++;
      }

      // finally, we can check the last pixel of the row as a special case
      if (*mask)
      {
         i = *(id - vw - 1);
         j = *(id - vw);         
         if (j>i) i = j;
         
         j = *(id - 1);
         if (j>i) i = j;

         if (i>0)
         {
             *id = i;
         }
         else 
         {
             *id = nextId++;
         }
      }
      else
      {
          *id = 0;
      }
      mask++;
      id++;

      if (nextId >= MAX_CC_IDS)
      {
         dbg("Error - not enough connected component ids (%d)\n", MAX_CC_IDS);
         return false;
      }
   }

   // pass 2 - update ids in label image according to equiv map
   id = (int*)pLabels->imageData;

   for( i = 0; i<nPixels; i++)
   {
      if (*id > 0) *id = map[*id];
      id++;
   }

   return true;
}

//////////////////////////////////////////////////////////////////////

/*
 Extracts blob info from labelled connected component image

 pLabels - 32x1 labelled image
 blobs - pointer to MAX_NUM_BLOBS blob object pointers
 nBlobs - pointer to int that will store number of valid blobs found
 minSize - minimum size for a blob to be valid
 */
bool Analyzer::FindBlobs(IplImage *pLabels, Blob **blobs, int *nBlobs, int minSize)
{
   if (!pLabels || !blobs || !nBlobs) return false;

   int *id = (int*)pLabels->imageData;

   int x, y, i, j, n;

   *nBlobs = 0;
   for( i = 0; i<MAX_CC_IDS; i++) map[i] = -1;

   n = 0;
   for(y=0; y<vh; y++)
   {
      // blob[X]->user will serve a dual purpose:
      //  1) flag to say whether the blob 'grew' due to the last row
      //  2) index into map array that references this blob
      for( i = 0; i<n; i++) blobs[i]->user = 0;

      for(x=0; x<vw; x++)
      {
         j = *id;         
         // is this a blob pixel?
         if (j > 0)
         {
            if (map[j] < 0)
            {
               if (n >= MAX_NUM_BLOBS)
               {
                  dbg("error: too many blobs (%d)\n", n);
                  *nBlobs = 0;
                  return false;
               }
               else
               {
                  // this is a new blob
                  map[j] = n;
                  blobs[n]->mass = 1;
                  blobs[n]->nHull = 0;
                  blobs[n]->centroid.x = x;
                  blobs[n]->centroid.y = y;
                  blobs[n]->bbBottomRight.x = x;
                  blobs[n]->bbBottomRight.y = y;
                  blobs[n]->bbTopLeft.x = x;
                  blobs[n]->bbTopLeft.y = y;
                  blobs[n]->id = j;
                  blobs[n]->user = j;
                  n++;
               }
            }
            else
            {
               // this blob already exists
               blobs[map[j]]->user = j;
               j = map[j];
               blobs[j]->mass++;
               blobs[j]->centroid.x += x;
               blobs[j]->centroid.y += y;
               if (x > blobs[j]->bbBottomRight.x) blobs[j]->bbBottomRight.x = x;
               else if (x < blobs[j]->bbTopLeft.x) blobs[j]->bbTopLeft.x = x;
               if (y > blobs[j]->bbBottomRight.y) blobs[j]->bbBottomRight.y = y;
               else if (y < blobs[j]->bbTopLeft.y) blobs[j]->bbTopLeft.y = y;               
            }
         }

         id++;
      }      

      // check for finished blobs
      for( i = 0; i<n; i++)
      {
         if ((blobs[i]->user == 0) && (blobs[i]->mass < minSize))
         {
            // kill this blob by moving it to the end of the list and
            //  decrement the list size
            n--;

            // we have to do a proper swap
            Blob *tblob = blobs[i];
            blobs[i] = blobs[n];
            blobs[n] = tblob;            

            // we also have to update the map            
            map[blobs[i]->user] = i;

            // since we changed the blob at position <i> we want to reprocess
            //  it next time through the loop
            i--;
         }
      }
   }

   // do some per blob post-processing
   for( i = 0; i<n; i++)
   {
      // we accumulated the sum of all (x, y) values
      //  now we need to divide by the mass to get the centroid
      blobs[i]->centroid.x /= blobs[i]->mass;
      blobs[i]->centroid.y /= blobs[i]->mass;
   }

   *nBlobs = n;

   return true;
}

//////////////////////////////////////////////////////////////////////

bool Analyzer::ColorTob(IplImage *pImg, IplImage *pLabels,
                        TrackedObject *tob, BYTE *color)
{
   if (!pImg || !tob) return false;

   BYTE *pi = (BYTE*)pImg->imageData;
   int *id = (int*)pLabels->imageData;
   int w, h, x, y, x3, pitch;

   Blob *blob = tob->blob;

   w = blob->bbBottomRight.x - blob->bbTopLeft.x;
   h = blob->bbBottomRight.y - blob->bbTopLeft.y;

   x = blob->bbTopLeft.y * vw + blob->bbTopLeft.x;
   pi += x*3;
   id += x;
   pitch = vw*3;

   for(y=0; y<h; y++)
   {
      x3 = 0;
      for(x=0; x<w; x++)
      {
         if (blob->id == *(id + x))
         {
            *(pi + x3) = color[2];
            *(pi + x3 + 1) = color[1];
            *(pi + x3 + 2) = color[0];
         }

         x3 += 3;
      }

      pi += pitch;
      id += vw;
   }

   return true;
}

//////////////////////////////////////////////////////////////////////

bool Analyzer::MatchBlobsTobs(Blob **blobs, int nBlobs,
                              TrackedObject **tobs, int *nTobs)
{
   if (!blobs || !tobs || !nTobs) return false;

   int i;
   int n = nBlobs;
   if (n > MAX_NUM_TOBS) n = MAX_NUM_TOBS;

   for( i = 0; i<n; i++)
   {
      tobs[i]->blob = blobs[i];
   }
   
   *nTobs = n;

   return true;
}

//////////////////////////////////////////////////////////////////////

void Analyzer::Handle(PFrameData pfd)
    {
   if (!pfd)
   {
      dbg("error: NULL pfd\n");
      return;
   }

   // store the raw data in a cv image
   CopyMemory(pFrame->imageData, pfd->pData, nBytes);

   // determine foreground pixels
   if (!bgm->Process(pFrame, pFgMask))
   {
      dbg("Error classifying current frame (%d) / updating bg model\n", (int)pfd->frame);
   }   

   // label each pixel with a blob-unique id
   FindConnectedComponents(pFgMask, pTempf1[0]);

   // now extract blobs
   int minSize = 8;
   FindBlobs(pTempf1[0], blobs, &nBlobs, minSize);   

   MatchBlobsTobs(blobs, nBlobs, tobs, &nTobs);
   //BuildLabelImageFromMask(pTempf1[0], pLabels);

   iplSet(pLabels, 0);

   for(int i=0; i<nTobs; i++)
   {      
      ColorTob(pLabels, pTempf1[0], tobs[i], colors[tobs[i]->id % NUM_COLORS]);
      blobs[i]->RenderBB(pFrame, RGB(0, 128, 200));
   }


   // render the data
   Call(vwnds[vRaw], SetData((BYTE*)pFrame->imageData));
   Call(vwnds[vSeg], SetData((BYTE*)pLabels->imageData));
}

//////////////////////////////////////////////////////////////////////

void Analyzer::Cleanup()
{
   AbstractAnalyzer::Cleanup();
   cvReleaseImage(&pLabels);
   cvReleaseImage(&pFgMask);
   int i;
   for( i = 0; i<MAX_NUM_BLOBS; i++) DeleteClean(blobs[i]);
   for( i = 0; i<MAX_NUM_TOBS; i++) DeleteClean(tobs[i]);
   DeleteClean(bgm);
   DeleteClean(points);
}

//////////////////////////////////////////////////////////////////////