import java.awt.*;
import java.awt.event.*;
import JSci.awt.*;
import JSci.maths.*;

/**
* Wavefunction numerically solves the time-independent
* Schr&ouml;dinger equation.
* @author Mark Hale
* @version 1.0
*/
public final class Wavefunction extends Frame {
        private final int N=200;
        private LineGraph graph;
        private DefaultGraph2DModel model=new DefaultGraph2DModel();
        private Label kLabel=new Label("Harmonic coupling");
        private TextField kField=new TextField("0.0",5);
        private Label hLabel=new Label("Anharmonic coupling");
        private TextField hField=new TextField("0.0",5);
        private Button fnButton=new Button("Probability");
        private Label statusLabel=new Label("-",Label.CENTER);
        private Button incButton=new Button("+");
        private Button decButton=new Button("-");
        private boolean showProbability=false;
        /**
        * Hamiltonian.
        */
        private DoubleTridiagonalMatrix H;
        /**
        * Energy level.
        */
        private int level=0;
        /**
        * Harmonic coupling.
        */
        private double k=0.0;
        /**
        * Anharmonic coupling.
        */
        private double h=0.0;
        /**
        * Runs an instance of Wavefunction.
        */
        public static void main(String arg[]) {
                new Wavefunction();
        }
        /**
        * Constructs Wavefunction.
        */
        public Wavefunction() {
                super("Wavefunction");
// generate x-axis values
                float xAxis[]=new float[N];
                for(int i=0;i<N;i++)
                        xAxis[i]=i*2.0f/(N-1)-1.0f;
                model.setXAxis(xAxis);
                model.addSeries(xAxis);
                model.addSeries(xAxis);
                model.addSeries(xAxis);
                graph=new LineGraph(model);
// register listeners
                addWindowListener(new WindowAdapter() {
                        public void windowClosing(WindowEvent evt) {
                                dispose();
                                System.exit(0);
                        }
                });
                incButton.addActionListener(new ActionListener() {
                        public void actionPerformed(ActionEvent evt) {
                                if(level<N-1) {
                                        level++;
                                        displayWavefunction();
                                }
                        }
                });
                decButton.addActionListener(new ActionListener() {
                        public void actionPerformed(ActionEvent evt) {
                                if(level>0) {
                                        level--;
                                        displayWavefunction();
                                }
                        }
                });
                kField.addActionListener(new ActionListener() {
                        public void actionPerformed(ActionEvent evt) {
                                k=Double.valueOf(kField.getText()).doubleValue();
                                H=calculateHamiltonian();
                                displayWavefunction();
                        }
                });
                hField.addActionListener(new ActionListener() {
                        public void actionPerformed(ActionEvent evt) {
                                h=Double.valueOf(hField.getText()).doubleValue();
                                H=calculateHamiltonian();
                                displayWavefunction();
                        }
                });
                fnButton.addActionListener(new ActionListener() {
                        public void actionPerformed(ActionEvent evt) {
                                if(showProbability) {
                                        showProbability=false;
                                        fnButton.setLabel("Probability");
                                        model.setSeriesVisible(0,true);
                                        model.setSeriesVisible(1,false);
                                } else {
                                        showProbability=true;
                                        fnButton.setLabel("Amplitude");
                                        model.setSeriesVisible(0,false);
                                        model.setSeriesVisible(1,true);
                                }
                                graph.setYExtrema(-0.5f,0.5f);
                        }
                });
                Panel levelPanel=new Panel();
                levelPanel.setLayout(new GridLayout(2,1));
                levelPanel.add(incButton);
                levelPanel.add(decButton);
                Panel optPanel=new Panel();
                optPanel.add(kLabel);
                optPanel.add(kField);
                optPanel.add(hLabel);
                optPanel.add(hField);
                optPanel.add(fnButton);
                add(statusLabel,"North");
                add(graph,"Center");
                add(levelPanel,"East");
                add(optPanel,"South");
                H=calculateHamiltonian();
                displayWavefunction();
                setSize(650,450);
                setVisible(true);
        }
        /**
        * Calculates and displays wavefunctions.
        */
        private void displayWavefunction() {
                double eval[]=null;
                DoubleVector evec[]=new DoubleVector[N];
                try {
                        eval=LinearMath.eigenSolveSymmetric(H,evec);
                } catch(MaximumIterationsExceededException e) {
                        System.err.println("Exceeded the maximum number of iterations for finding eigenvalues.");
                        System.exit(-1);
                }
                int sort[]=sortEigenvalues(eval);
// generate graph data
                float data[][]=new float[2][N];
                for(int i=0;i<N;i++) {
                        data[0][i]=(float)evec[sort[level]].getComponent(i);
                        data[1][i]=data[0][i]*data[0][i];
                }
                model.changeSeries(0,data[0]);
                model.changeSeries(1,data[1]);
                if(showProbability)
                        model.setSeriesVisible(0,false);
                else
                        model.setSeriesVisible(1,false);
                graph.setYExtrema(-0.5f,0.5f);
                statusLabel.setText("Energy ["+level+"] = "+(float)eval[sort[level]]);
        }
        /**
        * Calculates the Hamiltonian.
        */
        private DoubleTridiagonalMatrix calculateHamiltonian() {
                DoubleTridiagonalMatrix m=new DoubleTridiagonalMatrix(N);
                m.setElement(0,0,2.0+potential(0));
                m.setElement(0,1,-1.0);
                int i=1;
                for(;i<N-1;i++) {
                        m.setElement(i,i-1,-1.0);
                        m.setElement(i,i,2.0+potential(i));
                        m.setElement(i,i+1,-1.0);
                }
                m.setElement(i,i-1,-1.0);
                m.setElement(i,i,2.0+potential(i));
// generate graph data
                float data[]=new float[N];
                for(i=0;i<N;i++)
                        data[i]=(float)potential(i);
                model.changeSeries(2,data);
                return m;
        }
        /**
        * Potential energy function.
        */
        private double potential(int i) {
                double x=i*2.0/(N-1)-1.0;
                return k*x*x+h*x*x*x*x;
        }
        /**
        * Sorts the eigenvalues.
        * (bidirectional bubble sort).
        */
        private int[] sortEigenvalues(double eval[]) {
                int sort[]=new int[eval.length];
                for(int i=0;i<sort.length;i++)
                        sort[i]=i;
                int j;
                int limit=sort.length;
                int st=-1;
                while(st<limit) {
                        boolean flipped=false;
                        st++;
                        limit--;
                        for(j=st;j<limit;j++) {
                                if(eval[sort[j]]>eval[sort[j+1]]) {
                                        int t=sort[j];
                                        sort[j]=sort[j+1];
                                        sort[j+1]=t;
                                        flipped=true;
                                }
                        }
                        if(!flipped)
                                return sort;
                        for(j=limit;--j>=st;) {
                                if(eval[sort[j]]>eval[sort[j+1]]) {
                                        int t=sort[j];
                                        sort[j]=sort[j+1];
                                        sort[j+1]=t;
                                        flipped=true;
                                }
                        }
                        if(!flipped)
                                return sort;
                }
                return sort;
        }
}

