##################### Plot wavefunctions on band diagram ###################
# Inputs:
#  out: the result struct from mqwgain
#  wavefunction_plot: plot option, 1-conduction band only, 2-valence band only, 3-both
#  max_subband: maximum number of subbands to plot (if 0, plot all)
#
#  NOTE: To be able to use the function below first execute the script file in order to load
#        the function into the memory space.
function plot_wf_on_banddiagram(out,wavefunction_plot,max_subband)
{

if ((wavefunction_plot < 1) | (wavefunction_plot > 3)) {
    ?"Error: plotting options are 1-conduction band, 2-valence band, 3-both";
    return "";
}

bands = out.banddiagram; 
psi = out.wavefunction;
Ek = out.bandstructure;

nCBsubbands = size(Ek.conduction_band,2);
nVBUsubbands = size(Ek.valence_band_up,2);
nVBLsubbands = size(Ek.valence_band_lo,2);

maxBandEc = max(bands.Ec);
minBandEc = min(bands.Ec);
deltaBandEc = maxBandEc - minBandEc;

maxBandEv = max(bands.Ev);
minBandEv = min(bands.Ev);
deltaBandEv = maxBandEv - minBandEv;

deltaBand = max([deltaBandEc,deltaBandEv]);

if (max_subband < 1) {
    max_subband = max([nCBsubbands,nVBUsubbands,nVBLsubbands]);
}

z_nm = bands.z*1e9;

if ((wavefunction_plot == 1) | (wavefunction_plot == 3)) {
    plot(z_nm, bands.Ec, "z (nm)", "Energy (eV)");
    holdon;
    
    for (j = 1:min([nCBsubbands,max_subband])) {
        if (finite(Ek.conduction_band(1,j))) {
            psik = pinch(psi.conduction_band_1);   # nx, nkt, nsubbands
            psik = 50*deltaBand*abs(psik(:,1,j))^2 + real(Ek.conduction_band(1,j)); # kt = 0
            plot(z_nm, psik);
        }
    }
    if (wavefunction_plot != 3) {
        holdoff;
    }
}


if ((wavefunction_plot == 2) | (wavefunction_plot == 3)) {
    plot(z_nm, bands.Ev);
    holdon;
    
    for (j = 1:min([nVBUsubbands,max_subband])) {
        if (finite(Ek.valence_band_up(1,j))) {
            psik = pinch(psi.valence_band_up_1);   # nx, nkt, nsubbands
            psik = 50*deltaBand*abs(psik(:,1,j))^2 + real(Ek.valence_band_up(1,j)); # kt = 0
            plot(z_nm, psik);

            psik = pinch(psi.valence_band_up_2);   # nx, nkt, nsubbands
            psik = 50*deltaBand*abs(psik(:,1,j))^2 + real(Ek.valence_band_up(1,j)); # kt = 0
            plot(z_nm, psik);
        }
    }
    for (j = 1:min([nVBLsubbands,max_subband])) {
        if (finite(Ek.valence_band_lo(1,j))) {
            psik = pinch(psi.valence_band_lo_1);   # nx, nkt, nsubbands
            psik = 50*deltaBand*abs(psik(:,1,j))^2 + real(Ek.valence_band_lo(1,j)); # kt = 0
            plot(z_nm, psik);

            psik = pinch(psi.valence_band_lo_2);   # nx, nkt, nsubbands
            psik = 50*deltaBand*abs(psik(:,1,j))^2 + real(Ek.valence_band_lo(1,j)); # kt = 0
            plot(z_nm, psik);
        }
    }
}
setplot("show legend",false);
holdoff;
}