%% Hydrogen WaveFunction Generator
close all; clear all; clc
n = 5;
l = 3;
m = 0;
ao=1;
Box=110
%% Probability Density. run first
N=1000; %modify for number of points
threshold=.000005; % modify for the probability threshold
k=1

%find random point
while k < N
    u=rand(3,1);
    x=0;
    y=Box*rand(1,1);
    z=Box*rand(1,1);
    rho=2*sqrt(x^2+y^2+z^2)/(n*ao);

% spherical harmonics
    Lmn=legendre(l,cos(acos(z/sqrt(x^2+y^2+z^2))));
    if l~=0
        Lmn=squeeze(Lmn(m+1,:,:));
    end
    Ylm=sqrt(((2*l+1)/(4*pi))*factorial(l-m)/factorial(l+m))*Lmn.*exp(i*m*atan(y/x));

% laguerre polynomials
    Lnl=polyval(LaguerreGen(n-l-1,2*l+1),rho);
    P=abs(sqrt((2/(n.*ao)).^3*factorial(n-l-1)/(2*n*factorial(n+l))).*exp(-rho./2).*rho.^l.*Lnl.*Ylm).^2;
    if P>threshold
        a(k)=x;
        b(k)=y;
        c(k)=z;
        d(k)=P;
        k=k+1
    else
        continue
    end
end

%% Rotation Matricies. this will use the symmetry of the wavefunction to
%% generate a three-dimensional plot! some of these settings require fine
%% tuning
figure
clear C2

Ndiv=53;%number of divisions. larger is lighter
Normalized=d./threshold;
Div=(max(Normalized)-min(Normalized))/(Ndiv);%larger is lighter
for i = 2:Ndiv;
    C2(i)={[1-1/i,1-1/i,1-1/i]};
end
C2{1}=[.2,.2,.2];

Angle=linspace(0,pi/2,100);
hold on
for f = 1:length(Angle);
ap=a.*cos(Angle(f))+b.*sin(Angle(f));
bp=a.*(-Angle(f))+b.*cos(Angle(f));
for i = 1:Ndiv;
ind=find(Normalized>(i-1)*Div & Normalized<(i)*Div);
plot3(ap(ind),bp(ind),c(ind),'.','MarkerSize',5,'Color',C2{i})
plot3(-ap(ind),bp(ind),c(ind),'.','MarkerSize',5,'Color',C2{i})
plot3(-ap(ind),-bp(ind),c(ind),'.','MarkerSize',5,'Color',C2{i})
plot3(-ap(ind),-bp(ind),-c(ind),'.','MarkerSize',5,'Color',C2{i})
plot3(-ap(ind),bp(ind),-c(ind),'.','MarkerSize',5,'Color',C2{i})
plot3(ap(ind),bp(ind),-c(ind),'.','MarkerSize',5,'Color',C2{i})
end
end
axis  image; 
set(gcf,'color',[1,1,1]); 
set(gca,'color',[1,1,1]);
view([26,16]);
title(texlabel(sprintf('|Psi_%d_%d_%d|^2',n,l,m)),'FontWeight','bold',...
    'FontSize',13);
hold off

%% Contour Density Plot
clc
CBox=Box
[xMat,zMat]=meshgrid(linspace(-CBox,CBox,1500),linspace(-CBox,CBox,1500));
rho= 2*sqrt(0^2+xMat.^2+zMat.^2)/(n*ao);
 
% Angular wavefunction
    Lmn=legendre(l,cos(acos(zMat./sqrt(xMat.^2+zMat.^2))));
    if l~=0
        Lmn=squeeze(Lmn(m+1,:,:));
    end
    Ylm=sqrt(((2*l+1)/(4*pi))*factorial(l-m)/factorial(l+m))*Lmn.*exp(i*m*atan(0./xMat));

% Radial wavefunction
Lnl=polyval(LaguerreGen(n-l-1,2*l+1),rho);

%constants
a1=((2*l+1)/(4*pi));
a2=factorial(l-m)/factorial(l+m);
C=sqrt(a1*a2);
P=abs(sqrt((2/(n.*ao)).^3*factorial(n-l-1)/(2*n*factorial(n+l))).*exp(-rho./2).*rho.^l.*Lnl.*Ylm).^2;
P=P./max(max(P));

%contour3(zMat,xMat,P,20);
figure
surf(zMat,xMat,P,'FaceLighting','phong',...
    'EdgeLighting','phong',...
    'SpecularStrength',.7);
light;lighting phong;shading interp
axis  off; set(gcf,'color',[1,1,1]);
figure
axis off; set(gcf,'color',[1,1,1]);
[C,h]=contour('v6',xMat,zMat,P,20);axis tight equal
hold on;axis off
figure
imagesc(P);axis equal tight
%contour(zMat,xMat,P,100);
%light;lighting phong;

%% Get a new data set from contours

newind=find(P<threshold | P>1.9*threshold)
P2=P;
P2(newind)=0
newind2=find(P2>0)
y2=P2(newind2)
z2=zMat(newind2)
x2=xMat(newind2)


surf(zMat,xMat,P2)



%% Spherical Harmonic Plot
% Inputs
clear i
figure
[phi,theta] = meshgrid(linspace(0, pi, 110),linspace(0,2*pi, 110));
        %Create the Harmonics
        Lmn=legendre(l,cos(phi));

        if l~=0
            Lmn=squeeze(Lmn(m+1,:,:));
        end

        %C=sqrt(((2*l+1)/(4*pi))*factorial(l-m)/factorial(l+m));
        Ymn=Lmn.*exp(i*m*theta);
        %Composite Probability Density
        [xp,yp,zp]=sph2cart(theta,phi-pi/2,abs(Ymn).^2);
        
       surf((xp),(yp),(zp));
       shading flat;
       light;lighting phong;
       axis equal  off;
       title(sprintf('|(Y(%d,%d))|^2',m, l));   
       colormap cool;
       set(gcf,'color',[1,1,1])
hold on
%% Radial Wavefunction plot
r=linspace(0,Box,1000);
Lnl=polyval(LaguerreGen(n-l-1,2*l+1),2*r/n);
Rnl=Lnl.*(2*r/n).^l.*exp(-2*r./(2*n));hold on
C2=1/sqrt(2)*(2/ao)^(3/2)
Rnl=C2*Rnl
figure
plot(r,abs(Rnl).^2,'r','linewidth',2);
set(gcf,'color',[1,1,1])
title(sprintf('|(R(%d,%d))|^2',n, l));

%% Save Orbital Data
pointinfo={a,b,c,d};
A=sprintf('save n%d_l%d_m%d pointinfo',n,l,m)
eval(A)