#!/usr/bin/python -Wall

# ================================================================
# John Kerl
# kerl at math dot arizona dot edu
# 2008-02-06
#
# Ideas here are taken from a talk given by Itai Seggev at the AMS/MAA joint
# meetings in San Diego in January 2008.
# ================================================================
# This software is released under the terms of the GNU GPL.
# Please see LICENSE.txt in the same directory as this file.
# ================================================================

import sackmat_m
import sackmatc_m
import kerlutil
import sys
from math import *
from cplxreal_m import *

# Closer masses give slower oscillations.
# Eigenvalue degeneracy when a == b ...
# theta = pi/4 promotes mixing; theta = 0 does not.

# I am taking hbar = 1 and E_0 = 1.

# What about the density operator?

# ----------------------------------------------------------------
a     = 0.4
b     = 0.5
theta = 1.0 * pi/4

dt    = 0.1
N     = 100

initv = 0
pev   = 0

# ----------------------------------------------------------------
# Parse the command line.  Syntax: not -a 2.1 but a=2.1.
# This is crude (it's not syntax-error tolerant) but it's an easy-to-code
# hack.  It works in any language with an eval/exec (e.g. Python, Perl).
argc = len(sys.argv)
for argi in range(1, argc):
	if ((sys.argv[argi] == '-h') or (sys.argv[argi] == '--help')):
		print >> sys.stderr, \
			"Usage: %s [a=...] [b=...] [theta=...] [pev={0,1}]" \
			% (sys.argv[0])
		sys.exit(1)
	exec(sys.argv[argi])

# ----------------------------------------------------------------
c = cos(theta)
s = sin(theta)
a2 = a**2
b2 = b**2

H = sackmat_m.m([
	[a2 * c**2 + b2 * s**2, (a2-b2) * c * s ],
	[(a2-b2) * c * s,       a2 * s**2 + b2 * c**2 ]
	]);

ket0 = [1.0, 0.0]
ket1 = [0.0, 1.0]

[V,D] = sackmat_m.rs_eigensystem(H) # Only in sackmat_m

neg_iH0 = -H * 1j * 0
neg_iHt = -H * 1j * dt
A0  = neg_iH0.exp()
Adt = neg_iHt.exp()

if (initv):
	#psi0 = V.get_column(0)
	psi0 = V.get_column(1)
else:
	psi0 = ket0
	#psi0 = ket1

# ----------------------------------------------------------------
if pev:
	print "H:";            H.printf();    print
	print "A(0):";         A0.printf();   print
	print "A(dt):";        Adt.printf();  print
	print "Eigenvalues:";  D.printf();    print
	print "Eigenvectors (columns):"; V.printf();    print
	print "psi(0):";       sackmatc_m.print_row_vector(psi0); print

else:

	psit = psi0

	for t in kerlutil.mfrange(0.0, dt, N*dt):

		psit = Adt * psit

		p0 = abssq(psit[0])
		p1 = abssq(psit[1])

		q0 = phz(psit[0])
		q1 = phz(psit[1])

		r0 = real(psit[0])
		r1 = real(psit[1])

		i0 = imag(psit[0])
		i1 = imag(psit[1])

		#sackmat_m.print_row_vector_no_cr([abs(At.det())]);     print "  ", 

		#print "H:"
		#H.printf();
		#print

		#print "Adt:"
		#Adt.printf();
		#print

		sackmat_m.print_row_vector_no_cr([t]);
		print "  ", 

		#sackmatc_m.print_row_vector(psit);
		sackmat_m.print_row_vector([p0, p1]);