#!/usr/bin/env python3

from sympy import *
from sympy.physics.quantum.constants import hbar
from sympy.physics.quantum.operator import DifferentialOperator
from sympy.physics.quantum.qapply import qapply
from sympy.physics.quantum.state import Wavefunction

#from sympy.physics.qho_1d import psi_n

#Some general definitions
m = Symbol('m', nonzero=True, positive=True)
omega = Symbol('omega', nonzero=True, positive=True)
x_var = Symbol('x')
f = Function('f')

########################################################################
# This function calculates Bra Ket Products of form <n1| Operator |n2> #
# for the harmonic 1d Oscillator with potential V(x) = 1/2 omega^2 x^2 #
########################################################################

def ho1dExpectaion(n1, op, n2):

	def H(n,y):
		yDiff = Symbol('yp')
		out = (-1)**n
		out *= exp(y**2)
		out *= diff(exp(-(yDiff)**2), yDiff, n).subs(yDiff, y)
		
		return out

	def h1dWaveFunction(n, y):
		out = ((m*omega)/(pi *hbar))**Rational(1/4)
		out *= 1/sqrt((2**n)*factorial(n))
		out *= H(n,sqrt((m*omega)/hbar)*y)
		out *= exp(-((m*omega)/(2*hbar))*y**2)

		return out

	Hbc = DifferentialOperator(Derivative(f(x_var),x_var), f(x_var))

	w1 = Wavefunction(h1dWaveFunction(n1,x_var),x_var)
	w2 = Wavefunction(h1dWaveFunction(n2,x_var),x_var)
	v = qapply(op*w2)

	return integrate(w1(x_var)*v(x_var), (x_var,-oo,oo))


##################################
# Here are some useful Operators #
##################################

#Momentum Operator
p = DifferentialOperator(- I *hbar*Derivative(f(x_var),x_var), f(x_var))

#Position Operator
x = DifferentialOperator( x_var *f(x_var), f(x_var))

#Hamiltonian
H = DifferentialOperator(- hbar**2/(2*m) * Derivative(f(x_var),x_var,2) + (m/2) *omega**2 *x_var**2 *f(x_var), f(x_var))

#Potential Energy Operator
E_pot = DifferentialOperator((m/2) *omega**2 *x_var**2 *f(x_var), f(x_var))

#Kinetic Energy Operator
E_kin = DifferentialOperator(- hbar**2/(2*m) * Derivative(f(x_var),x_var,2), f(x_var))

#Lowering Operator
a = DifferentialOperator(sqrt((m * omega)/(2 * hbar))*(x_var*f(x_var) + ((hbar * Derivative(f(x_var),x_var))/(m * omega))), f(x_var))

#Raising Operator
ad = DifferentialOperator(sqrt((m * omega)/(2 * hbar))*(x_var*f(x_var) - (hbar * Derivative(f(x_var),x_var))/(m*omega)), f(x_var))



#################################
# Here you do your Calculations #
#################################


#Calculate for example <0| H |0>
pprint(ho1dExpectaion(1,E_kin,1))