#include "PPSuperregion.h"

#ifdef NICE_USELIB_ICE
#include <core/iceconversion/convertice.h>
#endif

#include "segmentation/RegionGraph.h"

using namespace std;
using namespace NICE;
using namespace OBJREC;

PPSuperregion::PPSuperregion()
{
  conf = new Config();
  Init();
}

PPSuperregion::PPSuperregion(const Config *_conf): conf(_conf)
{
  Init();
}

void PPSuperregion::Init()
{
  std::string section = "PostProcessSG";
  rf = new FPCRandomForests( conf, "ShapeRF" );
}

PPSuperregion::~PPSuperregion()
{
}


void PPSuperregion::optimizeShape(Examples &regions, NICE::Matrix &mask, NICE::MultiChannelImageT<double> & probabilities)
{
#ifdef NICE_USELIB_ICE
  vector<ice::Region> superregions;
  vector<double> probs;
  vector<int> classes;
  NICE::Matrix smask;
  getSuperregions(regions, mask, superregions, classes, smask);

  for (int i = 0; i < (int)superregions.size(); i++)
  {
    ice::Moments m;
    superregions[i].CalcMoments(m);

    NICE::Vector tmp = makeEVector(m.AffineHuInvariants());
    NICE::Vector *tmp2 = new NICE::Vector(tmp);
    Example tex(tmp2);

    ClassificationResult r = rf->classify ( tex );

    probs.push_back(r.scores[classes[i]]);
  }

  vector<ice::Region> orgregions;
  for (int i = 0; i < (int)regions.size(); i++)
  {
    orgregions.push_back(ice::Region());
  }

  for (int y = 0; y < (int)mask.cols(); y++)
  {
    for (int x = 0; x < (int)mask.rows(); x++)
    {
      int pos = mask(x, y);
      orgregions[pos].Add(x, y);
    }
  }

  // maps the regions to their superregions
  vector<int> regmapsreg(regions.size(), 0);
  for (int y = 0; y < (int)smask.cols(); y++)
  {
    for (int x = 0; x < (int)smask.rows(); x++)
    {
      int r = mask(x, y);
      int sr = smask(x, y);
      regmapsreg[r] = sr;
    }
  }

  RegionGraph g;
  g.computeGraph(regions, mask);

  vector<Node*> nodes;
  g.get(nodes);

  bool change = true;
  int k = 0;
  while (change && k < 100)
  {
    k++;
    change = false;
    int anders = 0;
    for (int i = 0; i < (int) nodes.size(); i++)
    {

      set<int> sr;
      int regnb = nodes[i]->getRegion();
      int orgreg = regmapsreg[regnb];

      if (nodes[i]->isAtBorder())
      {
        vector<Node*> nbs;
        nodes[i]->getNeighbors(nbs);
        for (int j = 0; j < (int)nbs.size(); j++)
          sr.insert(regmapsreg[nbs[j]->getRegion()]);
      }

      vector<double> otherprobs;

      ice::Region re = superregions[orgreg];
      re.Del(orgregions[regnb]);

      ice::Moments m;

      if (re.Area() > 0)
      {
        re.CalcMoments(m);

        NICE::Vector tmp = makeEVector( m.AffineHuInvariants());
        NICE::Vector *tmp2 = new NICE::Vector(tmp);
        Example tex(tmp2);
        ClassificationResult r = rf->classify ( tex );
        tex.vec = NULL;
        delete tmp2;

        double val = probabilities.get(regions[regnb].second.x, regions[regnb].second.y, classes[orgreg]) * r.scores[classes[orgreg]];

        otherprobs.push_back(val);
        if (otherprobs[0] < probs[orgreg])
          continue;
      }

      for ( set<int>::const_iterator iter = sr.begin();iter != sr.end();++iter )
      {
        ice::Moments m2;
        ice::Region re2 = superregions[regmapsreg[*iter]];
        re2.Add(orgregions[regnb]);
        re2.CalcMoments(m2);
        NICE::Vector tmp = makeEVector(m2.AffineHuInvariants());
        NICE::Vector *tmp2 = new NICE::Vector(tmp);
        Example tex(tmp2);
        ClassificationResult r2 = rf->classify ( tex );
        tex.vec = NULL;
        delete tmp2;

        double val = probabilities.get(regions[regnb].second.x, regions[regnb].second.y, classes[*iter]) * r2.scores[classes[*iter]];

        otherprobs.push_back(val);
      }

      int k = 1;
      int best = -1;
      double bestval = -1.0;
      for ( set<int>::const_iterator iter = sr.begin();iter != sr.end();++iter, k++ )
      {
        if (otherprobs[k] > probs[*iter])
        {
          if (bestval < otherprobs[k])
          {
            bestval = otherprobs[k];
            best = *iter;
          }
        }
      }

      if (best < 0 || bestval <= 0.0)
        continue;

      change = true;

      probs[best] = bestval;

      superregions[best].Add(orgregions[regnb]);

      probs[orgreg] = otherprobs[0];

      superregions[orgreg].Del(orgregions[regnb]);

      regmapsreg[regnb] = best;

      nodes[i]->setLabel(classes[best]);
      anders++;
    }
  }

  for (int i = 0; i < (int)regions.size(); i++)
  {
    regions[i].first = classes[regmapsreg[i]];
  }
#else
  throw("PPSuperRegion.cpp: please use ice library for this function");
#endif
}
#ifdef NICE_USELIB_ICE
void PPSuperregion::getSuperregions(const Examples &regions, const NICE::Matrix &mask, vector<ice::Region> &superregions, vector<int> &classes, NICE::Matrix &smask)
{
  NICE::Image tmp (mask.rows(), mask.cols());
  tmp.set(0);
  NICE::ColorImage m2 (tmp, tmp, tmp);
  for (int y = 0; y < (int)mask.cols(); y++)
  {
    for (int x = 0; x < (int)mask.rows(); x++)
    {
      int pos = mask(x, y);

      m2.setPixel(x, y, 0, regions[pos].first);
      m2.setPixel(x, y, 1, regions[pos].first);
      m2.setPixel(x, y, 2, regions[pos].first);
    }
  }

  RSMeanShift rs(conf);
  int count = rs.transformSegmentedImg( m2, smask);

  classes.resize(count);
  for (int i = 0; i < count; i++)
  {
    superregions.push_back(ice::Region());
  }

  for (int y = 0; y < (int)smask.cols(); y++)
  {
    for (int x = 0; x < (int)smask.rows(); x++)
    {
      int pos = smask(x, y);
      superregions[pos].Add(x, y);
      classes[pos] = regions[mask(x,y)].first;
    }
  }
}
#endif

void PPSuperregion::trainShape(Examples &regions, NICE::Matrix &mask)
{
#ifdef NICE_USELIB_ICE
  // bestimme Superregionen
  vector<ice::Region> superregions;
  vector<int> classes;
  // refactor-nice.pl: check this substitution
  // old: Image smask;
  NICE::Matrix smask;
  getSuperregions(regions, mask, superregions, classes, smask);

  // berechne die Momente der Superregionen und speichere diese als Merkmale ab
  for (int i = 0; i < (int)superregions.size(); i++)
  {
    ice::Moments m;
    superregions[i].CalcMoments(m);
    NICE::Vector tmp = makeEVector(m.AffineHuInvariants());
    NICE::Vector *tmp2 = new NICE::Vector(tmp);
    shapefeats.push_back(pair<int, Example>(classes[i], Example(tmp2)));
  }
#else
  throw("PPSuperRegion.cpp: please use ice library for this function");
#endif
}

void PPSuperregion::finishShape(ClassNames &cn)
{
  //Lerne Klassifikator mit dem den Formmerkmalen an
  FeaturePool fp;
  Feature *f = new VectorFeature ( 7 );
  f->explode ( fp );
  delete f;
  rf->train ( fp, shapefeats);
}

void PPSuperregion::restore (istream & is, int format)
{

}

void PPSuperregion::store (ostream & os, int format) const
{

}

void PPSuperregion::clear()
{

}