/* 

   TileTrain.java

   Vision training tool to be used with the Tekkotsu framework.

   Originally modified from ImageShow and VisionTrain

   Requires ImageData

   By Rob Salkin (salkin@cs.albany.edu) and Shawn Turner (st2750@albany.edu)

   University at Albany - May 2004

*/



import java.awt.image.*;

import java.awt.*;

import javax.swing.*;

import javax.swing.event.*;

import java.awt.event.*;

import java.io.*;

import java.util.*;

import java.awt.geom.*;



public class TileTrain extends Canvas implements KeyListener, MouseListener, MouseMotionListener{

    BufferedImage _image;

    String[] names;         //tracking color names

    byte[] tmap;            //threshold map

    int[] RGBdata;

    int[] YUVdata;

    ImageData imageDataYUV; //for YUV data

    ImageData imageDataRGB; //for RGB data

    int count;             //color counter



    int[] red;

    int[] green;

    int[] blue;



    //tiling layout

    //"prime" numbered layouts will run in a single row -- not very fun

    

    //dimensions                            W    H

    private static int [][]  sizetable = { {0  , 0},   //0 images  

					   {1  , 1},   //1

					   {2  , 1},   //2

					   {3  , 1},   //3

					   {2  , 2},   //4, etc.

					   {5  , 1},

					   {3  , 2},

					   {7  , 1},

					   {4  , 2},

					   {3  , 3},

					   {5  , 2},

					   {11 , 1},

					   {4  , 3},   //12

					   {13 , 1},

					   {7  , 2},

					   {5  , 3},

					   {4  , 4},   //16

					   {17 , 1},

					   {6  , 3},

					   {19 , 1},

					   {5  , 4},

					   {7  , 3},

					   {11 , 2},

					   {23 , 1},

					   {6  , 4},

					   {5  , 5}  };



    private static final int MAX_DIM = 25;             // maximum tiling area. Should match table



    //for storing YUV data

    public static final int size_y=16, size_u=64, size_v=64;

    public static final int size_total=size_y*size_u*size_v;





    //for marquee tool

    Area curArea;

    Polygon curPoly;

    int lastx, lasty;

    

    //for control

    public static final int MODIFIER_NONE=0;

    public static final int MODIFIER_SHIFT=1;

    public static final int MODIFIER_CTRL=2;

    public int curSelectModifier;

    public int curSelectModified;





    public static int WIDTH; // number of images wide

    public static int HEIGHT; // number of images high

    public static int TOTAL_IMGS; //total images



    // per tile. Set in main, based on argc

    public static int PIXEL_WIDTH;

    public static int PIXEL_HEIGHT;

    public static int TOTAL_PIXELS;



    // maximum number of allowable color names

    public static final int MAX_COLORS=8;



    //all of these should be set on the command line

    public static final int SMOOTH_FACTOR=2; //1 pixel used for (1 + (2*SMOOTH_FACTOR))^3 colors when same SMOOTH_FACTOR is used for Y,U,V

    public static final int Y_SMOOTH_FACTOR=3;//SMOOTH_FACTOR; //Y should have greatest smoothing for lighting variation

    public static final int U_SMOOTH_FACTOR=1;//SMOOTH_FACTOR; 

    public static final int V_SMOOTH_FACTOR=1;//SMOOTH_FACTOR; 

    

    

    public static void main(String args[]) {



	System.out.println

	    ("TileTrain is a vision training tool to be used with the Tekkotsu framework.\n"

	     +"Originally modified from ImageShow and VisionTrain\n"

	     +"Requires ImageData\n"

	     +"By Rob Salkin (salkin@cs.albany.edu) and Shawn Turner (st2750@albany.edu)\n"

	     +"University at Albany - May 2004\n");





	if(args.length-1 <= 0 || args.length-1 > MAX_DIM){

	    System.out.println("Number of images must be > 0 and <= "+MAX_DIM);

	    usageAndExit(args.length-1);

	}



	//determine proper image size for tiling by model

	if(args[0].equals("-is7")){

	    System.out.println("Processing images for ERS7");

	    PIXEL_WIDTH = 104;

	    PIXEL_HEIGHT = 80;

	}

	else if(args[0].equals("-is2xx")){

	    System.out.println("Processing images for ERS2xx");

	    PIXEL_WIDTH = 88;

	    PIXEL_HEIGHT = 72;

	}

	else{

	    usageAndExit(args.length);

	}



	



	TOTAL_PIXELS = PIXEL_WIDTH*PIXEL_HEIGHT;           //pixels per tile

	

	//set tile dimensions

	WIDTH = sizetable[args.length-1][0];

	HEIGHT = sizetable[args.length-1][1];

	TOTAL_IMGS = WIDTH*HEIGHT;

    

	String files[]=new String[args.length-1];

	for(int i=0; i<args.length-1; i++){

	    files[i]=args[i+1];

	}

	

	TileTrain tileTrain=new TileTrain(files);



	JFrame fraDisplay = new JFrame();

	fraDisplay.getContentPane().add(tileTrain);



	//ALWAYS PACK! OTHERWISE ARRAY OFFSET PROBLEMS HAPPEN!!!

	//	fraDisplay.setSize(WIDTH*PIXEL_WIDTH,HEIGHT*PIXEL_HEIGHT);

	fraDisplay.setResizable(false);

	fraDisplay.pack();

	fraDisplay.show();

		

	fraDisplay.addWindowListener(new WindowAdapter() {

		public void windowClosing(WindowEvent e) { System.exit(0); } });



    }

    

    

    public static void usageAndExit(int imgs) {

	System.out.println("usage: java TileTrain -is[7|2xx] raw_image (up to "+MAX_DIM+" raw images -- you gave me "+imgs+")");

	System.exit(1);

    }

    

    public int[] getRGBdata(){

	return RGBdata;

    }



    public int[] getYUVdata(){

	return YUVdata;

    }



    public TileTrain (String args[]) {

	ImageData imageDataRGB=new ImageData();

	ImageData imageDataYUV=new ImageData();



	tmap=new byte[size_total];



	red=new int[MAX_COLORS];

	green=new int[MAX_COLORS];

	blue=new int[MAX_COLORS];



	ImageData tempRGB;

	ImageData tempYUV;



	int iaRGB[];

	int iaYUV[];

	RGBdata=new int[TOTAL_IMGS*TOTAL_PIXELS];

	YUVdata=new int[TOTAL_IMGS*TOTAL_PIXELS];

	count=0;

	try{

	    for(int k=0;k<TOTAL_IMGS;k++){

		

		tempRGB=new ImageData();

		tempYUV=new ImageData();

		System.out.println("Loading " + args[k] );

		tempRGB.loadYUVFileAsRGB(args[k]);

		tempYUV.loadYUVFileAsYUV(args[k]);

		iaYUV=tempYUV.getPixels();

		iaRGB=tempRGB.getPixels();

		

		for(int h=0;h<PIXEL_HEIGHT;h++){

		    for(int w=0;w<PIXEL_WIDTH;w++){

			RGBdata[(k/WIDTH)*(WIDTH*TOTAL_PIXELS) //tile row offset

				+(k%WIDTH)*(PIXEL_WIDTH)       //tile col offset 

				+h*WIDTH*PIXEL_WIDTH           //pixel row offset

				+w]                            //pixel col offset

			    =iaRGB[h*PIXEL_WIDTH+w]; //RGB data for display

			

			YUVdata[(k/WIDTH)*(WIDTH*TOTAL_PIXELS) //tile row offset

				+(k%WIDTH)*(PIXEL_WIDTH)       //tile col offset 

				+h*WIDTH*PIXEL_WIDTH           //pixel row offset

				+w]                            //pixel col offset

			    =iaYUV[h*PIXEL_WIDTH+w]; //YUV data for seg use

		    }

		}

	    }	    

	}

	catch(Exception e){

	    System.out.println(e.getClass() + " thrown. \nPlease verify that the arguments to TileTrain" + 

			       " are consistent with the robot model and image filenames before trying again.\n\n");

	    

	    e.printStackTrace();

	    System.exit(1);

	}

   

	setBackground(Color.black);

	setSize(PIXEL_WIDTH*WIDTH, PIXEL_HEIGHT*HEIGHT);

	_image=new BufferedImage(PIXEL_WIDTH*WIDTH, PIXEL_HEIGHT*HEIGHT,

				 BufferedImage.TYPE_INT_RGB);

	

	showImage(RGBdata, tmap, PIXEL_WIDTH*WIDTH, PIXEL_HEIGHT*HEIGHT);

	

	names=new String[MAX_COLORS];

	curArea = new Area();

	addKeyListener(this);

	addMouseListener(this);

	addMouseMotionListener(this);     

	

	showHelp();



    }



    void showImage(int[] data, byte[] tmap, int width, int height) {

	_image.getRaster().setDataElements(0,0,width,height,data);

	repaint();

    }



    public void paint(Graphics graphics) {

	update(graphics);

    }



    public void keyPressed(KeyEvent e) {

	

	//modified to carry the shift/control selection features from them image to the canvas.

	if (e.getKeyCode()==KeyEvent.VK_SHIFT) {

	    curSelectModifier=MODIFIER_SHIFT;

	} else if (e.getKeyCode()==KeyEvent.VK_CONTROL) {

	    curSelectModifier=MODIFIER_CTRL;

	}

    }

    

    public static void showHelp(){

	JOptionPane.showMessageDialog

	    (null,"Start by selecting as much of a target color as possible to be segmented.\n"+

	     "Press N (=Name) to store the current selection and enter a name for the selection.\n"+

	     "Press S (=Save) to save the stored selections to .tm and .col files.\n"+

	     "Press C (=Clear) to clear ALL current mapped colors\n"+

	     "Press Q (=Quit) to quit this program\n" +

	     "Tips:\n"+

	     "\t-Press N after each selection -- don't forget to do so before pressing S.\n"+

	     "\t-Hold SHIFT to select multiple regions or append to regions.\n"+

	     "\t-Hold CTRL to deselect from regions.\n"+

	     "\t-On the color chooser, select a representative color that will stand out compared to other colors.\n"+

	     "\t-Click anywhere in the window to deselect all regions -- useful after pressing N.\n"+

	     "\t-All .tm and .col files should have names <=8 characters.\n"+

	     "Press H (=Help) for this message.\n", "TileTrain Help", JOptionPane.OK_OPTION);

    }

    

    public void keyReleased(KeyEvent e) { 

	if (e.getKeyCode()==KeyEvent.VK_SHIFT

	    ||e.getKeyCode()==KeyEvent.VK_CONTROL) {

	    curSelectModifier=MODIFIER_NONE;

	}

	else if (e.getKeyCode()==KeyEvent.VK_H){

	    showHelp();

	}

	else if (e.getKeyCode()==KeyEvent.VK_N){

	    //verify we have colors left to create

	    if(checkcolors() == false)

		return;



	    String tname = null;//temporary string for color name

	    do{

		tname=JOptionPane.showInputDialog("Enter name for color "+(count+1));

	    }while(tname!=null&&tname.trim().equals(""));

	    if(tname!=null){

		names[count] = tname;

		

		Color segcol = null;//segmentation will be seen as this color

		do{

		    segcol = JColorChooser.showDialog(null, "Segmented Color", Color.WHITE);

		}while(segcol == null);

		

		red[count]=segcol.getRed();

		green[count]=segcol.getGreen();

		blue[count]=segcol.getBlue();

		

		count++;

		doCalc(tmap, size_total);

		

		System.out.println(tname + " stored");

		

		//clear the area

		curArea.reset();

		repaint();

		

		checkcolors();

	    }

	    else{

		JOptionPane.showMessageDialog(null, "The current selection was not stored!");

	    }

	}	

	else if (e.getKeyCode()==KeyEvent.VK_S){



	    JFileChooser chooser=new JFileChooser();

	    

	    chooser.setCurrentDirectory(new File(System.getProperty("user.dir")));

	    chooser.setSelectedFile(new File("default.tm"));

	    int returnval=chooser.showSaveDialog(this);

	    if (returnval==JFileChooser.APPROVE_OPTION) {

		//bug fix for save path

		save(chooser.getCurrentDirectory() + "/" + chooser.getSelectedFile().getName());   

	    }   

	}

	else if (e.getKeyCode()==KeyEvent.VK_Q){

	    if( (JOptionPane.showConfirmDialog(null, "Are you Sure?\nAll unsaved work will be lost!", "Quit?", 

					       JOptionPane.OK_CANCEL_OPTION)) == JOptionPane.OK_OPTION)

		System.exit(0);

	}

	else if (e.getKeyCode()==KeyEvent.VK_C){

	    if( (JOptionPane.showConfirmDialog(null, "Are you Sure?\nAll Colors Will be lost!", "Clear Colors?", JOptionPane.OK_CANCEL_OPTION)) 

		== JOptionPane.OK_OPTION)

		count = 0;

	}

    

    }



    public void keyTyped(KeyEvent e) {



    }



    public boolean checkcolors(){

	//do we have any usuable colors left?



	if(count >= MAX_COLORS){

	    JOptionPane.showMessageDialog(null, "The maximum number of creatable colors ("+MAX_COLORS+

					  ") has been filled\nPlease save your work, or hit \"C\" to clear ALL colors");

	    return false;

	}

	return true;

    }

					  



    public void doCalc(byte[] tmap, int size){



	int size_total=size;



	int i,x,y,yuv,y_val,u_val,v_val,index;



	int rgb,r_val,g_val,b_val;





	if(count<=1){

	    for(i=0; i<size_total; i++){

		tmap[i]=(byte)0;

	    }

	}



	//scan the whole tiled image

	for(y=0;y<HEIGHT*PIXEL_HEIGHT;y++){

	    for(x=0;x<WIDTH*PIXEL_WIDTH;x++){

		//if the pixel is in the area selected

		if(curArea.contains(x,y)){	//could it be done without .contains?

		    //get the corresponding YUV pixel

		    yuv=YUVdata[y*WIDTH*PIXEL_WIDTH+x];



		    y_val=(yuv>>16)&0xff;

		    u_val=(yuv>>8)&0xff;

		    v_val=yuv&0xff;

		    y_val=y_val>>4;

		    u_val=u_val>>2;

		    v_val=v_val>>2;



		    //save the selected pixel color, or set it to conflict if previously set

		    int y1,u1,v1;

		    for(y1=y_val-Y_SMOOTH_FACTOR;y1<=y_val+Y_SMOOTH_FACTOR;y1++){

			for(u1=u_val-U_SMOOTH_FACTOR;u1<=u_val+U_SMOOTH_FACTOR;u1++){

			    for(v1=v_val-V_SMOOTH_FACTOR;v1<=v_val+V_SMOOTH_FACTOR;v1++){

				index=(y1*size_u+u1)*size_v+v1;

				

				if(index>=0&&index<=size_total-1){

				    if(tmap[index]==(byte)0){

					tmap[index]=(byte)count;

				    }

				    else if(tmap[index]!=(byte)count){

					tmap[index]=(byte)9;

				    }

				}

			    }

			}

		    }

		}

	    }

	}



    }



    public void save(String filename){

	int dotpos=filename.lastIndexOf('.');

	if (dotpos>0) filename=filename.substring(0,dotpos);

    

	try {

	    FileOutputStream file_tm_fos=new FileOutputStream(filename + ".tm");

	    OutputStreamWriter file_tm_osw=new OutputStreamWriter(file_tm_fos);

	    file_tm_osw.write("TMAP\nYUV8\n" +

			      size_y + " " + size_u + " " + size_v + "\n");

	    file_tm_osw.flush();

	    file_tm_fos.write(tmap);

	    file_tm_osw.close();

	    System.out.println(filename + ".tm saved");

	} catch (Exception ex) {

	    System.out.println("Error saving to "+filename +".tm: " + ex);

	}



	try {

	    FileWriter file_col_fw=new FileWriter(filename + ".col");



	    //grey

	    file_col_fw.write("0 (128 128 128) \"unclassified\" 8 1.00\n");



	    //selected colors

	    for(int colors=0;colors<count;colors++){

		//are 8 and 0.75 the best values?

		file_col_fw.write((1+colors)+" ("+red[colors]+" "+green[colors]+" "+blue[colors]+") \""+names[colors]+"\" 8 0.75\n");

	    }



	    //black

	    file_col_fw.write("9 (0 0 0) \"conflict\" 8 0.75\n");



	    file_col_fw.close();



	    System.out.println(filename + ".col saved");

	    

	} catch (Exception ex) {

	    System.out.println("Error saving to "+filename + ".col: " + ex);

	}







    }



    //ported mouse functions from TrainCanvas 

    public void mousePressed(MouseEvent e) {



	if (curArea==null) return;

      

	if (e.getButton()!=MouseEvent.BUTTON1) return;

      

	curSelectModified=curSelectModifier;

      

	if (curSelectModified!=MODIFIER_SHIFT

	    && curSelectModified!=MODIFIER_CTRL)

	    curArea.reset();

      

	curPoly=new Polygon();

      

	//standard

	lastx=e.getX();

	lasty=e.getY();

	curPoly.addPoint(e.getX(), e.getY());

      

    

	repaint();





    }



    public void mouseDragged(MouseEvent e) {

	if (curArea==null) return;

	try{

	    int x=e.getX();

	    int y=e.getY();

	    if ((Math.abs(x-lastx)+Math.abs(y-lasty))>2) {

		curPoly.addPoint(e.getX(), e.getY());

		lastx=x;

		lasty=y;

		repaint();

	    }

	

	}

	catch(Exception ex){

	    System.out.println("TileTrain::mouseDragged handled exception");

	    System.out.println("Probably due to the mouse accidently dragging off the preview window");

	}

	repaint();

      

    }



    public void mouseReleased(MouseEvent e) {

	if (curArea==null) return;

	if (e.getButton()!=MouseEvent.BUTTON1) return;

	try{

	    curPoly.addPoint(e.getX(), e.getY());



  

	    if (curPoly.npoints>=3) {

		if (curSelectModified==MODIFIER_CTRL) curArea.subtract(new Area(curPoly));

		else curArea.add(new Area(curPoly));

	    }

	}

	catch(Exception ex){

	    System.out.println("TileTrain::mouseReleased handled exception");

	    System.out.println("Probably due to the mouse accidently dragging off the preview window");

	}



	curPoly=null;

	repaint();

    }



    //pointing happening here

    public void mouseMoved(MouseEvent e){

      

	//      repaint();

      

    }



    public void mouseClicked(MouseEvent e){

	

    }



    public void mouseExited(MouseEvent e){



    }



    public void mouseEntered(MouseEvent e){}

    

    public void update(Graphics g) {

	Graphics2D g2d=(Graphics2D) g;



	Dimension sz=getSize();

      

	//draw image

	if (_image!=null)

	    g2d.drawImage(_image, 0, 0, sz.width, sz.height, null);

    

	//line color

	g2d.setColor(Color.white);



	//add shapes

	if (curArea!=null)

	    g2d.draw(curArea);



	if (curPoly!=null)

	    g2d.draw(curPoly);

  

	//I hate that thing...so it's off.

	//draw "hint" square

	//g2d.drawRect(lastx-2, lasty-2, 5, 5);



    }



}