#!/usr/bin/python3
# qbox_eig.py: extract eigenvalues from Qbox output
# use: qbox_eig.py kpoint n ispin file.r
# extracts eigenvalue n at (ispin,kpoint)
# note: ispin = 0..1, first eigenvalue is n=1


# This file is part of PyEPFD
# Copyright (c) 2023 Arpan Kundu
# See the LICENCE.md in root directory for full license information.

import xml.sax
import sys
import math

argc=len(sys.argv)
if ( not ( argc in [3,4,6,7] ) ):
  print("use: "+sys.argv[0]+" [ispin] [kx ky kz] n file.r")
  print(" ispin = 0..1, n = 1..neig")
  sys.exit()

if argc == 7:
  # ispin kx ky kz n file.r
  ispin = int(sys.argv[1])
  kx = float(sys.argv[2])
  ky = float(sys.argv[3])
  kz = float(sys.argv[4])
  n = int(sys.argv[5])
  infile = sys.argv[6]
elif argc == 6:
  # kx ky kz n file.r
  ispin = 0
  kx = float(sys.argv[1])
  ky = float(sys.argv[2])
  kz = float(sys.argv[3])
  n = int(sys.argv[4])
  infile = sys.argv[5]
elif argc == 4:
  # ispin n file.r
  ispin = int(sys.argv[1])
  kx = 0.0
  ky = 0.0
  kz = 0.0
  n = int(sys.argv[2])
  infile = sys.argv[3]
elif argc == 3:
  ispin = 0
  kx = 0.0
  ky = 0.0
  kz = 0.0
  n = int(sys.argv[1])
  infile = sys.argv[2]

print("# "+infile+" ispin="+str(ispin)+ " n="+ str(n)+ " k="+ str(kx) + str(ky) + str(kz))

# Qbox output handler to extract and process data
class QboxOutputHandler(xml.sax.handler.ContentHandler):
  def __init__(self):
    self.iter = 0
    self.readData = 0

  def startElement(self, name, attributes):
    if name == "eigenvalues":
      self.n = attributes["n"]
      self.spin = attributes["spin"]
      self.kpoint = attributes["kpoint"]
      self.readData = 1
      self.iter += 1
      self.buffer = ""

  def characters(self, data):
    if self.readData:
      self.buffer += data

  def endElement(self, name):
    if name == "eigenvalues":
      self.readData = 0
      isp = int(self.spin)
      self.kp = self.kpoint.split()
      dx = kx-float(self.kp[0])
      dy = ky-float(self.kp[1])
      dz = kz-float(self.kp[2])
      if isp == ispin and math.sqrt(dx*dx+dy*dy+dz*dz) < 1.e-6:
        self.print_eig()

  def print_eig(self):
    self.e = self.buffer.split()
    if n > int(self.n):
      print("n>neig: neig="+ str(self.n))
    else:
      print (str(self.e[n-1]))

parser = xml.sax.make_parser()
handler = QboxOutputHandler()
parser.setContentHandler(handler)
parser.parse(infile)
