from scipy.special import eval_legendre #importante para cargar la definicion de legendre

#defino el potencial para coordenadas r y theta, sumando la serie hasta lmax

def potencial(r,theta,lmax):

    pote = 0

    if r < a:

        for l in range (1,lmax+1):

            pote = pote + (r/a)**l*eval_legendre(l,cos(theta))*(eval_legendre(l+1,cos(alpha/2))-eval_legendre(l-1,cos(alpha/2)))/(2*l+1)

        pote = pote + (1 + cos(alpha/2))

    else:

        for l in range (1,lmax+1):

            pote = pote + (a/r)**(l+1)*eval_legendre(l,cos(theta))*(eval_legendre(l+1,cos(alpha/2))-eval_legendre(l-1,cos(alpha/2)))/(2*l+1)

        pote = pote + (1 + cos(alpha/2))*a/r

    

    pote = pote*2*pi*sigma*a

    return pote

    

potencialM=vectorize(potencial) # version que acepta arrays como input.

#defino el corte del potencial a lo largo del eje z, como funcion del lmax.

def potez(lmax):

    def f(z):

        if z > 0:

            return potencial(z,0,lmax)

        else:

            return potencial(-z,pi,lmax)

    return f

#y el corte a lo largo del eje x

def potex(lmax):

    def f(x):

        return potencial(abs(x),pi/2,lmax)

    return f

a=1;sigma=1;z=linspace(-4,4,300); # algunos valores para hacer los gráficos

fz=vectorize(potez(15));fx=vectorize(potex(15)); # especializo los cortes a lmax=15

#plot para varias alpha

alphas=[0,pi/3,2*pi/3,3*pi/3,4*pi/3,5*pi/3];

nalpha=size(alphas)

for i in range(nalpha):

    subplot(2,3,i+1);alpha=alphas[i];plot(z,fz(z),z,fx(z))

x=arange(-4,4,.1);z=x;[Z,X]=meshgrid(z,x);

for i in range(nalpha):

    subplot(2,3,i+1);alpha=alphas[i];

    contour(Z,X,potencialM(sqrt(Z**2+X**2),arccos(Z/sqrt(Z**2+X**2)),15),30)