def parse_kline(line, ik=None):
    from qharv.reel import ascii_out
    assert 'k(' in line
    ikt, kvect, wkt = line.split('=')
    myik = int(ascii_out.lr_mark(ikt, '(', ')'))
    if ik is not None:  # check k index
        assert ik == myik - 1  # fortran 1-based indexing
    wk = float(wkt)
    klist = ascii_out.lr_mark(kvect, '(', ')').split()
    kvec = np.array(klist, dtype=float)
    return kvec, wk
def read_out_cell(scf_out, ndim=3):
    axes = np.zeros([ndim, ndim])
    from qharv.reel import ascii_out
    mm = ascii_out.read(scf_out)
    idx = mm.find(b'crystal axes')
    mm.seek(idx)
    mm.readline()
    for idim in range(ndim):
        line = mm.readline()
        right = line.split('=')[-1]
        text = ascii_out.lr_mark(right, '(', ')')
        axes[idim, :] = map(float, text.split())
    return axes
def get_supertwists(qmc_out):
    """ read supercell twists from QMCPACK output

  Args:
    qmc_out (str): QMCPACK output, must contain "Super twist #"
  Return:
    np.array: an array of twist vectors (ntwist, ndim)
  """
    from qharv.reel import ascii_out
    mm = ascii_out.read(qmc_out)
    idxl = ascii_out.all_lines_with_tag(mm, 'Super twist #')
    lines = ascii_out.all_lines_at_idx(mm, idxl)
    data = []
    for line in lines:
        text = ascii_out.lr_mark(line, '[', ']')
        vec = np.array(text.split(), dtype=float)
        data.append(vec)
    mat = np.array(data)
    return mat
def retrieve_occupations(nscf_outfile, max_nbnd_lines=10):
    """ read the eigenvalues and occupations of DFT orbitals at every available kpoint in an non-scf output produced by pwscf """
    from qharv.reel import ascii_out
    span = 7

    def scanf_7f(line, n):
        """ implement scanf("%7.*f") """
        numl = []
        for i in range(n):
            token = line[span * i:span * (i + 1)]
            num = float(token)
            numl.append(num)
        return numl

    fhandle = open(nscf_outfile, 'r+')
    mm = mmap(fhandle.fileno(), 0)

    # read number of k points
    nk_prefix = b"number of k points="
    idx = mm.find(nk_prefix)
    mm.seek(idx)

    nk_line = mm.readline()
    nk = int(nk_line.strip(nk_prefix).split()[0])

    # skip to the end of band structure calculation
    idx = mm.find(b'End of self-consistent calculation')
    idx = mm.find(b'End of band structure calculation')
    mm.seek(idx)

    # read the eigenvalues and occupations at each kpoint
    kpt_prefix = "k ="
    data = []
    for ik in range(nk):
        idx = mm.find(kpt_prefix.encode())
        mm.seek(idx)
        kpt_line = mm.readline()
        kxkykz = ascii_out.lr_mark(kpt_line, '=', '(')
        kpt = scanf_7f(kxkykz, 3)

        mm.readline()  # skip empty line
        eval_arr = np.array([])
        for iline in range(max_nbnd_lines):
            tokens = mm.readline().split()
            if len(tokens) == 0:
                break
            # end if
            eval_arr = np.append(eval_arr, map(float, tokens))
        # end for iline

        idx = mm.find(b'occupation numbers')
        mm.seek(idx)
        mm.readline()  # skip current line
        occ_arr = np.array([])
        for iline in range(100):
            tokens = mm.readline().split()
            if len(tokens) == 0:
                break
            # end if
            occ_arr = np.append(occ_arr, map(float, tokens))
        # end for iline

        entry = {
            'ik': ik,
            'kpt': list(kpt),
            'eval': list(eval_arr),
            'occ': list(occ_arr)
        }
        data.append(entry)
    # end for
    mm.close()
    fhandle.close()
    return data
def parse_nscf_bands(nscf_out, span=7, trailer='occupation numbers'):
    data = {}  # build a dictionary as return value

    def scanf_7f(line, n):
        """ implement scanf("%7.*f") """
        numl = []
        for i in range(n):
            token = line[span * i:span * (i + 1)]
            num = float(token)
            numl.append(num)
        return numl

    def parse_float_body(body):
        """ parse a blob of floats """
        lines = body.split('\n')
        numl = []
        for line in lines:
            if len(line) == 0: continue
            numl += map(float, line.split())
        return numl

    from qharv.reel import ascii_out
    ndim = 3
    mm = ascii_out.read(nscf_out)
    alat = ascii_out.name_sep_val(mm, 'lattice parameter (alat)')
    blat = 2 * np.pi / alat

    # find the beginnings of each band
    bhead = ' k ='
    idxl = ascii_out.all_lines_with_tag(mm, bhead)
    nkpt = len(idxl)
    data['nkpt'] = nkpt

    # estimate the end of the last band
    idx1 = ascii_out.all_lines_with_tag(mm, trailer)[-1]

    # trick to use no if statement in the loop
    idxl = idxl + [idx1]

    kvecs = []  # (nkpt, ndim)
    mat = []  # (nkpt, nbnd)
    for ikpt in range(nkpt):
        # specify beginning and end of the band output
        idx0 = idxl[ikpt]
        idx1 = idxl[ikpt + 1]

        # parse band output
        #  first read header
        mm.seek(idx0)
        header = mm.readline().decode()
        if not 'bands (ev)' in header: continue
        kxkykz = ascii_out.lr_mark(header, '=', '(')
        kvec = scanf_7f(kxkykz, ndim)
        kvecs.append(kvec)
        #  then read body
        body = mm[mm.tell():idx1].decode().strip('\n')
        if trailer in body:
            idx2 = mm.find(trailer.encode())
            body = mm[mm.tell():idx2].strip('\n')
        row = parse_float_body(body)
        mat.append(row)
    # end for ikpt
    data['kvecs'] = blat * np.array(kvecs)
    data['bands'] = np.array(mat)
    return data
def read_sym_ops(scf_out, ndim=3):
    """ read symmetry operators

  Args:
    scf_out (str): QE output file
    ndim (int, optional): number of spatial dimensions, default is 3
  Return:
    list: all symmetry operators, each is represented as a dictionary
      isym is index, name is description, vec is shift, mat is rotation
  """
    from qharv.reel import ascii_out
    mm = ascii_out.read(scf_out)

    # find starting location of symmetry operator output
    idx = mm.find(b'Sym. Ops.')
    if idx == -1:
        msg = 'no symmetry operations printed in %s. Is verbosity high?' % scf_out
        raise RuntimeError(msg)
    # rewind to beginning of line
    idx0 = mm.rfind(b'\n', 0, idx)
    mm.seek(idx0 + 1)
    header = mm.readline().decode()
    nsym = int(header.split()[0])

    # check the number of symmetry outputs
    idxl = ascii_out.all_lines_with_tag(mm, 'isym = ')
    if len(idxl) != nsym:
        raise RuntimeError('found %d symm. expected %d' % (len(idxl), nsym))

    # parse symmetry operators
    symops = []
    for idx in idxl:
        mm.seek(idx)

        # read symmetry index and name: isym, name
        line0 = mm.readline().decode()
        text0 = line0.split('=')[1]
        tokens0 = text0.split()
        isym = int(tokens0[0])
        name = ' '.join(tokens0[1:])

        # read translation vector: vec
        vec = [0] * ndim
        if 'cart. axis' in name:
            vect = ascii_out.lr_mark(line0, '[', ']')
            vec[:] = list(map(float, vect.split(',')))

        # read rotation matrix: mat
        mat = []
        idx = mm.find(b'cryst.')
        mm.readline()  # skip empty line
        for idim in range(ndim):
            line = mm.readline().decode()
            if 'cryst.' in line:
                line = line.split('=')[1]
            text = ascii_out.lr_mark(line, '(', ')')
            mat.append(list(map(float, text.split())))
        entry = {'isym': isym, 'name': name, 'vec': vec, 'mat': mat}
        symops.append(entry)
    mm.close()
    return symops