/* LICENSE:
  =========================================================================
    CMPack'04 Source Code Release for OPEN-R SDK 1.1.5-r2 for ERS7
    Copyright (C) 2004 Multirobot Lab [Project Head: Manuela Veloso]
    School of Computer Science, Carnegie Mellon University
    All rights reserved.
  ========================================================================= */


#include <stdlib.h>
#include <stdio.h>
#include <iostream>
using namespace std;

#include <sys/types.h>
#include <sys/stat.h>
#include <unistd.h>

#include <FL/Fl.H>
#include <FL/Fl_Window.H>
#include <FL/fl_draw.H>

#include <sys/time.h>
#include <unistd.h>

char *model = "ERS7";

#include "../shared/image.h"

#include "../../agent/Motion/Kinematics.h"
#include "../../agent/Vision/VisionInterface.h"
#include "../../agent/Vision/Vision.h"

bool UsePlanarYUV = true;
bool ProfileMode = false;
bool SortMode    = false;

struct timeval tv1;
struct timeval tv2;

bool TestBWBallVision = false;

bool doAverageColor = false;
bool doRobotDist = true;
bool fixRobotColors = false;

bool Interactive = !(ProfileMode || SortMode);

int RobotModelNum = 7;

static const int NumSortClasses = 6;

static int width=208;//176
static int height=160;//144
Kinematics kin;

int RepeatCount=100;

extern bool ShouldThresholdImage;

int img_idx=0;

class MyDrawWidget : public Fl_Widget { 

  unsigned char * ptr;

  int x1,y1,x2,y2;

  void draw() {
    fl_draw_image(ptr,x(),y(),w(),h());
    // overlay rect
    //    fl_overlay_rect(x1,y1,x2-x1+1,y2-y1+1);
    // colored rec
    fl_color(255,0,255);
    // 4 width line
    fl_line_style(FL_SOLID, 4);
    fl_rect(x1,y1,x2-x1+1,y2-y1+1);
  }

  int handle(int e) {
    switch (e) {
    case FL_SHORTCUT:
      switch (Fl::event_key()) {
      case FL_Left:
      case 'b':
	// Go back one image
	img_idx--;
	break;
      case FL_Right:
      case ' ':
	// Go forward an image
	img_idx++;
	break;
      case 'q':
	exit(0);
	break;
      default:
	return 0;
      }
    default:
      return 0;
    }
    return 1;
  }

 public:
  MyDrawWidget(int x, int y, int w, int h, const char * s=NULL) :
    Fl_Widget(x,y,w,h,s) { ptr = new unsigned char[w*h*3];
  x1=y1=x2=y2=0;}

  void setPtr(unsigned char * p) {
    memcpy(ptr,p,width*height*3); 
    redraw();
  }

  void setBBox(int _x1, int _y1, int _x2, int _y2) {
    x1=_x1; y1=_y1; x2=_x2; y2=_y2; 
    redraw();
  }
};

Fl_Window * rle_win;
Fl_Window * rgb_win;

MyDrawWidget * rle_widget;
MyDrawWidget * rgb_widget;

rgb yuv2rgb(yuv c)
{
  rgb r;
  int y,u,v;

  y = c.y;
  u = c.u*2 - 255;
  v = c.v*2 - 255;

  r.red   = bound(y + u                     ,0,255);
  r.green = bound((int)(y - 0.51*u - 0.19*v),0,255);
  r.blue  = bound(y + v                     ,0,255);

  return(r);
}

void yuv2rgb(rgb *dest,yuv *src,int width,int height)
{
  int size = width*height;
  int i;

  for(i=0; i<size; i++){
    dest[i] = yuv2rgb(src[i]);
  }
}

void fix_image(SImage &img)
// Gets rid of garbage bits at the bottom of all
// dog images by copying an adjacent row over it
{
  int i;
  rgb r;

  for(i=0; i<16; i++){
    r = img.getpixel_fast(i,img.h-2);
    img.setpixel(i,img.h-1,r);
  }
}




int main(int argc, char *argv[])
{
  int next_opt;
  for(next_opt=1; next_opt < argc && argv[next_opt][0]=='-'; next_opt++) {
    char *opt;
    opt = argv[next_opt];
    if(strcmp(opt,"--sort")==0) {
      SortMode = true;
    } else if(strcmp(opt,"--profile")==0) {
      ProfileMode = true;
    } else if(strcmp(opt,"--model")==0) {
      next_opt++;
      if(!(next_opt < argc)) {
        fprintf(stderr,"must specify numeric argument after --model\n");
        exit(6);
      }
      RobotModelNum = atoi(argv[next_opt]);      
    } else if(strcmp(opt,"--repeat")==0) {
      next_opt++;
      if(!(next_opt < argc)) {
        fprintf(stderr,"must specify numeric argument after --repeat\n");
        exit(5);
      }
      RepeatCount = atoi(argv[next_opt]);
    }
  }
  argc -= next_opt;
  argv = &argv[next_opt];

  Interactive = !(ProfileMode || SortMode);

  switch(RobotModelNum){
    case 7:
      width  = 208;
      height = 160;
      kin = Kinematics(Kinematics::ERS7);
      break;
    case 210:
      width  = 176;
      height = 144;
      kin = Kinematics(Kinematics::ERS210);
      break;
    default:
      printf("invalid robot model number %d\n",RobotModelNum);
      exit(10);
  }

  FILE *sort_files[NumSortClasses];

  if(SortMode) {
    for(int sort_idx=0; sort_idx<NumSortClasses; sort_idx++) {
      char filename[256];
      sprintf(filename,"%01d.vfl",sort_idx);
      if((sort_files[sort_idx]=fopen(filename,"w"))==NULL) {
        fprintf(stderr,"couldn't open file '%s' for writing",filename);
        exit(7);
      }
    }
  }

  bool grab,run;

  grab = false;
  run = true;

  Vision vision;
  CMVision::image_idx<rgb> vis_img;
  CMVision::image<rgb> vis_img_raw;

  // Pass the image to vision in the exactly same format as on the
  // robot.
  CMVision::image_yuv<const uchar> vis_img_raw_robot;
  uchar yuv_planar_data[3*width*height];

  VisionInterface::ObjectInfo obj_info;

  vision.initialize("vision.cfg",width,height);
  vis_img.width =width;
  vis_img.height=height;
  vis_img.row_stride=width;

  vis_img_raw.width =width;
  vis_img_raw.height=height;
  vis_img_raw.row_stride=width;

  vis_img_raw_robot.width = width;
  vis_img_raw_robot.height = height;
  vis_img_raw_robot.row_stride = 3*width;
  vis_img_raw_robot.buf_y = yuv_planar_data;
  vis_img_raw_robot.buf_u = yuv_planar_data + width;
  vis_img_raw_robot.buf_v = yuv_planar_data + 2*width;

  double head_angles[3]={0.0,0.0,0.0};
  double body_angle=16.0 * M_PI / 180.0;
  double body_height=104.0;

  kin.getHeadAngles(head_angles,vector3d(1000.0,0.0,199.216774326),body_angle,body_height);

  char ang_filename[256],*img_filename;
  char **img_filenames;
  int num_imgs=0;
  struct stat stats;
  int last_frame_idx;
  SImage img;
  SImage img_processed;

  last_frame_idx = -1;
  img.allocate(width,height);
  img_processed.allocate(width,height);

  if(argc < 1) {
    fprintf(stderr,"must specify at least one image file\n");
    exit(1);
  }

  num_imgs = argc;
  img_filenames=&argv[0];

  rle_win = new Fl_Window(width,height,"Vision Test RLE");
  rle_widget = new MyDrawWidget(0,0,width,height);
  rle_win->show();
  rgb_win = new Fl_Window(width,height,"Vision Test RGB");
  rgb_widget = new MyDrawWidget(0,0,width,height);
  rgb_win->show();
  Fl::check();

  while(run) {
    memset(&obj_info,0,sizeof(obj_info));

    if(img_idx < 0) {
      printf("already at first image\n");
      img_idx = 0;
    }

    if(img_idx >= num_imgs) {
      printf("no more images\n");
      img_idx = num_imgs-1;
    }

    if(img_idx!=last_frame_idx) {
      char *img_basestart;
      bool raw_image=false;
      img_filename=img_filenames[img_idx];
      img_basestart = strrchr(img_filename,'/');
      if(img_basestart==NULL)
        img_basestart = img_filename;
      else
        img_basestart++;
      raw_image = (*img_basestart == 'i');

      strncpy(ang_filename,img_filename,256);
      ang_filename[255]=0;
      ang_filename[img_basestart-img_filename]='i';
      strcpy(&ang_filename[strlen(ang_filename)-3],"inf");
      
      if(stat(img_filename,&stats)==-1) {
        printf("could not stat %s\n",img_filename);
        break;
      }
      
      if(stat(ang_filename,&stats)==-1) {
        printf("could not stat %s\n",ang_filename);
        exit(123);
      } else {
        FILE *ang_file;
        double a0,a1,a2,a3,a4;
        
        if((ang_file=fopen(ang_filename,"rt"))==NULL)
          goto ang_error;
        
        if(fscanf(ang_file,"%lg %lg %lg %lg %lg",&a0,&a1,&a2,&a3,&a4)!=5)
          goto ang_error;
        
        body_angle     = a0;
        body_height    = a1;
        head_angles[0] = a2;
        head_angles[1] = a3;
        head_angles[2] = a4;
        goto ang_ok;
        
      ang_error:
        fprintf(stderr,"error reading angles from file\n");

      ang_ok:
        if(ang_file)
          fclose(ang_file);
      }
      
      if(!img.load(img_filename)) {
        printf("error loading file %s\n",img_filename);
        break;
      }
      
      if(img.h != height || img.w!=width){
        printf("image '%s' has wrong size w=%d h=%d, expected w=%d h=%d\n",
               img_filename,img.w,img.h,width,height);
        exit(5);
      }
      
      vis_img.buf=img.img;
      vis_img_raw.buf=img.img;
      
      // Build YUV planar image 
      {
	uchar *y_buf = (uchar *)vis_img_raw_robot.buf_y;
	uchar *u_buf = (uchar *)vis_img_raw_robot.buf_u;
	uchar *v_buf = (uchar *)vis_img_raw_robot.buf_v;

	for(int h=0; h<height; h++) {
	  for(int w=0; w<width; w++) {
	    y_buf[w] = img.getpixel_fast(w, h).red;
	    u_buf[w] = img.getpixel_fast(w, h).green;
	    v_buf[w] = img.getpixel_fast(w, h).blue;
	  }
	  y_buf += vis_img_raw_robot.row_stride;
	  u_buf += vis_img_raw_robot.row_stride;
	  v_buf += vis_img_raw_robot.row_stride;
	}
      }

      vision.setHeadAngles(head_angles);
      vision.setBodyParams(body_angle,body_height);

      ulong start_time,end_time;
      start_time = GetTime();
      
      for(int cnt=0; cnt<(ProfileMode ? RepeatCount : 1); cnt++) {
        if(ProfileMode || !(img_idx==last_frame_idx)) {
          printf("processing frame %d ('%s')\n",img_idx,img_filename);
          if(!Interactive && (!ProfileMode || cnt==0))
            fprintf(stderr,"processing frame %d ('%s')\n",img_idx,img_filename);
          if(raw_image) {
	    if(UsePlanarYUV)
	      vision.runLowLevelVision(vis_img_raw_robot);
	    else
	      vision.runLowLevelVision(vis_img_raw);
          } else {
            ShouldThresholdImage = (cnt==0);
            vision.runLowLevelVision(vis_img);
          }
          if(Interactive) {
	    //            CMVision::RmapToRgb(img_processed.img,vision.rmap,vision.num_runs,width,height,vision.color,vision.num_colors);
	    
	    //	    yuv2rgb(img.img,reinterpret_cast<yuv*>(img.img),width,height);
	    //	    rgb_widget->setPtr(reinterpret_cast<unsigned char*>(img.img));
	    //	    rle_widget->setPtr(reinterpret_cast<unsigned char*>(img_processed.img));
	    //rle_widget->setBBox(10,10,20,20);
          }
          if(raw_image) {
	    if(UsePlanarYUV)
	      vision.runHighLevelVision(&obj_info, vis_img_raw_robot);
	    else
	      vision.runHighLevelVision(&obj_info,vis_img_raw);
          }else{
            vision.runHighLevelVision(&obj_info,vis_img);
          }

	  if(TestBWBallVision) {
	    if(UsePlanarYUV)
	      vision.processBWBall(vis_img_raw_robot);
	    else
	      printf("Can only find BWBall with planar YUV images.\n");
	  }
	  
          if(!ProfileMode || cnt==0) {
            printf("done processing frame %d ('%s')\n",img_idx,img_filename);
	  }
          if(SortMode) {
            if(0 <= DebugClass && DebugClass < NumSortClasses) {
              fprintf(sort_files[DebugClass],"%02d %s %s\n",DebugClass,img_filename,DebugInfo);
            } else {
              fprintf(stderr,"invalid sort class %d (out of range)\n",DebugClass);
            }
          }
        }
        last_frame_idx = img_idx;
      }

      end_time = GetTime();

      double elapsed_time;
      elapsed_time = (end_time-start_time)/1.0e6;
      if(ProfileMode)
        elapsed_time /= RepeatCount;

      fprintf(stderr,"frame %d ('%s') took %g seconds\n",img_idx,img_filename,elapsed_time);

      //for(int robot_num=0; robot_num<3; robot_num++) {
      //  VisionInterface::VObject *obj;
      //  obj = &obj_info.red_robots[robot_num];
      //  printf("red robot %d\n",robot_num);
      //  printf("  conf     %10g\n",obj->confidence);
      //  printf("  distance %10g left %10g right %10g\n",obj->distance,obj->left,obj->right);
      //  printf("  loc (%10g,%10g,%10g)\n",obj->loc.x,obj->loc.y,obj->loc.z);
      //}
      //
      //for(int robot_num=0; robot_num<3; robot_num++) {
      //  VisionInterface::VObject *obj;
      //  obj = &obj_info.blue_robots[robot_num];
      //  printf("blue robot %d\n",robot_num);
      //  printf("  conf     %10g\n",obj->confidence);
      //  printf("  distance %10g left %10g right %10g\n",obj->distance,obj->left,obj->right);
      //  printf("  loc (%10g,%10g,%10g)\n",obj->loc.x,obj->loc.y,obj->loc.z);
      //}
    }

    if(Interactive) {
      img_idx++;
      if(img_idx >= num_imgs)
        run = false;
    }

 
    if(img_idx != last_frame_idx){
      //Set up red regions
      color_class_state *red;
      red=&vision.color[COLOR_ORANGE];
      region *r_reg;
      r_reg=red->list;

      //Set up stats
      int xmin,xmax,ymin,ymax,a;
      int isDog,n;
      double cen_y,ysum,xsum,cen_x;
      float d;
      xmin = 0;
      xmax = 0;
      ymin = 0;
      ymax = 0;
      ysum = 0;
      xsum = 0;
      a = 0;
      isDog = 0;
      cen_y = 0.0;
      
	
      //Fix Robot Colors
      if(fixRobotColors){
	yuv blob_color;
	bool red = true;
	
	//RED
	if(red){
	  int i;
	  color_class_state *orange;
	  orange=&vision.color[COLOR_ORANGE];
	  region *o_reg;
	  o_reg=orange->list;
	  
	  color_class_state *rred;
	  rred=&vision.color[COLOR_RED];
	  region *rr_reg;
	  rr_reg = rred->list;

	  region *prev_o_reg;
	  for (i=0; o_reg; o_reg = o_reg->next, i++){
	    blob_color =  AverageColor(vis_img_raw_robot,
				       vision.rmap,
				       o_reg->run_start);
	    int y = blob_color.y;
	    int u = blob_color.u;
	    int v = blob_color.v;
	    printf("y: %d u: %d v: %d\n",y,u,v);
	    //Check the region
	    if(v > 107){
	      printf("changed color!!!\n");
	      o_reg->color = 7;

	      a+= o_reg->area;
	      if(o_reg->area > 5){
		//a += o_reg->area;
		
		if (o_reg->x2 > xmax || (xmax ==0))
		  xmax = o_reg->x2;
		if (o_reg->x1 < xmin || (xmin == 0))
		  xmin = o_reg->x1;
		if (o_reg->y2 > ymax || (ymax == 0))
		  ymax = o_reg->y2;
		if (o_reg->y1 < ymin || (ymin == 0))
		  ymin = o_reg->y1;
	      }
	      
	      //if(xmin == 0)
	      //xmin = o_reg->x1;
	      //if(ymin == 0)
	      //ymin = o_reg->y1;
	      //rle_widget->setBBox(o_reg->x1,o_reg->y1,o_reg->x2,o_reg->y2);
	      //Fl::check();

	      //   printf("Should this region be not blue? (0 or 1)\n");
	      //scanf("%d",&isDog);

	      //vision.rle_t cur_run;
	      int run_idx = o_reg->run_start;
	      while(run_idx != 0){// < what?  not vision.num_runs
		//cur_run = vision.rmap[run_idx];
		//if((&vision.rmap[run_idx])->color == 1){
		vision.rmap[run_idx].color = 7;
		  //}
		//printf("changed cclass\n");
		run_idx=vision.rmap[run_idx].next;
	      }

	      //now move to red list
	      //prev_o_reg = o_reg->next;
	      //o_reg->next = rred->list;
	      //rred->list = o_reg;
	      
	    }
	  } 
	}//End RED




      }
      //End of Fix Robot Colors

      //vision.runHighLevelVision(&obj_info, vis_img);    

	 CMVision::RmapToRgb(img_processed.img,vision.rmap,vision.num_runs,width,height,vision.color,vision.num_colors);
	 yuv2rgb(img.img,reinterpret_cast<yuv*>(img.img),width,height);
	 rgb_widget->setPtr(reinterpret_cast<unsigned char*>(img.img));
	 rle_widget->setPtr(reinterpret_cast<unsigned char*>(img_processed.img));
	 Fl::check();



	 for (n=0; r_reg; r_reg = r_reg->next, n++){

	 
	   //Calculate stats for this blob	 
	   //if (r_reg->x2 > xmax)
	   //xmax = r_reg->x2;
	   //if (r_reg->x1 < xmin || (xmin == 0))
	   //xmin = r_reg->x1;
	   //if (r_reg->y2 > ymax)
	   //ymax = r_reg->y2;
	   //if (r_reg->y1 < ymin || (ymin == 0))
	   //ymin = r_reg->y1;
	   //a += r_reg->area;
	   ysum += r_reg->cen_y;
	   xsum += r_reg->cen_x;
	   //a=0;

	   //rle_widget->setBBox(r_reg->x1,r_reg->y1,r_reg->x2,r_reg->y2);
	   //Fl::check();
	   yuv blob_color;


 
	   //Average Color
	   if(doAverageColor){
	     //printf("Trying to get average color\n");
	     //printf("%d\n",vision.color[COLOR_RED].list->run_start);
	     

	     blob_color = AverageColor(vis_img_raw_robot, vision.rmap, r_reg->run_start);


	     // blob_color =  AverageColor(vis_img_raw_robot,
	     //				vision.rmap,
	     //				vision.color[COLOR_ORANGE].list->run_start);
	     //printf("Done!\n");
	     //	   printf("Average Color %d %d %d\n",
	     //  blob_color.y, blob_color.u, blob_color.v);
	     printf("Should this region be not blue? (0 or 1)\n");
	     scanf("%d",&isDog);
	 
	     printf("putting y: %d u: %d v: %d\n",blob_color.y,blob_color.u,blob_color.v);
	   
	     FILE *data;
	     data = fopen("regions_blue.data","a");
	     fprintf(data,"%d, %d, %d, %d\n",blob_color.y,blob_color.u,blob_color.v,isDog);    
	     fclose(data);
	     
	   }
	 }
       //End of region calculations
       
       if(doRobotDist){
	 int w;
	 int h;
	 int robot_x2, robot_x1, robot_y2, robot_y1;
	 int robot_w;
	 int robot_h;
	 int robot_a;
	 double dist;
	 vector3d robot_dir;
       
	 //w = xmax - xmin + 1;
	 //h = ymax - ymin + 1;
	 //cen_y = ysum/n;	    
	 //cen_x = xsum/n;
	 
	 // &obj_info.red_robots[0];
	 robot_x2 = (int)obj_info.red_robots[0].left;
	 robot_x1 = (int)obj_info.red_robots[0].right;
	 robot_y2 = (int)obj_info.red_robots[0].confidence;
	 robot_y1 = (int)obj_info.red_robots[0].orientation;
	 
 	 if(xmax > robot_x2 && xmax != 0){
	   printf("robot_x2: %d xmax: %d\n",robot_x2,xmax);
	   robot_x2 = xmax;
	 }
 	 if((xmin < robot_x1 && xmin !=0) || (robot_x1 == 0)){
	   printf("robot_x1: %d xmin %d\n",robot_x1,xmin);
 	   robot_x1 = xmin;
	 }
 	 if(ymax > robot_y2 && ymax !=0){
	   printf("robot_y2: %d ymax: %d\n",robot_y2,ymax);
 	   robot_y2 = ymax;
	 }
 	 if((ymin < robot_y1 && ymin !=0) || (robot_y1 == 0)){
	   printf("robot_y1: %d ymin: %d\n",robot_y1,ymin);
 	   robot_y1 = ymin;
	 }
        
	 
	  rle_widget->setBBox((int)robot_x1,(int)robot_y1,(int)robot_x2,(int)robot_y2);
	 Fl::check();

	 robot_w = robot_x2 - robot_x1 +1;
	 robot_h = robot_y2 - robot_y1 +1;
	 //robot_w = obj_info.red_robots[0].left;
	 //robot_h = obj_info.red_robots[0].right;
	 
	 
	 robot_a = (int)obj_info.red_robots[0].distance;
	 //robot_a = 0;
	 // printf("robot_dist: %f\n",obj_info.red_robots[0].distance);
	 printf("robot_w: %d robot_h: %d, robot_a: %d\n",robot_w,robot_h, robot_a);
	 a += robot_a;
	 //rle_widget->setBBox(r_reg->x1,r_reg->y1,r_reg->x2,r_reg->y2);

	 // robot_dir = vision.getPixelDirection(cen_x, cen_y);
	 //printf("X: %f, Y: %f, Z: %f\n", robot_dir.x, robot_dir.y, robot_dir.z);
	 //printf("dtree_distance: %f\n",obj_info.red_robots[0].distance);
	 //DTREE
	 
	 printf("a: %d robot_h: %d robot_w: %d\n",a,robot_h,robot_w);
	 if(a <= 59){
	   if(a <= 28){
	     if(robot_h <= 1){
	       dist = 0.0;
	     }
	     else{
	       dist = 2.0;
	     }
	   }
	   else{
	     if(robot_w <= 10){
	       dist = 2.0;
	     }
	     else{
	       dist = 1.5;
	     }
	   }
	 }
	 else{
	   if(robot_h <= 20){
	     dist = 1.0;
	   }
	   else{
	     if(a > 155){
	       dist = 0.5;
	     }
	     else{
	       if(robot_w <= 37){
		 dist = 0.5;
	       }
	       else{
		 dist = 1.0;
	       }
	     }
	   }
	 }
	 printf("dtree dist: %f\n",dist);
	 //END DTREE
	 printf("robot_a = %d a= %d\n",robot_a,a);
	 // printf("robot_w: %f robot_h: %f  w/h: %f\n",robot_w,robot_h,(robot_w/robot_h));
	 //printf("%f, %f, %f, %f, %f\n", robot_w, robot_h, robot_a, (robot_w/robot_h), d);
	 printf("%d, %d, %d, %f, %f\n", robot_w, robot_h, a, ((float)robot_w/(float)robot_h), d);
	 printf("How far is the dog? (increments of 0.5, up to 3.0)\n");
	 scanf("%f",&d);

	 if(a!=0){
	   FILE *data;
	   data = fopen("red_robots_dist.data","a");
	   fprintf(data,"%d, %d, %d, %f, %f\n", robot_w, robot_h, a, ((float)robot_w/(float)robot_h), d);
	 }
	 //if(a!=0){
	 //FILE *data;
	 //data = fopen("robots_dist.txt","a");
	 //fprintf(data,"%d, %d, %f, %f\n",w,h,cen_y, d);
	 //}
       }  
    }
    //end of blobstats

    if (0==Fl::check()) {
      run=false;
    }
  }

  vision.close();

  if(SortMode) {
    for(int sort_idx=0; sort_idx<NumSortClasses; sort_idx++) {
      fclose(sort_files[sort_idx]);
    }
  }

  return(0);
}