"""
Parsing for .xsf files. The goal is to extract volumetric data.

The .xsf file format offers a lot of information, we are only interested in
the volumetric data for now.
"""

import numpy as np


def read_xsf(filename):
    """
    Read xsf file into numpy array.

    Parameters
    ----------
    filename : string
        filename of xsf file.

    Returns
    -------
    data : numpy.array
        Data from xsf file.

    meta : dict
        Metadata from xsf file - empty for now, and provided for compatibility
        reasons.
    """
    f = open(filename)
    lines = f.readlines()

    found_datagrid = None
    first_data_line = None
    number_data_lines = None
    xiter = yiter = ziter = 0

    for idx, line in enumerate(lines):
        if found_datagrid is None:
            if "BEGIN_BLOCK_DATAGRID_3D" in line:
                found_datagrid = idx
                code = lines[idx+1].strip()

                # The specific formatting may, similar to .cube files.
                # So better to be specific.
                if code != "3D_PWSCF":
                    raise Exception("This .xsf parser can only read .xsf files"
                                    " generated by Quantum ESPRESSO")
        else:
            if idx == found_datagrid + 3:
                grid_dimensions = [int(line.split()[0]), int(line.split()[1]),
                                   int(line.split()[2])]
                data = np.zeros(grid_dimensions, dtype=np.float64)

                # Quantum ESPRESSO writes with 6 entries per line.
                number_data_lines = int(np.ceil(np.prod(grid_dimensions) / 6))
                last_entry = int(np.prod(data) % 6)
                first_data_line = found_datagrid + 8

        if first_data_line is not None:
            if first_data_line <= idx < number_data_lines+first_data_line:
                dataline = line.split()
                if idx == number_data_lines+first_data_line-1:
                    number_entries = last_entry
                else:
                    number_entries = 6

                # Each line contains 6 entries, except the last.
                for i in range(0, number_entries):
                    data[xiter, yiter, ziter] = float(dataline[i])
                    xiter += 1
                    if xiter == grid_dimensions[0]:
                        xiter = 0
                        yiter += 1
                    if yiter == grid_dimensions[1]:
                        yiter = 0
                        ziter += 1
    return data, {}
