import java.awt.*;
import java.awt.image.*;
import javax.swing.*;
import java.awt.event.*;
import java.lang.*;
//rem this line and below to run without output if dont have JAI
import javax.media.jai.JAI;

public class Diffuse  
{
	BufferedImage myImage; 			//image that is to be manipulated
	WritableRaster destinationRaster;	
	BufferedImage finalImage;		//final processed image
	Raster theImage;
	
	int kvalue;
    int numberOfBands=3;
    
    
	public Diffuse (String filename,int kvalue)
	{
	this.kvalue=kvalue;
	Image tempimage=(new ImageIcon(filename)).getImage();
	myImage=new BufferedImage(tempimage.getWidth(null),tempimage.getHeight(null),BufferedImage.TYPE_3BYTE_BGR);
	finalImage=new BufferedImage(tempimage.getWidth(null),tempimage.getHeight(null),BufferedImage.TYPE_3BYTE_BGR);
 	Graphics g=(myImage.createGraphics()); 
    	g.drawImage(tempimage,0,0, null);  //draw source image into bufferedimage for source
	theImage=myImage.getData();
	}


	public Image getFinalImage(){return finalImage;} 

	public Image getMyImage(){return myImage;}


	// /*  does step diffusion steps taking theImage as input and leaving output in finalImage on exit*/
	public void diffusionStep(int step)
	{
    	destinationRaster=theImage.createCompatibleWritableRaster(); 
	WritableRaster tempRaster=theImage.createCompatibleWritableRaster();

        //arrays for target and 4 nearest neighbor pixel values
  	float[] ij=new float[numberOfBands];
	float[] n=new float[numberOfBands];
	float[] s=new float[numberOfBands];
	float[] e=new float[numberOfBands];
	float[] w=new float[numberOfBands];
	float[] destination=new float[numberOfBands];
	float cCoefn;
	float cCoefs;
	float cCoefe;
	float cCoefw;
	float deltan;
	float deltas;
	float deltae;
	float deltaw;

	double localgradient;

	//do step number of diffusion passes on each call 
	for (int iteration=1;iteration<=step;iteration++){

		//iterate over  rows of image
		for (int i=0; i<theImage.getWidth(); i++){
			
			//iterate over a collumn of image
    			for (int j=0; j<theImage.getHeight(); j++){
			    
				ij=theImage.getPixel(i,j,ij);
				if (j==0)
					n=theImage.getPixel(i,j,n);
				else
					n=theImage.getPixel(i,j-1,n);
				if (j==theImage.getHeight()-1)
					s=theImage.getPixel(i,j,s);
				else
					s=theImage.getPixel(i,j+1,s);
				if (i==theImage.getWidth()-1)
					e=theImage.getPixel(i,j,e);
				else
					e=theImage.getPixel(i+1,j,e);
				if (i==0)
					w=theImage.getPixel(i,j,w);
				else
					w=theImage.getPixel(i-1,j,w);

				//perform op on each band sperately
				for ( int b=0;b<numberOfBands;b++){
				    deltan=n[b]-ij[b];
				    deltas=s[b]-ij[b];
				    deltae=e[b]-ij[b];
				    deltaw=w[b]-ij[b];
				    cCoefn=(float)(1/(1+Math.pow(deltan/kvalue,2)));
				    cCoefs=(float)(1/(1+Math.pow(deltas/kvalue,2)));
				    cCoefe=(float)(1/(1+Math.pow(deltae/kvalue,2)));
				    cCoefw=(float)(1/(1+Math.pow(deltaw/kvalue,2)));
				    destination[b]=ij[b]+(cCoefn*deltan+cCoefs*deltas+cCoefe*deltae+cCoefw*deltaw)/4;
				    if (destination[b]<0) destination[b]=0; if (destination[b]>255) destination[b]=255;
				}
				
    				destinationRaster.setPixel(i,j,destination);
				for (int b=0;b<numberOfBands;b++)
				    destination[b]=0;
			     

    			}//end height for
		}//end width for

		//set our source raster to product of past iteration 
		theImage=destinationRaster;

		//set destination for next pass to temp unless were on the last pass
		if (iteration<step-1)
			destinationRaster=tempRaster;

	}//end for for number of passes
	finalImage=new BufferedImage(destinationRaster.getWidth(),destinationRaster.getHeight(),BufferedImage.TYPE_3BYTE_BGR);
	finalImage.setData(destinationRaster);
	}//end function


	public static void main(String[] args)
	{

	if (args.length<3) {
		System.out.println("Usage: java Diffuse <imagefilename> <step size> <k value> <optional # step size iterations>");
		System.exit(-1);}
    	Diffuse mydif=new Diffuse(args[0],Integer.parseInt(args[2]));

	int number;
	if (args.length>3) number=Integer.parseInt( args[3]);
	else number=10;
	
	BufferedImage im;
	JFrame mainframe;
	ImageIcon desticon;
	for (int x=0;x<number;x++){
	     mydif.diffusionStep (Integer.parseInt(args[1]));
		
		        mainframe=new JFrame("diffusion test");
			mainframe.addWindowListener(new WindowAdapter()
				    {
					public void windowClosing(WindowEvent e)
					    {
						System.exit(0);
					    }
				    }
				    );
			im=(BufferedImage)mydif.getFinalImage();
			JAI.create("filestore", im, (new Integer(x)).toString()+".jpg", "JPEG", null);
			desticon=new ImageIcon(im);
			mainframe.setSize(desticon.getIconWidth(),desticon.getIconHeight());
			mainframe.getContentPane().add(new JLabel(desticon));
			mainframe.setVisible(true);
	
	}//end for 
    	}

}//end class

