# -*- coding: utf-8 -*-
###############################################################################
# Copyright (c), Forschungszentrum Jülich GmbH, IAS-1/PGI-1, Germany.         #
#                All rights reserved.                                         #
# This file is part of the Masci-tools package.                               #
# (Material science tools)                                                    #
#                                                                             #
# The code is hosted on GitHub at https://github.com/judftteam/masci-tools    #
# For further information on the license, see the LICENSE.txt file            #
# For further information please visit http://www.flapw.de or                 #
#                                                                             #
###############################################################################

"""
Io routines for band structure files
"""
from __future__ import absolute_import
from __future__ import print_function
ry_to_ev = 13.6056917253
hartree_to_ev = ry_to_ev * 2.

def read_fleur_banddos_hdf(filepath):
    """
    Reads in the banddos.hdf file from the FLEUR code
    
    returns a dictionary containing all datasets with multidim numpy arrays 
    and also containing the attributes of the groups 
    
    :param filepath: path to the banddos.hdf file
    
    :return: xcoord, bands, xlabels, band_character, band_char_label, 
    kpoints, weights, rep_cell, cell, positions, atomicnumbers, special_point_pos 
    
    
    Expected file content:
    datasets
    [u'bravaisMatrix',
     u'numFoundEigenvals',
     u'specialPointIndices',
     u'lLikeCharge',
     u'positions',
     u'atomicNumbers',
     u'coordinates',
     u'weights',
     u'reciprocalCell',
     u'eigenvalues',
     u'specialPointLabels',
     u'equivAtomsGroup']
    
    attributes:
    {u'lastFermiEnergy': array([0.20852455]),
     u'maxL': array([3], dtype=int32),
     u'nAtoms': array([2], dtype=int32),
     u'nSpecialPoints': array([7], dtype=int32),
     u'nTypes': array([1], dtype=int32),
     u'neigd': array([19], dtype=int32),
     u'nkpt': array([100], dtype=int32),
     u'spins': array([1], dtype=int32),
     u'version': array([1], dtype=int32)}
    """
    import numpy as np
    from masci_tools.io.io_hdf5 import read_hdf
    
    xcoord, bands, xlabels, band_character, band_char_label = [],[],[],[],[]
    
    # read in file
    datasets, groups, group_attrs = read_hdf(filepath)
    
    # correct eigenvalues with Fermi energy
    fermien_htr = group_attrs.read(u'lastFermiEnergy', None)
    
    eig = datasets['eigenvalues']
    print((len(eig), len(eig[-1]), len(eig[-1][-1])))
    bands_tmp = [eig[0].transpose()]
    if len(eig)==2: # wo spins
        bands_tmp.append(eig[1].transpose())
    
    for i, spinband in enumerate(bands_tmp):
        bands_s = []
        for j, band in enumerate(spinband):
            bands_s.append((np.array(band) - fermien_htr)*hartree_to_ev)
        bands.append(bands_s)
    
    print((len(bands)))
    # get special points
    nspecial_labels = group_attrs.read(u'nSpecialPoints', [0])[0]
    if nspecial_labels > 0:
        special_point_labels = list(datasets.read(u'specialPointLabels', []))
        special_point_pos = list(datasets.read(u'specialPointIndices', []))
        
        for i, special_label in enumerate(special_point_labels):
            label = special_label.rstrip()
            if label == u'g':
                label = u'Gamma'
            special_point_labels[i] = label
    else:
        special_point_pos = None
    
    # construct kpoint path
    kpoints = datasets.read('coordinates', None)
    if kpoints is None:
        return None
    
    rep_cell = datasets['reciprocalCell']
    repc = np.array(rep_cell)
    xc = 0.0 
    xcoord.append(xc)
    for i, kpt in enumerate(kpoints[1:]):
        if nspecial_labels > 0:
            if i+1 in special_point_pos:
                index = special_point_pos.index(i+1)
                xlabels.append([special_point_labels[index], xc])
        # get the right length for the path
        diff = np.matmul(repc,kpoints[i]) - np.matmul(repc,kpt)
        abst = np.sqrt(diff[0]**2+diff[1]**2+diff[2]**2)
        xc = xc + abst
        xcoord.append(xc)

    
    band_character = datasets['lLikeCharge'] # for each [kpoint (s,p,d,f), TODO, check because is optiona; output
    
    cell = datasets['bravaisMatrix']
    positions = datasets['positions']
    atomicnumbers = datasets['atomicNumbers']
    weights = datasets['weights']
    
    return xcoord, bands, xlabels, band_character, band_char_label, kpoints, weights, rep_cell, cell, positions, atomicnumbers, special_point_pos
    
