import edu.gwu.lintool.*;
import java.util.*;
import java.text.*;
import java.awt.*;
import java.awt.event.*;
import java.awt.geom.*;
import javax.swing.*;
import javax.swing.event.*;

public class DemoDFT2D extends JFrame {

    // N x N grid or image.
    int N = 5;

    JTextField vField = new JTextField (6);

    // The signal and spectrum are now 2D images or 2D arrays of complex numbers.
    ComplexLocal[][] signal, spectrum;

    // We're going to display the real and imaginary parts separately.
    Clickable2DGrid signalRe, signalIm, spectrumRe, spectrumIm;

    // For selection and change of value.
    Clickable2DGrid currentGrid;
    int selecti = -1, selectj = -1;
        

    public DemoDFT2D ()
    {
	// Build the GUI.
        this.setSize (800, 700);
        this.setTitle ("FFT in 2D");
        Container cPane = this.getContentPane();
        cPane.add (makeBottomPanel(), BorderLayout.SOUTH);
        cPane.add (makeCenterPanel(), BorderLayout.CENTER);
        this.setVisible (true);
    }
    

    // Apply() scales the values to the range [-127, 127] so that we can treat the
    // values as greyscale intensities.

    void apply ()
    {
        apply ("Signal", signal, signalRe, signalIm);
        apply ("Spectrum", spectrum, spectrumRe, spectrumIm);
        signalRe.repaint ();
        signalIm.repaint ();
        spectrumRe.repaint ();
        spectrumIm.repaint ();
    }
    
    void apply (String msg, ComplexLocal[][] data, Clickable2DGrid realGrid, Clickable2DGrid imGrid)
    {
        for (int i=0; i<N; i++) {
            for (int j=0; j<N; j++) {
                double re = data[i][j].re;
                double im = data[i][j].im;
                if (re < -127) {
                    re = -127;
                }
                if (re > 127) {
                    re = 127;
                }
                if (im < -127) {
                    im = -127;
                }
                if (im > 127) {
                    im = 127;
                }
                realGrid.grid[i][j] = (int) re;
                imGrid.grid[i][j] = (int) im;
            }
        }

    }
    

    void change ()
    {
        try {
            double v = Double.parseDouble (vField.getText());
            if (currentGrid == signalRe) {
                signal[selecti][selectj].re = v;
            }
            else if (currentGrid == signalIm) {
                signal[selecti][selectj].im = v;
            }
            else if (currentGrid == spectrumRe) {
                spectrum[selecti][selectj].re = v;
            }
            else if (currentGrid == spectrumIm) {
                spectrum[selecti][selectj].im = v;
            }
            apply ();
        }
        catch (NumberFormatException e) {
            System.out.println (e);
        }
    }


    void setTextField ()
    {
        double v = 0;
        if (currentGrid == signalRe) {
            v = signal[selecti][selectj].re;
        }
        else if (currentGrid == signalIm) {
            v = signal[selecti][selectj].im;
        }
        else if (currentGrid == spectrumRe) {
            v = spectrum[selecti][selectj].re;
        }
        else if (currentGrid == spectrumIm) {
            v = spectrum[selecti][selectj].im;
        }
        vField.setText ("" + v);
    }

    void showValues ()
    {
	signalRe.showValuesToggle ();
	signalIm.showValuesToggle ();
	spectrumRe.showValuesToggle ();
	spectrumIm.showValuesToggle ();
	apply ();
    }


    void clear ()
    {
        for (int i=0; i<N; i++) {
            for (int j=0; j<N; j++) {
		signal[i][j].re = signal[i][j].im = 0;
		spectrum[i][j].re = spectrum[i][j].im = 0;
	    }
	}
	apply ();
    }




    ////////////////////////////////////////////////////////////////////////////
    // DFT and inverse.

    void DFT ()
    {
	spectrum = new ComplexLocal [N][N];
	for (int u=0; u<N; u++) {
	    for (int v=0; v<N; v++) {

		ComplexLocal sum = new ComplexLocal (0,0);

		for (int x=0; x<N; x++) {
		    for (int y=0; y<N; y++) {
			double cosTerm = Math.cos (2*Math.PI*(u*x+v*y) / (double)N);
			double sinTerm = -Math.sin (2*Math.PI*(u*x+v*y) / (double)N);
			ComplexLocal cTerm = new ComplexLocal (cosTerm, sinTerm);
			ComplexLocal prod = signal[x][y].mult (cTerm);
			sum = sum.add (prod);
		    }
		}

		ComplexLocal term = sum.mult (1.0/(N));
		spectrum[u][v] = term;
	    }
	}

	apply ();
    }
    

    void inverseDFT () 
    {
	signal = new ComplexLocal [N][N];

	for (int x=0; x<N; x++) {
	    for (int y=0; y<N; y++) {

		ComplexLocal sum = new ComplexLocal (0,0);

		for (int u=0; u<N; u++) {
		    for (int v=0; v<N; v++) {
			double cosTerm = Math.cos (2*Math.PI*(u*x+v*y) / (double)N);
			double sinTerm = Math.sin (2*Math.PI*(u*x+v*y) / (double)N);
			ComplexLocal cTerm = new ComplexLocal (cosTerm, sinTerm);
			ComplexLocal prod = spectrum[u][v].mult (cTerm);
			sum = sum.add (prod);
		    }
		}

		ComplexLocal term = sum.mult (1.0/(N));
		signal[x][y] = term;
	    }
	}

	apply ();
    }
    

    ////////////////////////////////////////////////////////////////////////////
    // GUI construction.

    JPanel makeBottomPanel ()
    {
        JPanel panel = new JPanel ();
        
        panel.add (new JLabel ("Value: "));
        panel.add (vField);

	JButton changeB = new JButton ("Change");
	changeB.addActionListener (
	   new ActionListener () {
		   public void actionPerformed (ActionEvent a)
		   {
		       change ();
		   }
           }
        );
	panel.add (changeB);

        panel.add (new JLabel ("       "));
        

	JButton invB = new JButton ("invDFT");
	invB.addActionListener (
	   new ActionListener () {
		   public void actionPerformed (ActionEvent a)
		   {
		       inverseDFT ();
		   }
           }
        );
	panel.add (invB);

	JButton dftB = new JButton ("DFT");
	dftB.addActionListener (
	   new ActionListener () {
		   public void actionPerformed (ActionEvent a)
		   {
		       DFT ();
		   }
           }
        );
	panel.add (dftB);


	JButton showB = new JButton ("ShowValues");
	showB.addActionListener (
	   new ActionListener () {
		   public void actionPerformed (ActionEvent a)
		   {
		       showValues ();
		   }
           }
        );
	panel.add (showB);

        panel.add (new JLabel ("  "));
	JButton clearB = new JButton ("Clear");
	clearB.addActionListener (
	   new ActionListener () {
		   public void actionPerformed (ActionEvent a)
		   {
		       clear ();
		   }
           }
        );
	panel.add (clearB);

        panel.add (new JLabel ("  "));
	JButton quitB = new JButton ("Quit");
	quitB.addActionListener (
	   new ActionListener () {
		   public void actionPerformed (ActionEvent a)
		   {
		       System.exit(0);
		   }
           }
        );
	panel.add (quitB);
        
        return panel;
    }
    

    JPanel makeCenterPanel ()
    {
        JPanel panel = new JPanel ();
        panel.setLayout (new GridLayout (2,2));
        //panel.setLayout (new BorderLayout());
        
        signal = new ComplexLocal [N][N];
        spectrum = new ComplexLocal [N][N];
        for (int i=0; i<N; i++) {
            for (int j=0; j<N; j++) {
                signal[i][j] = new ComplexLocal (0,0);
                spectrum[i][j] = new ComplexLocal (0,0);
            }
        }

        signalRe = new Clickable2DGrid (this, N, N);
        signalRe.setBorder (BorderFactory.createLineBorder (Color.BLACK));
        signalIm = new Clickable2DGrid (this, N, N);
        signalIm.setBorder (BorderFactory.createLineBorder (Color.BLACK));
        spectrumRe = new Clickable2DGrid (this, N, N);
        spectrumRe.setBorder (BorderFactory.createLineBorder (Color.BLACK));
        spectrumIm = new Clickable2DGrid (this, N, N);
        spectrumIm.setBorder (BorderFactory.createLineBorder (Color.BLACK));

        panel.add (signalRe);
        panel.add (signalIm);
        panel.add (spectrumRe);
        panel.add (spectrumIm);
        return panel;
    }
    


    public static void main (String[] argv)
    {
        new DemoDFT2D ();
    }
    

}

class ComplexLocal {

    public double re, im;
    public double angle;

    public ComplexLocal ()
    {
    }

    public ComplexLocal (double re, double im)
    {
        this.re = re;
        this.im = im;
	angle = Math.atan2 (im, re);
	if (angle < 0) {
	    angle += Math.PI;
	}
    }

    public static ComplexLocal makeInstance (double re, double im)
    {
	return new ComplexLocal (re, im);
    }

    public double magnitude ()
    {
	return Math.sqrt (re*re + im*im);
    }

    public double angle ()
    {
	angle = Math.atan2 (im, re);
	if (angle < 0) {
	    angle = 2*Math.PI + angle;
	}
	return angle;
    }

    public ComplexLocal add (ComplexLocal b)
    {
	ComplexLocal c = (ComplexLocal) b;
        return new ComplexLocal (re + c.re, im + c.im);
    }

    public ComplexLocal sub (ComplexLocal b)
    {
	ComplexLocal c = (ComplexLocal) b;
	return new ComplexLocal (re - c.re, im - c.im);
    }
    
    public ComplexLocal mult (ComplexLocal b)
    {
	ComplexLocal c = (ComplexLocal) b;
        return new ComplexLocal (re*c.re - im*c.im, re*c.im + im*c.re);
    }

    public ComplexLocal mult (double a)
    {
        return new ComplexLocal (re*a, im*a);
    }
    

    public ComplexLocal pow (int n)
    {
        ComplexLocal result = new ComplexLocal (re, im);
        for (int i=2; i<=n; i++) {
            result = (ComplexLocal) result.mult (this);
        }
        return result;
    }

    public ComplexLocal conjugate ()
    {
	return new ComplexLocal (re, -im);
    }

}