from ATK.KohnSham import *
from numpy import array
from separations import *

def writeBandstructureDataFile (filename,symmetry_points,route,
                                k_points,bandstructure_up,bandstructure_dn):
    """
    Create a Pickle data file with the bandstructure
    Three objects are dumped: kx (1D values of the k-points
    for the horizontal axis), and energies for up and down spin
    # The last symmetry point is of course not included
    # except as endpoint for the last segment;
    # hence the -1 in the range commands
    """
    # Generate 1D k-points for the horizontal axis
    kx_points = []
    separations,separationsums = generateSeparations(symmetry_points,route)
    for point in range(len(route)-1):
        numKPoints = len(k_points[point])
        kx_points += convertVectorPointToX(numKPoints,point,separations,separationsums)
    # Make empty lists of the right size for the bandstructure data
    # Each list has as many elements as there are bands
    # Each element is a list, the length of which matches the total number of k-points
    energies_up = []
    energies_dn = []
    number_of_bands = bandstructure_up[0].numberOfBands()
    for band in range(number_of_bands):
        energies_up.append([])
        energies_dn.append([])
    # Extract the bands, strip off the unit, and store in the lists
    for point in range(len(route)-1):
        for band in range(number_of_bands):
            tmp = bandstructure_up[point].band(band)
            energies_up[band] = energies_up[band] + [E.inUnitsOf(eV) for E in tmp]
            tmp = bandstructure_dn[point].band(band)
            energies_dn[band] = energies_dn[band] + [E.inUnitsOf(eV) for E in tmp]
    # Store data in pickle file
    import pickle
    f = open(filename,"wb")
    pickle.dump(kx_points,f)
    pickle.dump(energies_up,f)
    pickle.dump(energies_dn,f)
    f.close()

def calculateBandstructure(calculation,symmetry_points,route,n_points,spin):
    """
    Calculate the bandstructure for the "calculation" object,
    using the provided symmetry points and the route that
    connects them. This routines defines the k-points,
    calculate the bandstructure, and return both. In principle
    each segment of the bandstructure plot can contain a
    different number of k-points, but for simplicity we choose
    an equal number for now.
    """
    band_k_points = []
    bandstructure = []
    # The last symmetry point is of course not included
    # except as endpoint for the last segment
    for point in range(len(route)-1):
        start_point = symmetry_points[route[point]]
        end_point = symmetry_points[route[point+1]]
        band_k_points.append(generateLineBetweenVectorsPoints(start_point,end_point,n_points))
        bandstructure.append(calculateEnergyBands(calculation,band_k_points[point],spin))
    return band_k_points, bandstructure

def writePyLabScript(plotfilename,datafilename,imagename,symmetry_points,route):
    # Create pyLab script to plot the data
    f = open(plotfilename,"w")
    separations,separationsums = generateSeparations(symmetry_points,route)
    # The text block to be written to the script
    s = """
import matplotlib
matplotlib.use('Agg')
from pylab import *

import pickle
f = open("%s","rb")
X    = pickle.load(f)
Y_up = pickle.load(f)
Y_dn = pickle.load(f)
f.close()

for i in range(len(Y_up)):
    plot(X,Y_up[i],'r-')
    plot(X,Y_dn[i],'b-')
  
ylabel('Energy (eV)')
xticks (%s,%s)
grid(True)
v1 = axis()
v2 = [v1[0],%g,v1[2],v1[3]]
axis(v2)
savefig('%s',dpi=100)
    """ % (datafilename,separationsums,\
    route,separationsums[-1]*1.001,imagename)
    f.write(s)

def calculateAndPlotSpinBandstructure(
        ncfilename,datafilename,plotfilename,imagename,
        symmetry_points,route,number_of_kpoints=25):

    calculation = restoreSelfConsistentCalculation(ncfilename)
    band_k_points,bandstructure_up = calculateBandstructure(
        calculation, symmetry_points, route,
        number_of_kpoints, Spin.Up)

    band_k_points,bandstructure_dn = calculateBandstructure(
        calculation, symmetry_points, route,
        number_of_kpoints, Spin.Down)

    writeBandstructureDataFile(datafilename,symmetry_points,route,
                               band_k_points,bandstructure_up,bandstructure_dn)

    writePyLabScript(plotfilename,datafilename,imagename,symmetry_points,route)