import numpy as np
import scipy.integrate as spi
import scipy.optimize as sop
import pylab as plt
from sympy import *


#constants
A = 1458.05 #MeV fm
B = 700.0 #MeV
mu_a = 3.11 #fm^-1
mu_b = 1.55 #fm^-1
K = 41.47 #MeV fm^2 (hbar^2/m)

def psisq(r,a,N):
    return N**2*exp(-a*(r**2))*r**2*4*np.pi


def psi_H_psi(r,a,N):
    V = A*exp(-mu_a*r)/r-B*exp(-mu_b*r)/r
    T = 1.5*K*a
    return psisq(r,a,N)*(V+T)

a = symbols('alpha')
r = symbols('r')
N = (a/np.pi)**(.75)

#varijacijska
#<psi|H|psi> -> d<H>/d(a) = 0 -> a_min -> H

norm = integrate(psisq(r,a,N), (r, 0, oo))
print norm.subs(a, .5)
raw_input("Press ENTER to continue.")

expr = integrate(psi_H_psi(r,a,N), (r, 0, oo), conds='none')
print expr
raw_input("Press ENTER to continue.")

a = None

sol = sop.fmin(expr, .5, xtol=1e-4, args=(a,))
print sol
#expr = diff(expr, a)
#print expr
#raw_input("Press ENTER to continue.")

#a_var = sop.fsolve(expr, 0.5)
#print a_var