# Copyright (c) 2013 Anton Kozhevnikov, Thomas Schulthess
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without modification, are permitted provided that
# the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this list of conditions and the
#    following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions
#    and the following disclaimer in the documentation and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED
# WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
# PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR
# ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR
# OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import json
import sys
import re
import io


def warning(*args, **kwargs):
    print(*args, file=sys.stderr, **kwargs)


def read_until(fin, tag):
    while True:
        line = fin.readline()
        if not line:
            warning("Unexpected end of file")
            sys.exit(-1)
        if tag in line:
            return


def read_mesh_data(in_file, npoints):
    out_data = []
    while True:
        s = in_file.readline().split()
        for k in range(len(s)):
            out_data.append(float(s[k]))
        if len(out_data) == npoints:
            break
    return out_data


def parse_header(upf_dict, upf_str):
    upf_dict["header"] = {}
    upf = io.StringIO(upf_str)
    read_until(upf, "<PP_HEADER>")

    s = upf.readline().split()
    # upf_dict["header"]["version"] = int(s[0]);

    s = upf.readline().split()
    upf_dict["header"]["element"] = s[0].strip()

    s = upf.readline().split()
    upf_dict["header"]["pseudo_type"] = s[0]

    s = upf.readline().split()
    upf_dict["header"]["core_correction"] = (s[0][0] == "T" or s[0][0] == "t") or 0

    s = upf.readline()
    # upf_dict["header"]["dft"] = s[:20]

    s = upf.readline().split()
    upf_dict["header"]["z_valence"] = float(s[0])

    s = upf.readline().split()
    # upf_dict["header"]["etotps"] = float(s[0])

    s = upf.readline().split()
    # upf_dict["header"]["ecutwfc"] = float(s[0])
    # upf_dict["header"]["ecutrho"] = float(s[1])

    s = upf.readline().split()
    upf_dict["header"]["l_max"] = int(s[0])

    s = upf.readline().split()
    upf_dict["header"]["mesh_size"] = int(s[0])

    s = upf.readline().split()
    upf_dict["header"]["number_of_wfc"] = int(s[0])
    upf_dict["header"]["number_of_proj"] = int(s[1])

    upf.readline()

    # upf_dict["header"]["wavefunctions"] = []
    # for i in range(upf_dict["header"]["nwfc"]):
    #    upf_dict["header"]["wavefunctions"].append({})
    #    s = upf.readline().split()
    #    upf_dict["header"]["wavefunctions"][i]["els"] = s[0]
    #    upf_dict["header"]["wavefunctions"][i]["lchi"] = int(s[1])
    #    upf_dict["header"]["wavefunctions"][i]["oc"] = float(s[2])

    upf.close()


#
# QE subroutine read_pseudo_mesh
#
def parse_mesh(upf_dict, upf_str):
    upf = io.StringIO(upf_str)
    read_until(upf, "<PP_R>")

    # =================================================================
    # call scan_begin (iunps, "R", .false.)
    # read (iunps, *, err = 100, end = 100) (upf%r(ir), ir=1,upf%mesh )
    # call scan_end (iunps, "R")
    # =================================================================
    upf_dict["radial_grid"] = read_mesh_data(upf, upf_dict["header"]["mesh_size"])

    read_until(upf, "<PP_RAB>")

    # ===================================================================
    # call scan_begin (iunps, "RAB", .false.)
    # read (iunps, *, err = 101, end = 101) (upf%rab(ir), ir=1,upf%mesh )
    # call scan_end (iunps, "RAB")
    # ===================================================================
    # upf_dict["mesh"]["rab"] = read_mesh_data(upf, upf_dict["header"]["nmesh"])

    upf.close()


#
# QE subroutine read_pseudo_nlcc
#
def parse_nlcc(upf_dict, upf_str):
    upf = io.StringIO(upf_str)
    read_until(upf, "<PP_NLCC>")

    # =======================================================================
    # ALLOCATE( upf%rho_atc( upf%mesh ) )
    # read (iunps, *, err = 100, end = 100) (upf%rho_atc(ir), ir=1,upf%mesh )
    # =======================================================================
    upf_dict["core_charge_density"] = read_mesh_data(
        upf, upf_dict["header"]["mesh_size"]
    )

    upf.close()


#
# QE subroutine read_pseudo_local
#
def parse_local(upf_dict, upf_str):
    upf_dict["local_potential"] = []
    upf = io.StringIO(upf_str)

    read_until(upf, "<PP_LOCAL>")

    # =================================================================
    # ALLOCATE( upf%vloc( upf%mesh ) )
    # read (iunps, *, err=100, end=100) (upf%vloc(ir) , ir=1,upf%mesh )
    # =================================================================
    vloc = read_mesh_data(upf, upf_dict["header"]["mesh_size"])
    upf_dict["local_potential"] = [v / 2 for v in vloc]

    upf.close()


#
# QE subroutine read_pseudo_nl
#
def parse_non_local(upf_dict, upf_str):
    upf_dict["beta_projectors"] = []
    upf_dict["D_ion"] = []

    upf = io.StringIO(upf_str)
    read_until(upf, "<PP_NONLOCAL>")

    # ======================================================================
    #   do nb = 1, upf%nbeta
    #      call scan_begin (iunps, "BETA", .false.)
    #      read (iunps, *, err = 100, end = 100) idum, upf%lll(nb), dummy
    #      read (iunps, *, err = 100, end = 100) ikk
    #      upf%kbeta(nb) = ikk
    #      upf%kkbeta = MAX ( upf%kkbeta, upf%kbeta(nb) )
    #      read (iunps, *, err = 100, end = 100) (upf%beta(ir,nb), ir=1,ikk)
    #
    #      read (iunps, *, err=200,iostat=ios) upf%rcut(nb), upf%rcutus(nb)
    #      read (iunps, *, err=200,iostat=ios) upf%els_beta(nb)
    #      call scan_end (iunps, "BETA")
    # 200  continue
    #   enddo
    # ======================================================================
    for i in range(upf_dict["header"]["number_of_proj"]):
        read_until(upf, "<PP_BETA>")
        upf_dict["beta_projectors"].append({})
        upf_dict["beta_projectors"][i]["label"] = ""
        s = upf.readline().split()
        upf_dict["beta_projectors"][i]["angular_momentum"] = int(s[1])
        s = upf.readline()
        # upf_dict['beta_projectors'][i]['cutoff_radius_index'] = int(s)
        nr = int(s)
        beta = read_mesh_data(upf, nr)

        upf_dict["beta_projectors"][i][
            "radial_function"
        ] = beta  # read_mesh_data(upf, upf_dict['beta_projectors'][i]['cutoff_radius_index'])

        upf_dict["beta_projectors"][i]["cutoff_radius"] = 0.0
        upf_dict["beta_projectors"][i]["ultrasoft_cutoff_radius"] = 0.0

        # line = upf.readline()
        # if not "</PP_BETA>" in line:
        # s = line.split()
        # upf_dict['beta_projectors'][i]['cutoff_radius'] = float(s[0])
        # upf_dict['beta_projectors'][i]['ultrasoft_cutoff_radius'] = float(s[1])
        # s = upf.readline().split()
        # if not "</PP_BETA>" in line: upf_dict['beta_projectors'][i]['els_beta'] = float(s[0])

    # ================================================================
    # call scan_begin (iunps, "DIJ", .false.)
    # read (iunps, *, err = 101, end = 101) nd, dummy
    # do icon = 1, nd
    #    read (iunps, *, err = 101, end = 101) nb, mb, upf%dion(nb,mb)
    #    upf%dion (mb,nb) = upf%dion (nb,mb)
    # enddo
    # call scan_end (iunps, "DIJ")
    # ================================================================
    nb = upf_dict["header"]["number_of_proj"]
    dij = [0 for i in range(nb * nb)]
    read_until(upf, "<PP_DIJ>")
    s = upf.readline().split()
    nd = int(s[0])
    for k in range(nd):
        s = upf.readline().split()
        i = int(s[0]) - 1
        j = int(s[1]) - 1
        dij[i * nb + j] = float(s[2]) / 2  # convert to Hartree
        dij[j * nb + i] = float(s[2]) / 2  # convert to Hartree

    upf_dict["D_ion"] = dij

    if upf_dict["header"]["pseudo_type"] != "US":
        return

    upf_dict["augmentation"] = []

    # =============================================
    # call scan_begin (iunps, "QIJ", .false.)
    # read (iunps, *, err = 102, end = 102) upf%nqf
    # upf%nqlc = 2 * upf%lmax  + 1
    # =============================================
    read_until(upf, "<PP_QIJ>")
    s = upf.readline().split()
    num_q_coef = int(s[0])
    # nqlc = upf_dict['header']['l_max'] * 2 + 1

    # =======================================================================
    # if ( upf%nqf /= 0) then
    #    call scan_begin (iunps, "RINNER", .false.)
    #    read (iunps,*,err=103,end=103) ( idum, upf%rinner(i), i=1,upf%nqlc )
    #    call scan_end (iunps, "RINNER")
    # end if
    # =======================================================================
    if num_q_coef != 0:
        R_inner = []
        read_until(upf, "<PP_RINNER>")
        for i in range(2 * upf_dict["header"]["l_max"] + 1):
            s = upf.readline().split()
            R_inner.append(float(s[1]))
        read_until(upf, "</PP_RINNER>")

    # ==================================================================================
    # do nb = 1, upf%nbeta
    #    do mb = nb, upf%nbeta
    #       read (iunps,*,err=102,end=102) idum, idum, ldum, dummy
    #       !"  i    j   (l)"
    #       if (ldum /= upf%lll(mb) ) then
    #         call errore ('read_pseudo_nl','inconsistent angular momentum for Q_ij', 1)
    #       end if
    #       read (iunps,*,err=104,end=104) upf%qqq(nb,mb), dummy
    #       ! "Q_int"
    #       upf%qqq(mb,nb) = upf%qqq(nb,mb)
    #       ! ijv is the combined (nb,mb) index
    #       ijv = mb * (mb-1) / 2 + nb
    #       IF (upf%q_with_l .or. upf%tpawp) THEN
    #          l1=upf%lll(nb)
    #          l2=upf%lll(mb)
    #          DO l=abs(l1-l2),l1+l2
    #             read (iunps, *, err=105, end=105) (upf%qfuncl(n,ijv,l), &
    #                                                n=1,upf%mesh)
    #          END DO
    #       ELSE
    #          read (iunps, *, err=105, end=105) (upf%qfunc(n,ijv), n=1,upf%mesh)
    #       ENDIF
    # ==================================================================================

    nb = upf_dict["header"]["number_of_proj"]
    l_max = upf_dict["header"]["l_max"]

    for i in range(nb):
        li = upf_dict["beta_projectors"][i]["angular_momentum"]
        for j in range(i, nb):
            lj = upf_dict["beta_projectors"][j]["angular_momentum"]

            s = upf.readline().split()
            if int(s[2]) != lj:
                warning("inconsistent angular momentum")
                sys.exit(-1)

            if int(s[0]) != i + 1 or int(s[1]) != j + 1:
                print("inconsistent ij indices")
                sys.exit(-1)

            s = upf.readline().split()

            qij = read_mesh_data(upf, upf_dict["header"]["mesh_size"])

            # ======================================================================
            # if ( upf%nqf > 0 ) then
            #   call scan_begin (iunps, "QFCOEF", .false.)
            #   read (iunps,*,err=106,end=106) &
            #             ( ( upf%qfcoef(i,lp,nb,mb), i=1,upf%nqf ), lp=1,upf%nqlc )
            #   do i = 1, upf%nqf
            #      do lp = 1, upf%nqlc
            #         upf%qfcoef(i,lp,mb,nb) = upf%qfcoef(i,lp,nb,mb)
            #      end do
            #   end do
            #   call scan_end (iunps, "QFCOEF")
            # end if
            # =======================================================================
            if num_q_coef > 0:
                read_until(upf, "<PP_QFCOEF>")

                q_coefs = read_mesh_data(upf, num_q_coef * (2 * l_max + 1))

                read_until(upf, "</PP_QFCOEF>")

            if num_q_coef > 0:
                # constuct Qij(r) for each l
                for l in range(abs(li - lj), li + lj + 1):
                    if (li + lj + l) % 2 == 0:
                        qij_fixed = [
                            qij[k] for k in range(upf_dict["header"]["mesh_size"])
                        ]

                        for ir in range(upf_dict["header"]["mesh_size"]):
                            x = upf_dict["radial_grid"][ir]
                            x2 = x * x
                            if x < R_inner[l]:
                                qij_fixed[ir] = q_coefs[0 + l * num_q_coef]
                                for n in range(1, num_q_coef):
                                    qij_fixed[ir] += (
                                        q_coefs[n + l * num_q_coef] * x2**n
                                    )
                                qij_fixed[ir] *= x ** (l + 2)

                        qij_dict = {}
                        qij_dict["radial_function"] = qij_fixed
                        qij_dict["i"] = i
                        qij_dict["j"] = j
                        qij_dict["angular_momentum"] = l
                        upf_dict["augmentation"].append(qij_dict)

    upf.close()


#
# QE subroutine read_pseudo_pswfc
#
def parse_pswfc(upf_dict, upf_str):
    upf_dict["atomic_wave_functions"] = []

    upf = io.StringIO(upf_str)
    read_until(upf, "<PP_PSWFC>")

    # =======================================================================
    # ALLOCATE( upf%chi( upf%mesh, MAX( upf%nwfc, 1 ) ) )
    # upf%chi = 0.0_DP
    # do nb = 1, upf%nwfc
    #    read (iunps, *, err=100, end=100) dummy  !Wavefunction labels
    #    read (iunps, *, err=100, end=100) ( upf%chi(ir,nb), ir=1,upf%mesh )
    # enddo
    # =======================================================================
    for i in range(upf_dict["header"]["number_of_wfc"]):
        wf = {}
        s = upf.readline().split()
        wf["label"] = s[0]
        wf["angular_momentum"] = int(s[1])
        wf["occupation"] = float(s[2])

        wf["radial_function"] = read_mesh_data(upf, upf_dict["header"]["mesh_size"])
        upf_dict["atomic_wave_functions"].append(wf)
    upf.close()


#
# QE subroutine read_pseudo_rhoatom
#
def parse_rhoatom(upf_dict, upf_str):
    upf = io.StringIO(upf_str)
    read_until(upf, "<PP_RHOATOM>")

    # ================================================================
    # ALLOCATE( upf%rho_at( upf%mesh ) )
    # read (iunps,*,err=100,end=100) ( upf%rho_at(ir), ir=1,upf%mesh )
    # ================================================================
    upf_dict["total_charge_density"] = read_mesh_data(
        upf, upf_dict["header"]["mesh_size"]
    )

    upf.close()


def parse_upf1_from_string(upf1_str):
    upf_dict = {}

    parse_header(upf_dict, upf1_str)
    parse_mesh(upf_dict, upf1_str)
    if upf_dict["header"]["core_correction"] == 1:
        parse_nlcc(upf_dict, upf1_str)
    parse_local(upf_dict, upf1_str)
    parse_non_local(upf_dict, upf1_str)
    parse_pswfc(upf_dict, upf1_str)
    parse_rhoatom(upf_dict, upf1_str)

    pp_dict = {}
    pp_dict["pseudo_potential"] = upf_dict

    return pp_dict


def main():
    with open(sys.argv[1], "r") as fh:
        upf = fh.read()
    pp_dict = parse_upf1_from_string(upf)

    fout = open(sys.argv[1] + ".json", "w")

    # Match comma, space, newline and an arbitrary number of spaces ',\s\n\s*' with the
    # following conditions: a digit before (?<=[0-9]) and a minus or a digit after (?=[-|0-9]).
    # Replace found sequence with comma and space.
    fout.write(
        re.sub(r"(?<=[0-9]),\s\n\s*(?=[-|0-9])", r", ", json.dumps(pp_dict, indent=2))
    )
    fout.close()


if __name__ == "__main__":
    main()
