#!/usr/bin/env python
# Author: Samuel Ponc\'e
# Date: 24/04/2013
# Classes needed for the temperature.py script

import sys
import os
try:
  import numpy as N
except ImportError:
  import warnings
  warnings.warn("The numpy module is missing!")
  raise
try:
  import netCDF4 as nc
except ImportError:
  import warnings
  warnings.warn("The netCDF4 module is missing!")
  raise


# Read _EIG.nc file
class EIG_file_open:
  EIG = None
  Kptns = None
  nband = None
  def __init__(self,filefullpath):
    print filefullpath
    if not (os.path.isfile(filefullpath)):
      raise Exception('The file "%s" does not exists!' %filefullpath)
    root = nc.Dataset(filefullpath,'r')
    self.EIG = root.variables['Eigenvalues'][:,:,:]
    self.Kptns = root.variables['Kptns'][:,:]
    self.NBandK = root.variables['NBandK'][:,:]
    root.close()



# Interaction with the user
print '\n##############################'
print '# Merge eigenvalue EIG.nc files #'
print '#################################'
print '\nThis script merges different EIG.nc files into one. \n\
This is usefull to later create a full electronic bandstructure. \n'

# Define the output file name
user_input = raw_input('Enter name of the EIG.nc files separated by a space that you want to merge\n')
eig_list = user_input.split()

# Define the output file name
user_input = raw_input('Enter name of the output file\n')
output = user_input

out = nc.Dataset(output,'w')
nsppol = out.createDimension('nsppol',None)
nkpt = out.createDimension('nkpt',None)
nkpt2 = out.createDimension('nkpt2',None)
nkpt3 = out.createDimension('nkpt3',None)
mband = out.createDimension('mband',None)
xyz = out.createDimension('xyz',None)

Eigenvalues  = out.createVariable('Eigenvalues','f8',('nsppol','nkpt','mband',))
Eigenvalues.units = "Hartree"
Kptns  = out.createVariable('Kptns','f8',('nkpt2','xyz',))
Kptns.units = "Dimensionless"
NBandK  = out.createVariable('NBandK','i4',('nsppol','nkpt3',))
NBandK.units = "Dimensionless"

first = True
for ii in eig_list:
  current = EIG_file_open(ii)
  if first:
    Eigenvalues[:,:,:] = current.EIG[:,:,:]
    Kptns[:,:] = current.Kptns[:,:]
    NBandK[:,:] = current.NBandK[:,:]
    first = False
  else:
    Eigenvalues[:,:,:] = N.append(Eigenvalues[:,:,:],current.EIG[:,:,:],axis=1)
    Kptns[:,:] = N.append(Kptns[:,:],current.Kptns[:,:],axis=0)
    NBandK[:,:] = N.append(NBandK[:,:],current.NBandK[:,:],axis=1)

out.close()