/* 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 <stdio.h>

#include <iostream>

using namespace std;

#include "../../agent/Vision/colors.h"
#include "../../agent/Vision/cmv_types.h"
#include "../../agent/Vision/Vision.h"

#include "../log_processing/shared/LogReader.h"
#include "../log_processing/shared/PacketDecoder.h"
#include "../log_processing/shared/VisionRawPacketDecoder.h"

char* model="ERS7";

class SImage;
SImage *DebugImage=NULL;

class RedVOrangeCalibrator {
private:
  static const int ers210_ImageXSize = 176;
  static const int ers210_ImageYSize = 144;

  static const int MaxPacketLength=120*1024;

  /* We need to delay initialization until we determine if we're
     reading old, 210 packets which don't contain an image size or
     the new format of packets which do include the image size (and
     come in different sizes). If it's the former, we init with good
     ol' 176x144. In the later case, we use what's in the first packet.
  */
  bool initialized;

  int width, height;

  RobotDataPacket packet;
  VisionRawPacketDecoder vis_raw_decoder;
  // need non const version to write to
  CMVision::image_yuv<uchar> img;
  Vision vision;

  yuv_generic<ulong> sum_color;
  int example_count;

  void get_vision_info(double *angles);

  yuv red;
  yuv orange;

public:
  LogReader log_reader;

  RedVOrangeCalibrator();
  ~RedVOrangeCalibrator();

  void init(int _width, int _height);

  void run_vision();
  void get_red_data(const char *filename);
  void gen_calibration();
};

RedVOrangeCalibrator::RedVOrangeCalibrator()
{
  // We have not called init yet.
  initialized = false;

  // packet.data = NULL;
  packet.data = new uchar[MaxPacketLength];
    

  red.y = 0;
  red.u = 0;
  red.v = 0;

  orange.y = 0;
  orange.u = 0;
  orange.v = 0;
}

RedVOrangeCalibrator::~RedVOrangeCalibrator()
{
  delete[] packet.data;
  packet.data = NULL;
}

void RedVOrangeCalibrator::init(int _width, int _height) {
  
  if(!initialized) {
    initialized = true;
    
    width = _width;
    height = _height;
    
    if(width==176)
      vision = Vision(false);
    else
      vision = Vision(true);

    img.row_stride = 3*width;
    img.width  = width;
    img.height = height;
    img.buf_y = new uchar[width*height*3];
    img.buf_u = img.buf_y+  width;
    img.buf_v = img.buf_y+2*width;
    
    vision.initialize("vision.cfg", width, height);
  }
}

void RedVOrangeCalibrator::get_vision_info(double *angles)
{
  static int image_num=-1;

  image_num++;

  ObjectInfo obj_info;
  VisionObjectInfo extra_obj_info;

  vision.setHeadAngles(angles+2);
  vision.setBodyParams(angles[0],angles[1]);

  vision.runLowLevelVision(img);
  vision.runHighLevelVision(&obj_info,img);

  //printf("ball conf %g, dist %g\n",obj_info.ball.confidence,obj_info.ball.distance);

  if(obj_info.ball.confidence > .4){
    yuv avg_color;

    if(obj_info.ball.distance < 1500.0){
      printf("skipping image %d because ball is too close (%g mm)\n",image_num,obj_info.ball.distance);
      return;
    }

    extra_obj_info = vision.vobj_info[VisionInterface::BALL];
    if(extra_obj_info.reg->area >= 128){
      printf("skipping image %d because ball is too big\n",image_num);
      return;
    }

    CMVision::color_class_state *red_state;
    CMVision::region *red_reg;
    red_state = &vision.color[COLOR_RED];
    red_reg = red_state->list;
    if(red_reg!=NULL && red_reg->area >= 64){
      printf("skipping image %d because too much red\n",image_num);
      return;
    }

    avg_color = AverageColor(img,vision.rmap,extra_obj_info.reg->run_start);

    sum_color.y += avg_color.y;
    sum_color.u += avg_color.u;
    sum_color.v += avg_color.v;
    example_count++;

    printf("using image %d, avg color (%d,%d,%d)\n",image_num,avg_color.y,avg_color.u,avg_color.v);
  }else{
    printf("skipping image %d because ball confidence of %g is too low\n",image_num,obj_info.ball.confidence);
  }
}

void RedVOrangeCalibrator::run_vision()
{
  sum_color.y = 0;
  sum_color.u = 0;
  sum_color.v = 0;
  example_count = 0;

  while(log_reader.readNextPacket(&packet,MaxPacketLength)){
    switch(packet.dataType){
      case PacketDecoder::VisionRawLead:
        {
	  init(ers210_ImageXSize, ers210_ImageYSize);

          //printf("raw vision packet\n");
          double angles[5];
          vis_raw_decoder.decode(angles,&packet);
          vis_raw_decoder.decode(ers210_ImageXSize,ers210_ImageYSize,&img,&packet);

          get_vision_info(angles);
          break;
        }
      case PacketDecoder::VisionRaw2Lead:
        {
          //printf("raw vision packet\n");
          double angles[5];
          vis_raw_decoder.decode(angles,&packet);
	  vis_raw_decoder.decode(width, height, &packet);

	  init(width, height);

	  int w, h;
          vis_raw_decoder.decode(w, h, &img, width, height, &packet);
	  if(w!=width || h!=height) {
	    printf("Ummm... You didn't read the image size that you were expecting\n");
	  }
	  
          get_vision_info(angles);
          break;
        }
    }
  }

  yuv avg_color;
  avg_color.y = (uchar)((double)sum_color.y/example_count);
  avg_color.u = (uchar)((double)sum_color.u/example_count);
  avg_color.v = (uchar)((double)sum_color.v/example_count);

  //printf("avg color is (%d,%d,%d)\n",avg_color.y,avg_color.u,avg_color.v);
  orange = avg_color;
}

void RedVOrangeCalibrator::get_red_data(const char *filename)
{
  FILE *red_file;

  red_file = fopen(filename,"rb");
  if(red_file==NULL)
    return;

  fread(&red.y,sizeof(red.y),1,red_file);
  fread(&red.u,sizeof(red.y),1,red_file);
  fread(&red.v,sizeof(red.y),1,red_file);

  fclose(red_file);
}

//static double NormalizedLocation(vector3d loc,vector3d dir,double base,double scale){
//  double val;
//
//  val = dir.dot(loc);
//  val -= base;
//  val *= scale;
//
//  return val;
//}

void RedVOrangeCalibrator::gen_calibration()
{
  if(red.y==0 || orange.y==0){
    printf("missing data needed for calibration, no calibration generated\n");
    return;
  }

  printf("red    is: (%3d,%3d,%3d)\n",red.y,red.u,red.v);
  printf("orange is: (%3d,%3d,%3d)\n",orange.y,orange.u,orange.v);

  vector3d red_v;
  vector3d orange_v;

  red_v.x = red.y;
  red_v.y = red.u;
  red_v.z = red.v;

  orange_v.x = orange.y;
  orange_v.y = orange.u;
  orange_v.z = orange.v;

  vector3d orange_to_red;
  vector3d thresh_dir;
  double thresh_base;
  double thresh_scale;

  orange_to_red = red_v - orange_v;
  thresh_dir = orange_to_red.norm();
  thresh_base = thresh_dir.dot(orange_v);
  thresh_scale = 1.0 / (thresh_dir.dot(red_v)-thresh_base);

  printf("(%g,%g,%g) base %g scale %g\n",V3COMP(thresh_dir),thresh_base,thresh_scale);
  //printf("orange %g red %g\n",
  //       NormalizedLocation(orange_v,thresh_dir,thresh_base,thresh_scale),
  //       NormalizedLocation(red_v,   thresh_dir,thresh_base,thresh_scale));
  //printf("%g %g %g\n",
  //       NormalizedLocation(orange_v-orange_to_red*.5,thresh_dir,thresh_base,thresh_scale),
  //       NormalizedLocation(orange_v+orange_to_red*.5,thresh_dir,thresh_base,thresh_scale),
  //       NormalizedLocation(orange_v+orange_to_red*1.5,thresh_dir,thresh_base,thresh_scale));

  FILE *redvor_file;
  redvor_file = fopen("red_v_or.prm","wb");
  if(redvor_file==NULL){
    printf("couldn't create file 'redvor.prm', skipping\n");
    return;
  }
  fwrite(&thresh_dir,  sizeof(thresh_dir  ),1,redvor_file);
  fwrite(&thresh_base, sizeof(thresh_base ),1,redvor_file);
  fwrite(&thresh_scale,sizeof(thresh_scale),1,redvor_file);
  fclose(redvor_file);
}

int main(int argc, char *argv[])
{
  if(argc < 3) {
    fprintf(stderr,"usage:vis_val <log file> <red.prm>\n");
    exit(1);
  }

  RedVOrangeCalibrator calib;
  FILE *log_file;
  char *log_filename=NULL;
  
  // We do a just-in-time init once we know what type of
  // log we're dealing with. (repeated calls beyond the
  // first do not do anything - we set a flag the first
  // time through.
  //  calib.init();

  log_filename=argv[1];

  printf("log name='%s'\n",log_filename);

  log_file = fopen(log_filename,"r");
  if(log_file == NULL){
    cerr << "Couldn't open log file '" << log_filename << "'" << endl;
    exit(2);
  }
  calib.log_reader.setSource(log_file);

  calib.run_vision();
  
  calib.get_red_data(argv[2]);

  calib.gen_calibration();

  fclose(log_file);
}