#include <opencv2/opencv.hpp>

int K=2,sample_count=100,response;
CvSize imsize=cvSize(500,500);
CvMat trainData1,trainData2,trainClasses1,trainClasses2;

void drawImage(IplImage *image,void* classifier,bool knn=true)
{
  CvMat *sample=cvCreateMat(1,2,CV_32FC1);
  for(int i=0;i<imsize.height;i++)
    for(int j=0;j<imsize.width;j++)
    {
      sample->data.fl[0]=(float)j;
      sample->data.fl[1]=(float)i;
      if (knn)
        response=((CvKNearest*)classifier)->find_nearest(sample,K,0,0,NULL,0);
      else
	      response=((CvSVM*)classifier)->predict(sample);
      cvSet2D(image,i,j,(response==1?CV_RGB(0,0,255):CV_RGB(0,150,0)));
    }

  for(int i =0;i<sample_count/2;i++)
  {
    CvPoint pt;
    pt.x=cvRound(trainData1.data.fl[i*2]);
    pt.y=cvRound(trainData1.data.fl[i*2+1]);
    cvCircle(image,pt,3,CV_RGB(0,0,150),CV_FILLED);
    pt.x=cvRound(trainData2.data.fl[i*2]);
    pt.y=cvRound(trainData2.data.fl[i*2+1]);
    cvRectangle(image,cvPoint(pt.x-2,pt.y-2),cvPoint(pt.x+2,pt.y+2),CV_RGB(0,255,0),CV_FILLED);
  }
  cvReleaseMat(&sample);
}

int main(int arc,char**argv)
{
  CvRNG rng=cvRNG(-1);
  CvMat* trainData=cvCreateMat(sample_count,2,CV_32FC1 );
  CvMat* trainClasses=cvCreateMat(sample_count,1,CV_32FC1 );
  IplImage* img_knn=cvCreateImage(imsize,IPL_DEPTH_8U,3);
  IplImage* img_svm=cvCreateImage(imsize,IPL_DEPTH_8U,3);

  cvGetRows(trainData,&trainData1,0,sample_count/2);
  cvRandArr(&rng,&trainData1,CV_RAND_NORMAL,cvScalar(200,200),cvScalar(50,50));
  cvGetRows(trainData,&trainData2,sample_count/2,sample_count);
  cvRandArr(&rng,&trainData2,CV_RAND_NORMAL,cvScalar(300,300),cvScalar(50,50));
  cvGetRows(trainClasses,&trainClasses1,0,sample_count/2);
  cvSet(&trainClasses1,cvScalar(1));
  cvGetRows(trainClasses,&trainClasses2,sample_count/2,sample_count);
  cvSet(&trainClasses2,cvScalar(2) );

	CvSVMParams params(CvSVM::C_SVC,CvSVM::POLY,5,1,0,5,0,0,NULL,cvTermCriteria(CV_TERMCRIT_ITER,sample_count,0.001));
	CvSVM svm(trainData,trainClasses,NULL,NULL,params);
  CvKNearest knn(trainData,trainClasses,0,0,K);

  drawImage(img_knn,(void*)&knn,true);
  drawImage(img_svm,(void*)&svm,false);

  cvReleaseImage(&img_knn);cvReleaseImage(&img_svm);
  cvReleaseMat(&trainClasses);cvReleaseMat(&trainData);
  return EXIT_SUCCESS;
}
