# Path: vasppy/procar.py
# class KPoint():
# class Procar:
#     def __init__( self, index, frac_coords, weight ):
#     def cart_coords( self, reciprocal_lattice ):
#     def __eq__( self, other ):
#     def __repr__( self ):
# def get_numbers_from_string( string ):
# def k_point_parser( string ):
# def projections_parser( string ):
# def area_of_a_triangle_in_cartesian_space( a, b, c ):
# def points_are_in_a_straight_line( points, tolerance=1e-7 ):
# def two_point_effective_mass( cartesian_k_points, eigenvalues ):
# def least_squares_effective_mass( cartesian_k_points, eigenvalues ):
#     def __init__( self, spin=1, negative_occupancies='warn' ):
#     def occupancy( self ):
#     def __add__( self, other ):
#     def parse_projections( self ):
#     def parse_k_points( self ):
#     def parse_bands( self ):
#     def sanity_check( self ):
#     def from_files( cls, filenames, **kwargs ):
#     def from_file( cls, filename, negative_occupancies='warn',
#                    select_zero_weighted_k_points=False ):
#     def read_from_file( self, filename ):
#     def _read_from_file( self, filename ):
#     def number_of_k_points( self ):
#     def number_of_bands( self ):
#     def spin_channels( self ):
#     def number_of_ions( self ):
#     def number_of_projections( self ):
#     def print_weighted_band_structure( self, spins=None, ions=None, orbitals=None, scaling=1.0, e_fermi=0.0, reciprocal_lattice=None ):
#     def weighted_band_structure( self, spins=None, ions=None, orbitals=None, scaling=1.0, e_fermi=0.0, reciprocal_lattice=None ):
#     def effective_mass_calc( self, k_point_indices, band_index, reciprocal_lattice, spin=1, printing=False ):
#     def x_axis( self, reciprocal_lattice=None ):
#     def bands( self ):
#     def k_points( self ):
#     def select_bands_by_kpoint( self, band_indices ):
#     def select_k_points( self, band_indices ):
# 
# Path: vasppy/outcar.py
# def reciprocal_lattice_from_outcar( filename ): # from https://github.com/MaterialsDiscovery/PyChemia
#     """
#     Finds and returns the reciprocal lattice vectors, if more than
#     one set present, it just returns the last one.
#     Args:
#         filename (Str): The name of the outcar file to be read
# 
#     Returns:
#         List(Float): The reciprocal lattice vectors.
#     """
#     outcar = open(filename, "r").read()
#     # just keeping the last component
#     recLat = re.findall(r"reciprocal\s*lattice\s*vectors\s*([-.\s\d]*)",
#                         outcar)[-1]
#     recLat = recLat.split()
#     recLat = np.array(recLat, dtype=float)
#     # up to now I have, both direct and rec. lattices (3+3=6 columns)
#     recLat.shape = (3, 6)
#     recLat = recLat[:, 3:]
#     return recLat


from vasppy import procar
from vasppy.outcar import reciprocal_lattice_from_outcar
import argparse

def orbitals_with_l( l ):
    to_return = { 's' : [ 0 ],
                  'p' : [ 1, 2, 3 ],
                  'd' : [ 4, 5, 6, 7, 8 ],
                  'f' : [ 9, 10, 11, 12, 13 ],
                  'all' : None }
    return to_return[ l ]
 
def main():
    parser = argparse.ArgumentParser()
    parser.add_argument( '-i', '--ions', help='ion indices for band projection (default: sum over all ions)', nargs='+', type=int )
    parser.add_argument( '-s', '--spins', help='spin indices for band projection (default [ 1 ])', nargs='+', type=int, default=[ 1 ] )
    parser.add_argument( '-o', '--orbitals', help='orbital indices for band projection (default: sum over all orbitals)', nargs='+', type=int )
    parser.add_argument( '-e', '--efermi', help='set fermi energy as reference for energy scale', type=float, default=0.0 )
    parser.add_argument( '-l', '--l-angular-momentum', help='select all orbitals with angular momentum L for band projection. This supercedes the --orbitals option', choices=[ 's', 'p', 'd', 'f', 'all' ] )
    parser.add_argument( '-f', '--procar', help='PROCAR filename (default PROCAR)', type=str, default='PROCAR' )
    parser.add_argument( '--scaling', help='Energy scaling for band widths (default 0.2 eV)', type=float, default=0.2 )
    parser.add_argument( '-x', '--xscaling', help='Automatic scaling of x-axis using reciprocal lattice vectors read from OUTCAR', action='store_true', default=False )
    args = parser.parse_args()

    if args.l_angular_momentum:
        args.orbitals = orbitals_with_l( args.l_angular_momentum )

    if args.xscaling:
        reciprocal_lattice = reciprocal_lattice_from_outcar( 'OUTCAR' ) # Move reading the reciprocal lattice to procar.py
    else:
        reciprocal_lattice = None

    pcar = procar.Procar()

# Path: vasppy/procar.py
# class KPoint():
# class Procar:
#     def __init__( self, index, frac_coords, weight ):
#     def cart_coords( self, reciprocal_lattice ):
#     def __eq__( self, other ):
#     def __repr__( self ):
# def get_numbers_from_string( string ):
# def k_point_parser( string ):
# def projections_parser( string ):
# def area_of_a_triangle_in_cartesian_space( a, b, c ):
# def points_are_in_a_straight_line( points, tolerance=1e-7 ):
# def two_point_effective_mass( cartesian_k_points, eigenvalues ):
# def least_squares_effective_mass( cartesian_k_points, eigenvalues ):
#     def __init__( self, spin=1, negative_occupancies='warn' ):
#     def occupancy( self ):
#     def __add__( self, other ):
#     def parse_projections( self ):
#     def parse_k_points( self ):
#     def parse_bands( self ):
#     def sanity_check( self ):
#     def from_files( cls, filenames, **kwargs ):
#     def from_file( cls, filename, negative_occupancies='warn',
#                    select_zero_weighted_k_points=False ):
#     def read_from_file( self, filename ):
#     def _read_from_file( self, filename ):
#     def number_of_k_points( self ):
#     def number_of_bands( self ):
#     def spin_channels( self ):
#     def number_of_ions( self ):
#     def number_of_projections( self ):
#     def print_weighted_band_structure( self, spins=None, ions=None, orbitals=None, scaling=1.0, e_fermi=0.0, reciprocal_lattice=None ):
#     def weighted_band_structure( self, spins=None, ions=None, orbitals=None, scaling=1.0, e_fermi=0.0, reciprocal_lattice=None ):
#     def effective_mass_calc( self, k_point_indices, band_index, reciprocal_lattice, spin=1, printing=False ):
#     def x_axis( self, reciprocal_lattice=None ):
#     def bands( self ):
#     def k_points( self ):
#     def select_bands_by_kpoint( self, band_indices ):
#     def select_k_points( self, band_indices ):
# 
# Path: vasppy/outcar.py
# def reciprocal_lattice_from_outcar( filename ): # from https://github.com/MaterialsDiscovery/PyChemia
#     """
#     Finds and returns the reciprocal lattice vectors, if more than
#     one set present, it just returns the last one.
#     Args:
#         filename (Str): The name of the outcar file to be read
# 
#     Returns:
#         List(Float): The reciprocal lattice vectors.
#     """
#     outcar = open(filename, "r").read()
#     # just keeping the last component
#     recLat = re.findall(r"reciprocal\s*lattice\s*vectors\s*([-.\s\d]*)",
#                         outcar)[-1]
#     recLat = recLat.split()
#     recLat = np.array(recLat, dtype=float)
#     # up to now I have, both direct and rec. lattices (3+3=6 columns)
#     recLat.shape = (3, 6)
#     recLat = recLat[:, 3:]
#     return recLat


# Path: vasppy/scripts/fat_bands.py
from vasppy import procar
from vasppy.outcar import reciprocal_lattice_from_outcar
import argparse


def orbitals_with_l( l ):
    to_return = { 's' : [ 0 ],
                  'p' : [ 1, 2, 3 ],
                  'd' : [ 4, 5, 6, 7, 8 ],
                  'f' : [ 9, 10, 11, 12, 13 ],
                  'all' : None }
    return to_return[ l ]
 
def main():
    parser = argparse.ArgumentParser()
    parser.add_argument( '-i', '--ions', help='ion indices for band projection (default: sum over all ions)', nargs='+', type=int )
    parser.add_argument( '-s', '--spins', help='spin indices for band projection (default [ 1 ])', nargs='+', type=int, default=[ 1 ] )
    parser.add_argument( '-o', '--orbitals', help='orbital indices for band projection (default: sum over all orbitals)', nargs='+', type=int )
    parser.add_argument( '-e', '--efermi', help='set fermi energy as reference for energy scale', type=float, default=0.0 )
    parser.add_argument( '-l', '--l-angular-momentum', help='select all orbitals with angular momentum L for band projection. This supercedes the --orbitals option', choices=[ 's', 'p', 'd', 'f', 'all' ] )
    parser.add_argument( '-f', '--procar', help='PROCAR filename (default PROCAR)', type=str, default='PROCAR' )
    parser.add_argument( '--scaling', help='Energy scaling for band widths (default 0.2 eV)', type=float, default=0.2 )
    parser.add_argument( '-x', '--xscaling', help='Automatic scaling of x-axis using reciprocal lattice vectors read from OUTCAR', action='store_true', default=False )
    args = parser.parse_args()

    if args.l_angular_momentum:
        args.orbitals = orbitals_with_l( args.l_angular_momentum )

    if args.xscaling:
        reciprocal_lattice = reciprocal_lattice_from_outcar( 'OUTCAR' ) # Move reading the reciprocal lattice to procar.py
    else:
        reciprocal_lattice = None

    pcar = procar.Procar()

    pcar.read_from_file( args.procar )