#!/usr/bin/env python2
#=========================== Wave2spn =============================
#                                                                 #
#    An utility to convert WAVECAR to .spn file for WANNIER90     #
#                                                                 #
#=========================== Wave2spn =============================
#
# Written by Stepan Tsirkin (University of the Basque Country)
#    now at : Iniversity of Zurich
""" An utility to calculate the ``.spn`` file for wannier90 from ``WAVECAR`` file generated by `VASP <https://www.vasp.at/>`_.

    usage : ::

        python3 -m wannierberri.utils.vaspspn   option=value

    Options
        -h
            |  print this help message

        fin
            |  inputfile  name.
            |  default: WAVECAR
        fout
            |  outputfile name
            |  default: wannier90.spn
        IBstart
            |  the first band to be considered (counting starts from 1).
            |  default: 1
        NB
            |  number of bands in the output. If NB<=0 all bands are used.
            |  default: 0
        norm
            |  how to normalize the eigenstates, if they are not perfectly orthonormal
            |  norm=norm  (D) -   normalize each state individually
            |  norm=none      -   do not normalize WFs, take as they are.

"""

from scipy.io import FortranFile
import numpy as np
import datetime

class FortranFileW(FortranFile):

    def __init__(self, filename):
        print("using scipy.io to write")
        super().__init__(filename, mode='w')


def vasp_to_spn(fin='WAVECAR', fout='wannier90.spn', NBout=0, IBstart=1, normalize='norm', formatted=False):

    print(fout)
    print("reading {0}\n writing to {1}".format(fin, fout))

    WAV = open(fin, "rb")
    RECL = 3

    def record(irec, cnt=np.Inf, dtype=float):
        WAV.seek(irec * RECL)
        return np.fromfile(WAV, dtype=dtype, count=min(RECL, cnt))

    RECL, ispin, iprec = [int(x) for x in record(0)]

    print(RECL, ispin, iprec)

    if iprec != 45200: raise RuntimeError('double precision WAVECAR is not supported')
    if ispin != 1: raise RuntimeError('WAVECAR does not contain spinor wavefunctions. ISPIN={0}'.format(ispin))

    NK, NBin = [int(x) for x in record(1, 2)]

    IBstart -= 1
    if IBstart < 0: IBstart = 0
    if NBout <= 0: NBout = NBin
    if NBout + IBstart > NBin:
        print(
            ' WARNING: NB+IBstart-1=', NBout + IBstart,
            ' exceeds the number of bands in WAVECAR NBin=' + str(NBin) + '. We set NBout=' + str(NBin - IBstart))
        NBout = NBin - IBstart

    print(
        "WAVECAR contains {0} k-points and {1} bands.\n Writing {2} bands in the output starting from".format(
            NK, NBin, NBout))

    SPN = FortranFileW(fout)
    header = "Created from wavecar at {0}".format(datetime.datetime.now().isoformat())
    header = header[:60]
    header += " " * (60 - len(header))
    if formatted:
        with open(fout, 'w') as fw:
            fw.write(header)
            fw.write('\n')
            np.savetxt(fw, np.array([NBout, NK], dtype=np.int32))
    else:
        SPN.write_record(bytearray(header, encoding='ascii'))
        SPN.write_record(np.array([NBout, NK], dtype=np.int32))

    for ik in range(NK):
        npw = int(record(2 + ik * (NBin + 1), 1))
        npw12 = npw // 2
        if npw != npw12 * 2:
            raise RuntimeError(f"odd number of coefs {npw}")
        print("k-point {0:3d} : {1:6d} plane waves".format(ik, npw))
        WF = np.zeros((npw, NBout), dtype=complex)
        for ib in range(NBout):
            WF[:, ib] = record(3 + ik * (NBin + 1) + ib + IBstart, npw, np.complex64)
        overlap = WF.conj().T.dot(WF)
        assert np.max(np.abs(overlap - overlap.T.conj())) < 1e-15

        if normalize == "norm":
            WF = WF / np.sqrt(np.abs(overlap.diagonal()))

        SIGMA = np.array(
            [
                [
                    np.einsum("ki,kj->ij",
                              WF.conj()[npw12 * i:npw12 * (i + 1), :], WF[npw12 * j:npw12 * (j + 1), :])
                    for j in (0, 1)
                ] for i in (0, 1)
            ])
        SX = SIGMA[0, 1] + SIGMA[1, 0]
        SY = -1.j * (SIGMA[0, 1] - SIGMA[1, 0])
        SZ = SIGMA[0, 0] - SIGMA[1, 1]
        A = np.array([s[n, m] for m in range(NBout) for n in range(m + 1) for s in (SX, SY, SZ)], dtype=np.complex128)
        if formatted:
            with open(fout, 'a') as fw:
                np.savetxt(fw, A.reshape(1, A.shape[0]))
        else:
            SPN.write_record(A)
