// ComPhys File: TDSE.java
// Chapter 18: The Time-Dependent Schroedinger Equation

import comphys.*;
import java.applet.*;
import java.awt.*;
import java.awt.event.*;
import java.util.*;

public class TDSE extends Applet implements ActionListener, Runnable {

    // arrays to hold wavefunction

    double xmax = 20;
    double xmin = -xmax;
    double dx = 0.4;

    int n;
    double[] Re;
    double[] Im;
    double[] Imold;

    void allocateArrays () {

        n = (int) ( (xmax - xmin) / dx );

        if (Re == null || Re.length < n + 2) {

            Re = new double[n + 2];
            Im = new double[n + 2];
            Imold = new double [n + 2];

        }

    }

    // step potential
    
    double V0 = 2;
    double a = 1;

    double V (double x) {

        if (x > a)
            return V0;
        else return 0;

    }

    double probability () {

        double Psum = 0;
        for (int i = 0; i < n; i++)
            Psum += Re[i] * Re[i] + Im[i] * Imold[i];

        return Psum * dx;

    }

    // Visscher et al., algorithm for time evolution

    double t = 0;
    double dt = 0.1;

    void evolve () {

        double dx2 = dx * dx;

        for (int i = 1; i <= n; i++) {

            double x = xmin + (i - 1) * dx;
            double HIm = V(x) * Im[i] -
                0.5 * (Im[i + 1] - 2 * Im[i] + Im[i - 1]) / dx2;
            Re[i] += HIm * dt;
            
        }

        for (int i = 1; i <= n; i++) {

            double x = xmin + (i - 1) * dx;
            Imold[i] = Im[i];
            double HRe = V(x) * Re[i] -
                0.5 * (Re[i + 1] - 2 * Re[i] + Re[i - 1]) / dx2;
            Im[i] -= HRe * dt;

        }

        t += dt;

    }

    // initial Gaussian wave packet

    double x0 = -15;
    double k0 = 2;
    double width = 1;

    void initial_packet () {

        allocateArrays();

        double delta2 = width * width;
        double A = 2 * Math.PI * delta2;
        A = Math.pow(A, -0.25);
        double b = 0.5 * k0 * dt;

        Re[0] = Im[0] = Imold[0] = 0;

        for (int i = 1; i <= n; i++) {

            double x = xmin + (i - 1) * dx;
            double arg = 0.25 * (x - x0) * (x - x0) / delta2;
            double e = Math.exp(-arg);
            Re[i] = A * Math.cos(k0 * (x - x0)) * e;
            Im[i] = A * Math.sin(k0 * (x - x0 - 0.5 * b)) * e;
            Imold[i] = Im[i];

        }

        Re[n + 1] = Im[n + 1] = Imold[n + 1] = 0;

    }

    class Picture extends Canvas {

        int xSize = 400;
        int ySize = 400;
        Color bgColor = new Color(255, 250, 240);

        Picture () {

            setSize(xSize, ySize);
            setBackground(bgColor);

        }

        public void paint (Graphics g) {

            update(g);

        }

        Image offScreen;
        Graphics osg;

        void draw_potential () {

            osg.setColor(bgColor);
            osg.drawLine(0, ySize / 2, xSize, ySize / 2);
            
            osg.setColor(Color.green);
            int iyOld = 0;
            for (int ix = 0; ix <= xSize; ix++) {

                double x = xmin + ix * (xmax - xmin) / (double) xSize;
                double y = V(x);
                y *= ySize / (2 * 1.2 * V0);
                int iy = ySize / 2 - (int) y;
                osg.drawLine(ix, iy, ix, iy);
                if (ix > 0 && Math.abs(iy - iyOld) > 1)
                    osg.drawLine(ix, iyOld, ix, iy);
                iyOld = iy;

            }

        }

        void plots () {

            int iy0 = ySize / 2;
            int iyPold = iy0;
            int iyRold = iy0;
            int iyIold = iy0;
            for (int ix = 0; ix <= xSize; ix++) {

                int i = (int) Math.round(ix / (double) xSize * n);
                double P = Re[i] * Re[i] + Im[i] * Imold[i];
                int iy = iy0 - (int) (2 * ySize * P);
                osg.setColor(Color.yellow);
                osg.drawLine(ix, iy0, ix, iy);
                osg.setColor(Color.gray);
                osg.drawLine(ix, iyPold, ix, iy);
                iyPold = iy;
                iy = iy0 - (int) (0.5 * ySize * Re[i]);
                osg.setColor(Color.red);
                osg.drawLine(ix, iyRold, ix, iy);
                iyRold = iy;
                iy = iy0 - (int) (ySize * Im[i]);
                osg.setColor(Color.blue);
                osg.drawLine(ix, iyIold, ix, iy);
                iyIold = iy;

            }

        }
        
        public void update (Graphics g) {

            if (offScreen == null) {
                
                offScreen = createImage(xSize, ySize);
                osg = offScreen.getGraphics();

            }

            osg.clearRect(0, 0, xSize, ySize);
            plots();
            draw_potential();

            int ix = 10;
            int iy = ySize - 10;
            osg.setColor(Color.black);
            osg.drawString("t = " + CP.format(t, 4), ix, iy);
            ix = xSize / 2 + 10;
            osg.drawString("Probability = " + CP.format(probability(), 6),
                           ix, iy);
            ix = 10;
            iy = 15;
            int len = 40;
            osg.setColor(Color.red);
            osg.drawLine(ix, iy, ix + len, iy);
            ix += len + 5;
            osg.drawString("Re(psi)", ix, iy);
            ix = xSize / 2 + 10;
            osg.setColor(Color.blue);
            osg.drawLine(ix, iy, ix + len, iy);
            ix += len + 5;
            osg.drawString("Im(psi)", ix, iy);

            g.drawImage(offScreen, 0, 0, null);

        }

    }

    Picture picture;
    TextField widthTextField;
    TextField k0TextField;
    TextField dxTextField;
    TextField dtTextField;
    TextField skipTextField;
    int nSkip = 0;
    Button runButton;
    Button resetButton;

    public void init () {

        initial_packet();

        add(picture = new Picture());

        Panel p = new Panel();
        p.setLayout(new GridLayout(6, 2));

        p.add(new Label("Packet width:"));
        p.add(widthTextField = new TextField("" + width));
        widthTextField.addActionListener(this);

        p.add(new Label("Group velocity:"));
        p.add(k0TextField = new TextField("" + k0));
        k0TextField.addActionListener(this);

        p.add(new Label("Grid size dx:"));
        p.add(dxTextField = new TextField("" + dx));
        dxTextField.addActionListener(this);

        p.add(new Label("Time step dt:"));
        p.add(dtTextField = new TextField("" + dt));
        dtTextField.addActionListener(this);

        p.add(new Label("Steps to skip:"));
        p.add(skipTextField = new TextField("" + nSkip));
        skipTextField.addActionListener(this);

        p.add(runButton = new Button("Start"));
        runButton.addActionListener(this);

        p.add(resetButton = new Button("Reset"));
        resetButton.addActionListener(this);

        add(p);

    }

    Thread runThread;
    boolean running;

    public void actionPerformed (ActionEvent event) {

        boolean doReset = false;

        if (event.getSource() == widthTextField) {

            try {
                
                width = new Double(event.getActionCommand()).doubleValue();
                
            } catch (NumberFormatException nfe) { }
            
            widthTextField.setText("" + width);

        } else if (event.getSource() == k0TextField) {

            try {
                
                k0 = new Double(event.getActionCommand()).doubleValue();
                
            } catch (NumberFormatException nfe) { }
            
            k0TextField.setText("" + k0);

        } else if (event.getSource() == dxTextField) {

            double dxSave = dx;
            try {
                
                dx = new Double(event.getActionCommand()).doubleValue();
                if (dx < 0)
                    dx = -dx;
                if (dx == 0)
                    dx = dxSave;
                
            } catch (NumberFormatException nfe) { }
            
            dxTextField.setText("" + dx);
            
            if (dx != dxSave)
                doReset = true;
            
        } else if (event.getSource() == dtTextField) {

            try {
                
                dt = new Double(event.getActionCommand()).doubleValue();
                
            } catch (NumberFormatException nfe) { }
            
            dtTextField.setText("" + dt);

        } else if (event.getSource() == skipTextField) {

            try {
                
                nSkip = Integer.parseInt(event.getActionCommand());
                if (nSkip < 0)
                    nSkip = 0;
                
            } catch (NumberFormatException nfe) { }
            
            widthTextField.setText("" + nSkip);

        } else if (event.getSource() == runButton) {

            if (!running) {
                
                running = true;
                runButton.setLabel("Stop");
                runThread = new Thread(this);
                runThread.start();

            } else {

                running = false;

            }


        } else if (event.getSource() == resetButton) {

            doReset = true;

        }

        if (doReset) {

            running = false;
            try {
                Thread.sleep(100);
            } catch (InterruptedException e) { }
            t = 0;
            initial_packet();
            picture.repaint();

        }

        if (!running)
            runButton.setLabel("Start");

    }

    public void run () {

        runThread.setPriority(Thread.MIN_PRIORITY);

        while (running) {

            evolve();
            picture.repaint();
            for (int skip = 0; skip < nSkip; skip++)
                evolve();

            try {
                Thread.sleep(30);
            } catch (InterruptedException e) { }

        }

    }

    public static void main (String[] args) {

        TDSE tdse = new TDSE();
        CPFrame frame = new CPFrame("Time-Dependent Schroedinger Equation");

        frame.add(tdse);
        tdse.init();
        frame.setSize(700, 450);
        frame.setLocation(50, 50);
        frame.setVisible(true);

    }

}