"""ABINIT-MP の入出力処理モジュール。

AJFファイルの生成・読み込み、IFIE/PIEDAの解析、
エネルギー抽出など ABINIT-MP 関連の I/O を提供する。
"""

from __future__ import annotations

import ast
import sys
import os
import math
import re
import tarfile
import glob
import logging
from typing import Any, Callable

from .mol_io import mol_io as mi

logger = logging.getLogger(__name__)
try:
    import collections
except ImportError:
    pass
try:
    import pandas as pd
except ImportError:
    pass
try:
    import itertools
except ImportError:
    pass


class abinit_io(mi):
    """ABINIT-MP の入出力を管理するクラス。

    AJFファイルの生成、フラグメント情報の読み書き、
    FMO計算結果（エネルギー、IFIE等）の解析機能を提供する。
    """

    def __init__(self):
        # print('## load abinit io')
        super().__init__()
        # gen_rand: gen_coord

        # submit param
        self.npro = 4
        self.memory = "1800"
        self.memorymp2 = "0"
        self.queue = None
        self.solv_flag = False
        self.abinitmp_path = 'abinitmp'
        self.mpipath = 'mpirun'

        # abinitmp param
        self.ajf_method = "MP2"
        self.ajf_basis_set = "6-31G*"
        self.abinit_ver = 'rev22' #rev10, 11, 15, 17, 22, 23, mizuho
        self.pbmolrad = 'vdw'
        self.pbcnv = 1.0
        self.piedaflag = True
        self.cpfflag = False
        self.cpfver = 23
        self.cmmflag = False
        self.nofzc_flag = False
        self.bsseflag = False
        self.readgeom = ""
        self.writegeom = ""
        self.submit_system = None
        self.autofrag = False
        self.dgemm = False
        self.nbo = True
        self.resp = True
        self.ligchg = None
        self.rsolv = None
        self.mldatfrag = False
        self.mldatname = None
        self.mllimit = None
        self.disp = False
        self.is_xyz = False
        self.natom = 0
        self.xyzstr = ''
        # distlitname

        # frag table
        self.fatomnums = []
        self.fchgs = []
        self.fbaas = []
        self.fatminfos = []
        self.connects = []

        return

    def chkdepth(self, k: Any) -> int:
        """リストのネスト深度を返す。"""
        if not k:
            return 0
        else:
            if isinstance(k, list):
                return 1 + max(self.chkdepth(i) for i in k)
            else:
                return 0

    def get_fragsection(self) -> str:
        """フラグメント情報からAJFの&FRAGMENTセクション文字列を生成する。"""
        if self.chkdepth(self.fatomnums) == 1:
            frag_atom = [self.fatomnums]
            frag_charge = [self.fchgs]
            frag_connect_num = [self.fbaas]
            frag_connect = [self.connects]
            seg_info = [self.fatminfos]
        else:
            frag_atom = self.fatomnums
            frag_charge = self.fchgs
            frag_connect_num = self.fbaas
            frag_connect = self.connects
            seg_info = self.fatminfos

        ajf_fragment = ''
        # fragment atom
        for i in range(len(frag_atom)):
            icount = 0
            for j in range(len(frag_atom[i])):
                icount += 1
                ajf_fragment += '%8d' % (frag_atom[i][j])
                if icount % 10 == 0:
                    icount = 0
                    ajf_fragment += '\n'
        if icount % 10 != 0:
            ajf_fragment += '\n'

        # fragment charge
        for i in range(len(frag_charge)):
            icount = 0
            for j in range(len(frag_charge[i])):
                icount += 1
                ajf_fragment += '%8d' % (frag_charge[i][j])
                if icount % 10 == 0:
                    icount = 0
                    ajf_fragment += '\n'
        if icount % 10 != 0:
            ajf_fragment += '\n'

        # fragment connect num
        for i in range(len(frag_connect_num)):
            icount = 0
            for j in range(len(frag_connect_num[i])):
                icount += 1
                ajf_fragment += '%8d' % (frag_connect_num[i][j])
                if icount % 10 == 0:
                    icount = 0
                    ajf_fragment += '\n'
        if icount % 10 != 0:
            ajf_fragment += '\n'

        # fragment body
        atom_count = 0
        # i: mol j: frag, k:frag atom
        for i in range(len(seg_info)):
            for j in range(len(seg_info[i])):
                icount = 0
                for k in range(len(seg_info[i][j])):
                    icount += 1
                    ajf_fragment += '%8d' % (seg_info[i][j][k] + atom_count)
                    # print(seg_info[i][j][k] + atom_count, 'icount', icount)
                    if icount % 10 == 0:
                        icount = 0
                        ajf_fragment += '\n'

                if icount % 10 != 0:
                    ajf_fragment += '\n'
            atom_count += sum(frag_atom[0])

        atom_count = 0
        # connect info
        for i in range(len(frag_connect)):
            for j in range(len(frag_connect[i])):
                ajf_fragment += '%8d' % (frag_connect[i][j][0] + atom_count)
                ajf_fragment += '%8d' % (frag_connect[i][j][1] + atom_count)
                if len(frag_connect[i][j]) == 3:
                    ajf_fragment += '%8d' % (frag_connect[i][j][2])
                ajf_fragment += '\n'
            atom_count += sum(frag_atom[0])
        return ajf_fragment


    def gen_ajf_bodywrap(self, inname: str) -> str:
        """AJF本体部分を生成するラッパーメソッド。

        Args:
            inname: 入力ファイル名（拡張子付き）。writegeom未設定時のCPFファイル名生成に使用。

        Returns:
            str: AJFファイルの本体文字列。
        """
        # i, j ,k (mol, frag, component_perfrag) dimension
        if self.autofrag :
            ajf_fragment = ""
        else:
            ajf_fragment = self.get_fragsection()[:-1]

        ajf_charge = sum(self.fchgs)
        ajf_parameter = [ajf_charge, ajf_fragment, len(self.fatomnums), 0]

        if self.writegeom == "":
            self.writegeom = "'" + os.path.splitext(inname)[0] + '-' + self.ajf_method + '-' + self.ajf_basis_set.replace('*', 'd') + ".cpf'"

        ajf_body = self.gen_ajf_body(ajf_parameter)

        return ajf_body


    def gen_ajf_body(self, param: list[Any], fzctemp: str = "YES") -> str:
        """AJFファイルの全セクション（&CNTRL, &SCF, &MP2等）を生成する。

        Args:
            param: [電荷, フラグメント文字列, フラグメント数, MOMEフラグ] のリスト。
            fzctemp: FZC（frozen core）設定。デフォルトは'YES'。

        Returns:
            str: AJFファイル本体の文字列。
        """
        ajf_charge, ajf_fragment, num_fragment, MOME_flag = param

        if self.autofrag:
            auto = 'ON'
        else:
            auto = 'OFF'

        if not self.is_xyz:
            self.natom = 0

        if MOME_flag:
            frag_section = """
NF= """ + str(num_fragment) + """
LMOTYP='MOME' """
        else:
            frag_section = """
NF= """ + str(num_fragment) + """
LMOTYP='ANO' """

        if num_fragment == 1:
            use_FMO = 'OFF'
            frag_section = ''
        else:
            use_FMO = 'ON'

        if self.submit_system in ['K', 'OFP', 'Fugaku']:
            np = '1'
        else:
            np = str(int(self.npro/self.para_job))

        # CNTRL NAMELIST
        ajf_body = """&CNTRL
ElecState='S1'
Method='""" + str(self.ajf_method) + """'
Nprint=3
Memory=""" + str(self.memory) + """
Natom=""" + str(self.natom) + """
Charge=""" + str(ajf_charge) + "\n"
        if not self.is_xyz:
            ajf_body += """ReadGeom='""" + str(self.readgeom) + "'\n"
        if self.cpfflag:
            ajf_body += "WriteGeom=" + str(self.writegeom) + "\n"
            if self.cpfver != 23:
                ajf_body += "CPFVER=" + str(self.cpfver) + "\n"
        if self.mldatfrag:
            ajf_body += "WriteMLdata=" + str(self.mldatname) + "\n"
        if self.mllimit not in [0, None]:
            ajf_body += "MLfraglimit=" + str(self.mllimit) + "\n"

        ajf_body += """Gradient='NO'
Vector='OFF'
CPFBIN='NO'
THOVL=1.0E-12
E_THSWZ=1.0E-12
G_THSWZ=1.0E-12
/

&FMOCNTRL
FMO='""" + str(use_FMO) + """'
NBody=2
AutoFrag='"""+ auto + """' """ + str(frag_section) + """
Laoc=0.0
Lptc=2.0
esp_ptc_multipole='NO'
Ldimer=2.0
NP=""" + np + """
MaxSCCcyc=250
MaxSCCenergy=5.0E-7
"""
        if  self.ligchg is not None:
            ligchgstr = "LigandCharge='"
            for i in range(len(self.ligchg)):
                ligchgstr += self.ligchg[i][0] + '=' + self.ligchg[i][1]
                if i == len(self.ligchg) - 1:
                    ligchgstr += "'"
                else:
                    ligchgstr += ','

            logger.debug(ligchgstr)
            ajf_body += ligchgstr + """
"""

        if  self.rsolv is not None:
            rsolvstr = "RSolv='"
            for i in range(len(self.rsolv)):
                rsolvstr += self.rsolv[i][0] + '=' + self.rsolv[i][1]
                if i == len(self.rsolv) - 1:
                    rsolvstr += "'"
                else:
                    rsolvstr += ','

            logger.debug(rsolvstr)
            ajf_body += rsolvstr + """
"""

        if self.cmmflag:
            ajf_body += """Dimer_es_multipole='YES'
Ldimer_cmm=5.0
/
"""
        else:
            ajf_body += """/
"""

        ajf_body += """
&SCF
MaxSCFenergy=1.0E-8
MaxSCFdensity=1.0E-6
MaxSCFcyc=150
DIISTYPE='C2_OLD'
THINTEG=1.0E-12
IFCD='NO'
/

&BASIS
BasisSet='""" + str(self.ajf_basis_set) + """'
DiffuseOn='NO'
/

&OPTCNTRL
OPT='OFF'
/

"""
        if self.abinit_ver in ['rev22', 'rev23', 'v2rev4', 'v2rev8']:
            ajf_body += """
&MFMO
/
"""

        if self.abinit_ver in ['v2rev8']:
            ajf_body += """
&EXPROP
/
"""

        ajf_body += """&SCZV
DimerResponseTerm='NO'
/
"""
        if self.abinit_ver in ['rev22', 'rev23', 'v2rev4', 'v2rev8']:
            ajf_body += """
&XUFF
/
"""

        ajf_body += """
&MP2"""

        if self.dgemm:
            ajf_body += """
MOD1ST='GEMM'
MOD2ND='GEMM'
MOD3RD='GEMM'
MOD4TH='GEMM'"""

        ajf_body += """
NP_MP2_IJ=1
NP_MP2_S=0
MemoryMP2=""" + str(self.memorymp2) + """
IFSCS='YES'
IFPRNM='YES'
OSSCAL=1.0
PSSCAL=1.0
NBODY=2
CHKFZC='""" + (fzctemp) + """'
LPRINT=2"""
        ajf_body += """
/

&MP2DNS
/

&MP2GRD
/

&MP3"""
        if self.ajf_method == 'MP3':
            ajf_body += """
  IFSCS='YES'
  IFMANU='YES'
  SCL2OS=1.0D0
  SCL2PS=1.0D0
  SCL3=0.5D0
"""

        ajf_body += """
/
&LMP2
/
"""
        # define new section
        if self.abinit_ver in ['rev15', 'rev17', 'rev22', 'rev23', 'v2rev4', 'v2rev8']:
            new_section = """
&LRD
"""
            if self.disp:
                new_section += """  DISP='ON'
"""
            new_section += """/
"""
        else:
            new_section = ""

        if self.abinit_ver in ['rev10', 'rev11', 'rev15', 'rev17', 'rev22', 'rev23', 'v2rev4', 'v2rev8' ]:
            new_section += """
&DFT
/

&ANALYSIS """
            if self.piedaflag:
                new_section += """
PIEDA='YES'"""
            else:
                new_section += """
PIEDA='NO'"""

            new_section += """
/
"""

        elif self.abinit_ver == 'mizuho':
            new_section = """
&DFT
/

&PIEDA"""
            if self.piedaflag:
                new_section += """
EnergyDecomposition='YES'"""
            new_section += """
/
"""

        else:
            new_section = ""

        ajf_body += str(new_section) + """
&BSSE"""
        if self.bsseflag:
            ajf_body += """
CP='ON'"""
        ajf_body += """
/

&FRAGPAIR
/

"""
        # define solv section
        if self.solv_flag:
            solv_section = """
&SOLVATION
EFFECT='ON'
ITRMOD='Normal'
MAXITR=100
THICNV=""" + str(self.pbcnv) + """
IGSCC='Core'
INIEHF='OFF'
PRBRAD=1.4
EPSOUT=80.0
EPSIN=1.0
SRFTNS=0.0072
SRFOFF=0.0
NSPHER=1000"""
            if self.abinit_ver in ['rev10', 'rev11', 'rev15', 'mizuho']:
                solv_section += "\nSCREEN='ES+NP' \n"
        else:
            solv_section = """
&SOLVATION
EFFECT='OFF' """

        ajf_body += str(solv_section) + """
/

&PBEQ
"""
        if self.solv_flag:
            ajf_body += """
MAXITR=1000
JDGCNV='RMS'
THRCNV=1.0E-5
"""
            if self.abinit_ver in ['rev17', 'rev22', 'rev23', 'v2rev4', 'v2rev8']:
                ajf_body += "ATMRAD='" + self.pbmolrad + "'"
            else:
                ajf_body += "MOLRAD='" + self.pbmolrad + "'"
        ajf_body +="""
/

&POP
"""
        if self.resp :
            ajf_body += "ESPFIT='ON'\n"
            ajf_body += "ESPTYP='RESP'\n"
        if self.nbo :
            ajf_body +="NBOANL='ON'"
        ajf_body +="""
/
"""

        if self.abinit_ver in ['rev22', 'rev23', 'v2rev4', 'v2rev8']:
            ajf_body += """
&COUPLING
/
"""

        ajf_body += """
&GRIDCNTRL
GRID='NO'
/

&MCP
/

&CIS
/

&CISGRD
/

&CAFI
METLOC='PIPE'
IFLOC='OCC'
CHKFZC='NO'
LPRINT=2
/

&POL
/

&GF2
/
"""
        # new section2
        if self.abinit_ver in ['rev11', 'rev15', 'rev17', 'rev22', 'rev23', 'v2rev4', 'v2rev8']:
            new_section2 = """
&CCPT
/
"""
        else:
            new_section2 = ""

        ajf_body += str(new_section2) + """
&XYZ
""" + self.xyzstr + """/

&FRAGMENT
""" + str(ajf_fragment) + """
/

&MDCNTRL
MD='OFF'
/

&VEL
/

&NHC
/

&TYPEFRAG
/"""
        return ajf_body


    def getfraginfo(self, seg1_conf: dict[str, Any], seg2_conf: dict[str, Any]) -> list[list[int]]:
        """2つのセグメント設定からフラグメント番号リストを取得する。

        Args:
            seg1_conf: セグメント1の設定辞書（seg_infoキーを含む）。
            seg2_conf: セグメント2の設定辞書（seg_infoキーを含む）。

        Returns:
            list: 各セグメントのフラグメント番号を格納した2次元リスト。
        """
        frag1_num = len(seg1_conf['seg_info'])
        frag2_num = len(seg2_conf['seg_info'])

        fcount = 1
        frag = [[], []]
        for i in range(frag1_num):
            frag[0].append(fcount)
            fcount += 1
    #         print frag[0][i]

        for i in range(frag2_num):
            frag[1].append(fcount)
            fcount += 1
    #         print frag[1][i]

        return frag

    def getifie(self, target_dir: str, frag: list[list[int]], skipbsse: bool = False) -> None:
        """ABINIT-MP出力ファイル群からIFIEを読み取り、集計結果をファイルに書き出す。

        Args:
            target_dir: 出力ファイルが格納されたディレクトリパス。
            frag: フラグメント番号の2次元リスト。
            skipbsse: TrueならBSSE補正をスキップする。
        """
        ielist = []
        if not self.pbflag:
            # print "**** 1.get ifie running*****"
            for i in range(1, self.total_num+1):
                target = target_dir + '/%05d' % i + ".out"
            #    if i % 1000 == 0:
            #        print target + " end"
                energy = self.read_ifie(target, skipbsse)
                ifiesum = self.getifiesum(energy, frag)
                ielist.append(ifiesum)
            # print energies
            out = target_dir + "/ielist_ifie" + self.mp2temp + self.solvtype + "_detail"
            with open(out, 'w') as f:
                # print ielist
                print("HF-IFIE", "MP2-IFIE", "PB-term", "PB-polar", "PB-nonpolar", file=f)
                for i in range(len(ielist)):
                    print(ielist[i][0], ielist[i][1], 0.0, 0.0, 0.0, file=f)

            out = target_dir + "/ielist_ifie" + self.mp2temp + self.solvtype
            with open(out, 'w') as f:
                # print ielist
                for i in range(len(ielist)):
                    print(ielist[i][0] + ielist[i][1], file=f)

        if self.pbflag:
            # print("*** 1.get ifie-pb running ***")
            for i in range(1, self.total_num+1):
                target = target_dir + '/%05d' % i + ".out"
                # if i % 1000 == 0:
                    # print(target)
                energy = self.read_ifiepb(target)
                ifiesum = self.getifiesum(energy[0], frag)
                pbsum = self.getifiesumpb(energy[1], frag)
                ielist.append([ifiesum, pbsum])
            # print ielist
            out = target_dir + "/ielist_ifie" + self.mp2temp + self.solvtype + "_detail"
            with open(out, 'w') as f:
                # print ielist
                print("HF-IFIE", "MP2-IFIE", "PB-term", "PB-polar", "PB-nonpolar", file=f)
                for i in range(len(ielist)):
                    print(ielist[i][0][0], ielist[i][0][1], \
                            ielist[i][1][0] + ielist[i][1][1], \
                            ielist[i][1][0], ielist[i][1][1], file=f)

            out = target_dir + "/ielist_ifie" + self.mp2temp + self.solvtype
            with open(out, 'w') as f:
                # print ielist
                for i in range(len(ielist)):
                    print(ielist[i][0][0] + ielist[i][0][1] + \
                            ielist[i][1][0] + ielist[i][1][1], file=f)

    def getifiesum(self, energy: list[list[Any]], frag: list[list[int]]) -> list[float]:
        """IFIEエネルギーのHFおよびMP2成分を合計し、kcal/mol単位で返す。

        Args:
            energy: IFIEエネルギーリスト。
            frag: フラグメント番号の2次元リスト。

        Returns:
            list: [HF合計値, MP2合計値]（kcal/mol）。
        """
        if len(energy) == 0:
            return [0, 0]
        ifiesum = []
        hf = 0.0
        mp2 = 0.0
        total_frag = len(frag[0]) + len(frag[1])
        pairnum = int((total_frag * (total_frag - 1)) / 2)
        for i in range(pairnum):
            for j in(len(frag[0]) + 1, total_frag + 1):  # seg2の要素を動かすループ
                for k in(1, len(frag[0]) + 1):  # seg1の要素を動かすループ
                    if int(energy[i][1]) == int(k) and int(energy[i][0]) == int(j):
                        hf += float(energy[i][4])
                        if self.PR_flag:
                            mp2 += float(energy[i][6])
                        else:
                            mp2 += float(energy[i][5])
        return [hf * 627.5095, mp2 * self.mp2fac * 627.5095]

    def getifiesumpb(self, energy: list[list[Any]], frag: list[list[int]]) -> list[float]:
        """PB溶媒和エネルギーのES項およびNP項を合計し、kcal/mol単位で返す。

        Args:
            energy: PBエネルギーリスト。
            frag: フラグメント番号の2次元リスト。

        Returns:
            list: [ES合計値, NP合計値]（kcal/mol）。
        """
        # print energy
        if len(energy) == 0:
            return [0, 0]
        ifiesum = []
        pb_es = 0.0
        pb_np = 0.0
        total_frag = len(frag[0]) + len(frag[1])
        pairnum = int((total_frag * (total_frag - 1)) / 2)
        for i in range(pairnum):
            for j in(len(frag[0]) + 1, total_frag + 1):  # seg2の要素を動かすループ
                for k in(1, len(frag[0]) + 1):  # seg1の要素を動かすループ
                    if int(energy[i][2]) == int(k) and int(energy[i][1]) == int(j):
                        pb_es += float(energy[i][4])
                        pb_np += float(energy[i][5])
        return [pb_es * 627.5095, pb_np * 627.5095]


    def read_ifie(self, fname: str, skipbsse: bool = False, debug: bool = False) -> list[list[Any]]:
        '''
        read ifie from abinitmp output file
        Args:
            fname: file name
            skipbsse: skip bsse term
        Returns:
            ifie: ifie list
        '''

        ifie = []
        count = 0
        bsseflag = False
        bssecount = 0
        bsse = []
        readflag = False
        try:
            with open(fname, "r") as f:
                text = f.readlines()
        except IOError:
            logger.error("can't open %s", fname)
            return ifie
        flag = False
        # print text
        for i in range(len(text)):
            itemList = text[i][:-1].split()
            # print itemList
            if len(itemList) < 2:
                continue
            if itemList[1] == 'MP2-IFIE':
                flag = True
                readflag = True
                # head.append(itemList)
                # print('READ IFIE start!')
                continue
            if itemList[1] in ['PIEDA']:
                flag = False
                continue
            # after pieda or BSSE (break)
            if itemList[1] == 'Mulliken':
                break
            # for BSSE
            if itemList[:5] == ['##', 'BSSE', 'for', 'non-bonding', 'MP2-IFIE']:
                if skipbsse:
                    break
                flag = False
                # print('pieda end! next is BSSE')
                continue
            if itemList[:4] == ['##', 'BSSE', 'for', 'MP2-IFIE']:
                if skipbsse:
                    break
                bsseflag = True
                # print('BSSE start!')
                continue
            if bsseflag:
                bssecount += 1
            if bsseflag and bssecount > 2:
                bsse.append(itemList)
            if flag:
                count += 1
            if flag and count > 2:
                ifie.append(itemList)
                if debug:
                    logger.debug(itemList)
        if not readflag:
            try:
                logger.warning("can't read ifie %s", fname.split("/")[1])
            except (ValueError, IndexError):
                pass

        if debug:
            logger.debug("ifie %s", ifie)

        for i in range(len(ifie)):
            if float(ifie[i][4]) < -2 or float(ifie[i][5]) < -2:
                ifie[i][4] = 0.0
                ifie[i][5] = 0.0
                ifie[i][6] = 0.0

        if bsseflag:
            # print(bsse)
            for i in range(len(ifie)):
                if not (float(ifie[i][4]) < -2 or float(ifie[i][5]) < -2):
                    ifie[i][4] = float(ifie[i][4]) + float(bsse[i][4])
                    ifie[i][5] = float(ifie[i][5]) + float(bsse[i][5])
                    ifie[i][6] = float(ifie[i][6]) + float(bsse[i][6])

        # print ifie
        return ifie

    def read_ifiepb(self, fname: str) -> list[list[list[Any]]]:
        """ABINIT-MP出力ファイルからIFIEとPB溶媒和項を読み取る。

        Args:
            fname: 出力ファイルパス。

        Returns:
            list: [IFIEリスト, PB項リスト]。
        """
        ifie = []
        pbterm = []
        count = 0
        count2 = 0
        hit = 0
        try:
            with open(fname, "r") as f:
                text = f.readlines()
        except (IOError, OSError):
            logger.error("can't open %s", fname)
            return [ifie, pbterm]
        flag = False
        flag2 = False
        # print text
        for i in range(len(text)):
            itemList = text[i][:-1].split()
            # print itemList
            if len(itemList) < 2:
                continue
            if itemList[1] == 'MP2-IFIE':
                hit += 1
                if hit == 2:
                    flag = True
                # head.append(itemList)
                continue
            if flag:
                if itemList[1] == 'Mulliken' or itemList[1] == 'PIEDA' or itemList[1] == 'NATURAL':
                    flag = False
                # break
            if flag:
                count += 1
            if flag and count > 2:
                ifie.append(itemList)

            if len(itemList) >= 4:
                if itemList[1] == 'SUMMARY' and itemList[3] == 'ELECTROSTATIC':
                    flag2 = True
                    continue
            if len(itemList) >= 4:
                if itemList[1] == 'ESTIMATION' and itemList[3] == 'NON-POLAR':
                    flag2 = False
                    break
            if flag2:
                count2 += 1
            if flag2 and count2 >= 4:
                pbterm.append(itemList)
        if hit != 2 and not flag2:
            logger.warning("can't read pb %s", fname.split("/")[1])

        for i in range(len(ifie)):
            if float(ifie[i][4]) < -2 or float(ifie[i][5]) < -2:
                ifie[i][4] = 0.0
                ifie[i][5] = 0.0
                ifie[i][6] = 0.0
                pbterm[i][4] = 0.0
                pbterm[i][5] = 0.0

        return [ifie, pbterm]

    @staticmethod
    def unpack_tar(target_dir: str, total_num: int) -> None:
        """ディレクトリ内のout_files.tarを展開する。

        出力ファイル数が期待値の90%未満の場合のみ展開を実行する。

        Args:
            target_dir: tarファイルが格納されたディレクトリパス。
            total_num: 期待される出力ファイルの総数。
        """
        # (前) target_dir 内の out_files.tar を展開
        out_tar = os.path.join(target_dir, 'out_files.tar')
        # outfile = os.path.join(target_dir, '*sh*.out')

        outfile = [
            fn for fn in glob.glob(os.path.join(target_dir, '*.out'))
            if 'zz_submit' not in os.path.basename(fn)
        ]

        logger.info("Checking for %s in %s", os.path.basename(out_tar), target_dir)
        logger.info("Total number of expected output files: %s", total_num)
        num_outfile = len(outfile)
        logger.info("num of outfile %s", num_outfile)
        if os.path.exists(out_tar) and num_outfile/total_num < 0.9:
            with tarfile.open(out_tar, 'r') as tar:
                tar.extractall(path=target_dir)
            logger.info("Extracted %s into %s", os.path.basename(out_tar), target_dir)
        return

    @staticmethod
    def del_out(target_dir: str) -> None:
        """ディレクトリ内の出力ファイル（*.out）を一括削除する。

        Args:
            target_dir: 出力ファイルが格納されたディレクトリパス。
        """
        out_files = [
            fn for fn in glob.glob(os.path.join(target_dir, '*.out'))
            if 'zz_submit' not in os.path.basename(fn)
        ]
        removed = 0
        for out_file in out_files:
            try:
                os.remove(out_file)
                removed += 1
            except OSError:
                pass
        if removed:
            logger.info("Removed %d outfiles", removed)
        else:
            logger.info("No out files to remove")
        return

    def getTE(self, target_dir: str, molname: str, mode: str, fzcflag: bool) -> None:
        """全エネルギー（TE）を取得してファイルに書き出す。

        Args:
            target_dir: 出力ファイルが格納されたディレクトリパス。
            molname: 分子名。
            mode: 'batch'（複数ファイル処理）または'single'（単一ファイル処理）。
            fzcflag: Trueの場合nofzcディレクトリから読み込む。
        """
        elistname = target_dir + "/energylist_" + self.solvtype
        if mode == "batch":
            energies = []
            if not self.pbflag:
                logger.info("**** 1.capt te running ****")

                self.unpack_tar(target_dir, self.total_num)

                # 処理
                for i in range(1, self.total_num+1):
                    target = target_dir + '/%05d' % i + ".out"
                    if i % 500 == 0:
                        logger.info(target)
                    energy = self.captfmomp2e(target)
                    energies.append(energy)
                # print energies
                with open(elistname, "w") as f:
                    for i in energies:
                        try:
                            print(i[0], i[1], file=f)
                        except TypeError:
                            print(0, 0, file=f)
                # print("create", elistname)
                # self.del_out(target_dir)

            if self.pbflag:
                logger.info("*** 1.capt bepb running ***")

                self.unpack_tar(target_dir, self.total_num)

                # 本処理
                for i in range(1, self.total_num+1):
                    target = target_dir + '/%05d' % i + ".out"
                    if i % 500 == 0:
                        logger.info(target)
                    energy = self.getfmopbenergy(target)
                    energies.append(energy)
                # print energies
                with open(elistname, "w") as f:
                    for i in energies:
                        try:
                            print(i[0], i[1], i[2], i[3], i[4], file=f)
                        except TypeError:
                            print(0, 0, 0, 0, 0, file=f)
                # print("*****create", elistname)
                # self.del_out(target_dir)

        if mode == "single":
            target = target_dir + "/" + molname + ".out"
            if fzcflag:
                # print (molname + ": nofzc")
                target = target_dir + "/nofzc/" + molname + ".out"
            if not self.pbflag:
                self.captmom_single(target)
            elif self.pbflag:
                self.captpb_single(target)

        return

    def captmom_single(self, target: str) -> None:
        """単一ファイルからMO/FMOのMP2エネルギーを取得し、結果ファイルに書き出す。

        Args:
            target: ABINIT-MP出力ファイルのパス。
        """
        out = os.path.splitext(target)[0]
        fmoflag = self.getmo_or_fmo(target)  # FMO計算かMO計算かの判定
        # print(target)
        # print(fmoflag)
        if fmoflag:
            hf, mp2 = self.captfmomp2e(target)
        else:
            hf, mp2 = self.getmomp2ene(target)
        with open(out + "_te" + self.solvtype + ".dat", "w") as f:
            # print("hf:", hf, "mp2:", mp2)
            print(hf, mp2, file=f)

    def captpb_single(self, target: str) -> None:
        """単一ファイルからPB溶媒和エネルギーを取得し、結果ファイルに書き出す。

        Args:
            target: ABINIT-MP出力ファイルのパス。
        """
        out = os.path.splitext(target)[0]
        fmoflag = self.getmo_or_fmo(target)  # FMO計算かMO計算かの判定
        if fmoflag:
            eg, cor, dg, dges, dgnp = self.getfmopbenergy(target)
        else:
            eg, cor, dg, dges, dgnp = self.getmopbenergy(target)
        with open(out + "_te" + self.solvtype + ".dat", "w") as f:
            # print("eg:" + eg, "cor:" + cor, "dg:" + dg, "dges:" + dges, "dgnp:" + dgnp)
            print(eg, cor, dg, dges, dgnp, file=f)

    def getmo_or_fmo(self, target: str) -> bool:
        """出力ファイルがFMO計算かMO計算かを判定する。

        Args:
            target: ABINIT-MP出力ファイルのパス。

        Returns:
            bool: FMO計算の場合True、MO計算の場合False。
        """
        with open(target, "r") as f:
            text = f.readlines()
        fmoflag = False
        for i in range(len(text)):
            itemList = text[i][:-1].split()
            # print itemList
            if "FMO" in itemList and "ON" in itemList:
                fmoflag = True  # FMO mom

                # print(target + " is FMO")
                break
        return fmoflag

    def captfmomp2e(self, target: str) -> list[float]:
        """FMO計算出力ファイルからHFエネルギーとMP2エネルギーを抽出する。

        Args:
            target: ABINIT-MP出力ファイルのパス。

        Returns:
            list: [HFエネルギー, MP2エネルギー]。取得失敗時は[0, 0]。
        """
        mp2 = 0
        hf = 0
        try:
            with open(target, "r") as f:
                text = f.readlines()
            for i in range(len(text)):
                itemList = text[i][:-1].split()
                if itemList == ['##', 'FMO', 'TOTAL', 'ENERGY']:
                    hf = text[i + 6].split()
                    if self.PR_flag:
                        mp2 = text[i + 13].split()
                    else:
                        mp2 = text[i + 8].split()
                    return [float(hf[3]), float(mp2[2])]
        except (OSError, ValueError, IndexError):
            logger.warning("can't get result: %s", target)
            return [0, 0]

    def captfmomp2etar(self, target: str) -> list[float]:
        """
        target: '/path/to/00001.out' のようなパスを受け取る。
        常に同ディレクトリの out_files.tar から member を取り出して解析する。
        """
        out_tar = os.path.join(os.path.dirname(target), 'out_files.tar')
        try:
            with tarfile.open(out_tar, 'r') as tar:
                member_name = os.path.basename(target)
                member = tar.getmember(member_name)
                fobj = tar.extractfile(member)
                if fobj is None:
                    msg = (
                        f"Failed to extract {member_name} "
                        f"from {out_tar}"
                    )
                    raise IOError(msg)
                raw = fobj.read()
                text = raw.decode('utf-8')
                lines = text.splitlines()
        except Exception as e:
            logger.warning("can't get result: %s %s", target, e)
            return [0.0, 0.0]

        # エネルギー抽出
        for idx, line in enumerate(lines):
            cols = line.split()
            if cols == ['##', 'FMO', 'TOTAL', 'ENERGY']:
                hf_vals = lines[idx + 6].split()
                mp2_vals = (
                    lines[idx + 13].split()
                    if self.PR_flag
                    else lines[idx + 8].split()
                )
                hf = float(hf_vals[3])
                mp2 = float(mp2_vals[2])
                return [hf, mp2]

        logger.warning("can't parse energies in: %s", target)
        return [0.0, 0.0]

    def getmomp2ene(self, target: str) -> list[float]:
        """MO計算出力ファイルからHFエネルギーとMP2エネルギーを抽出する。

        Args:
            target: ABINIT-MP出力ファイルのパス。

        Returns:
            list: [HFエネルギー, MP2エネルギー]。取得失敗時は[0, 0]。
        """
        mp2 = 0
        hf = 0
        try:
            with open(target, "r") as f:
                text = f.readlines()
        except (IOError, OSError):
            logger.error("can't open monomer file: %s", target)
            return [0,0]
            # sys.exit()

        hfflag = False
        for i in range(len(text)):
            itemList = text[i][:-1].split()
            # print itemList
            if itemList == ['SCF', 'COMPLETED']:
                # print itemList
                if not hfflag:
                    hf = text[i + 4].split()
                    hfflag = True
            if itemList == ['##', 'MP2', 'ENERGY']:
                # print itemList
                if self.PR_flag:
                    mp2 = text[i + 7].split()
                else:
                    mp2 = text[i + 2].split()
                break

        try:
            return [float(hf[3]), float(mp2[2])]
        except (IndexError, TypeError):
            logger.error("can't get monomer result: %s", target)
            return [0,0]
            # sys.exit()


    def getfmopbenergytar(self, target: str) -> tuple[Any, Any, Any, Any, Any]:
        """tarアーカイブ内のFMO出力ファイルからPB溶媒和エネルギーを抽出する。

        Args:
            target: tarアーカイブ内の出力ファイルパス。

        Returns:
            tuple: (気相エネルギー, 溶液中エネルギー, 溶媒和自由エネルギー, ES項, NP項)。
                   取得失敗時は (0, 0, 0, 0, 0)。
        """

        try:
            # out_files.tar のパス
            tar_path = os.path.join(os.path.dirname(target), 'out_files.tar')

            # tar を開いて member を取り出す
            with tarfile.open(tar_path, 'r') as tar:
                member_name = os.path.basename(target)  # '00001.out' 等
                member = tar.getmember(member_name)
                fobj = tar.extractfile(member)
                if fobj is None:
                    msg = (
                        f"Failed to extract {member_name} "
                        f"from {tar_path}"
                    )
                    raise IOError(msg)

                # バイト列を読み込み→文字列化→行リスト化
                raw = fobj.read()
                text = raw.decode('utf-8')
                lines = text.splitlines()
                fobj.close()

        except (OSError, KeyError, UnicodeDecodeError):
            logger.error("can't open %s", target)
            return 0, 0, 0, 0, 0

        if self.abinit_ver in ['rev17', 'rev22', 'rev23', 'v2rev4', 'v2rev8']:
            pbdone =False
            getflag = False
            for line in lines:
                itemList = line.split()
                if len(itemList) < 3:
                    continue
                if itemList[0:4] == ['##', 'SOLUTE', 'TOTAL', 'ENERGY']:
                    getflag = True
                if getflag:
                    if itemList[0:3] == ['FMO2', 'in', 'vacuo']:
                        # print('vacuo')
                        eg = itemList  # [5]
                    if itemList[0:2] == ['in', 'solvent']:
                        # print('in solv')
                        esol = itemList
                    if itemList[0] == 'difference':
                        # print('diff')
                        dif = itemList
                    if itemList[0] == 'Electrostatic' or \
                            itemList[0:3] == ["ES", "term", "(FMO2-MP2)"]:
                        # print('ES')
                        dges = itemList
                    if itemList[0] == 'Nonpolar' or \
                            itemList[0:2] == ["NP", "term"]:
                        # print('NP')
                        dgnp = itemList
                    if itemList[0] == 'Total':
                        # print('total')
                        dg = itemList
                        pbdone =True
                if pbdone:
                    break
            # egas, cor(insolv), dgtotal, dges, dgnp
            # print(eg[3], esol[3], dg[2], dges[2], dgnp[2])
            try:
                return eg[3], esol[3], dg[-1], dges[-1], dgnp[-1]
            except (NameError, IndexError):
                logger.warning("can't get result: %s", target)
                return 0, 0, 0, 0, 0

            '''
            ## SOLUTE TOTAL ENERGY

                                       SCF / Hartree   MP2 corr. / Hartree       Total / Hartree
               FMO2 in vacuo           -236.52210938*          -0.80641280         -237.32852218
                    in solvent         -236.52238631           -0.80642693*        -237.32881324
                    difference           -0.00027693           -0.00001413           -0.00029106
                  ( ES correction        -0.00028778 )


             ## SOLVATION FREE ENERGY

                                       FMO2 / Hartree       FMO2 / kcal/mol
                Electrostatic             -0.04769392          -29.92838623*
                Nonpolar                   0.00548501            3.44189687*
                Total                     -0.04220891          -26.48648936*
            '''



    def getfmopbenergy(self, target: str) -> tuple[Any, Any, Any, Any, Any]:
        """FMO出力ファイルからPB溶媒和エネルギーを抽出する。

        Args:
            target: ABINIT-MP出力ファイルのパス。

        Returns:
            tuple: (気相エネルギー, 溶液中エネルギー, 溶媒和自由エネルギー, ES項, NP項)。
                   取得失敗時は (0, 0, 0, 0, 0)。
        """
        try:
            with open(target, "r") as f:
                text = f.readlines()
        except (IOError, OSError):
            logger.error("can't open %s", target)
            return 0, 0, 0, 0, 0
        if self.abinit_ver in ['rev11', 'rev15', 'mizuho']:
            index = 0
            pbdone = False
            for i in range(len(text)):
                itemList = text[i][:-1].split()
                if len(itemList) == 0:
                    continue
                if itemList[0:4] == ['Energy', 'in', 'gas', 'phase']:
                    eg = itemList  # [5]
                if itemList[0:4] == ['ElectroStatic', '[', 'DGes=Es-Eg-DE', ']']:
                    dges = itemList
                if itemList[0:4] == ['Non-polar', '[', 'DGnp', ']']:
                    dgnp = itemList
                if itemList[0:4] == ['Total', '[', 'DG=DGes+DGnp', ']']:
                    dg = itemList  # [5]
                    pbdone =True
                if itemList == ['##', 'FMO', 'TOTAL', 'ENERGY']:
                    index = i
            if not pbdone:
                return 0, 0, 0, 0, 0
            cor = text[index+8].split()
            # print eg, cor, dg
            try:
                return eg[8], cor[2], dg[5], dges[5], dgnp[5]
            except (NameError, IndexError):
                logger.warning("can't get result: %s", target)
                return 0, 0, 0, 0, 0

        if self.abinit_ver in ['rev17', 'rev22', 'rev23', 'v2rev4', 'v2rev8']:
            index = 0
            pbdone =False
            getflag = False
            for i in range(len(text)):
                itemList = text[i][:-1].split()
                if len(itemList) < 3:
                    continue
                if itemList[0:4] == ['##', 'SOLUTE', 'TOTAL', 'ENERGY']:
                    getflag = True
                if  getflag:
                    if itemList[0:3] == ['FMO2', 'in', 'vacuo']:
                        # print('vacuo')
                        eg = itemList  # [5]
                    if itemList[0:2] == ['in', 'solvent']:
                        # print('in solv')
                        esol = itemList
                    if itemList[0] == 'difference':
                        # print('diff')
                        dif = itemList
                    if itemList[0] == 'Electrostatic' or \
                            itemList[0:3] == ["ES", "term", "(FMO2-MP2)"]:
                        # print('ES')
                        dges = itemList
                    if itemList[0] == 'Nonpolar' or \
                            itemList[0:2] == ["NP", "term"]:
                        # print('NP')
                        dgnp = itemList
                    if itemList[0] == 'Total':
                        # print('total')
                        dg = itemList
                        pbdone =True
                if pbdone:
                    break
            # egas, cor(insolv), dgtotal, dges, dgnp
            # print(eg[3], esol[3], dg[2], dges[2], dgnp[2])
            try:
                return eg[3], esol[3], dg[-1], dges[-1], dgnp[-1]
            except (NameError, IndexError):
                logger.warning("can't get result: %s", target)
                return 0, 0, 0, 0, 0

            '''
            ## SOLUTE TOTAL ENERGY

                                       SCF / Hartree   MP2 corr. / Hartree       Total / Hartree
               FMO2 in vacuo           -236.52210938*          -0.80641280         -237.32852218
                    in solvent         -236.52238631           -0.80642693*        -237.32881324
                    difference           -0.00027693           -0.00001413           -0.00029106
                  ( ES correction        -0.00028778 )


             ## SOLVATION FREE ENERGY

                                       FMO2 / Hartree       FMO2 / kcal/mol
                Electrostatic             -0.04769392          -29.92838623*
                Nonpolar                   0.00548501            3.44189687*
                Total                     -0.04220891          -26.48648936*
            '''


    def getmopbenergy(self, target: str) -> tuple[Any, Any, Any, Any, Any]:
        """モノマーPB（Poisson-Boltzmann）エネルギーをログファイルから取得する。

        Args:
            target: ログファイルのパス。

        Returns:
            tuple: (気相エネルギー, 相関エネルギー, 全溶媒和自由エネルギー,
                静電項, 非極性項) の5要素タプル。取得失敗時はすべて0。
        """
        count = 0
        try:
            with open(target, "r") as f:
                text = f.readlines()
        except (IOError, OSError):
            logger.error("can't open %s", target)
            return 0, 0, 0, 0, 0

        if self.abinit_ver in ['rev11', 'rev15', 'mizuho']:

            for i in range(len(text)):
                itemList = text[i][:-1].split()
        #         print itemList
                if len(itemList) == 0:
                    continue
                if itemList[0:4] == ['Energy', 'in', 'gas', 'phase']:
                    eg = itemList  # [5]
                if itemList[0:4] == ['ElectroStatic', '[', 'DGes=Es-Eg-DE', ']']:
                    dges = itemList
                if itemList[0:4] == ['Non-polar', '[', 'DGnp', ']']:
                    dgnp = itemList
                if itemList[0:4] == ['Total', '[', 'DG=DGes+DGnp', ']']:
                    dg = itemList  # [5]
                if itemList == ['##', 'MP2', 'ENERGY']:
                    count += 1
                    if count == 2:
                        cor = text[i+2].split()
            # print eg, cor, dg
            try:
                return eg[8], cor[2], dg[5], dges[5], dgnp[5]
            except (NameError, IndexError):
                logger.error("can't get monomer result: %s", target)
                # sys.exit()
                return 0, 0, 0, 0, 0

        if self.abinit_ver in ['rev17', 'rev22', 'rev23', 'v2rev4', 'v2rev8']:
            a, b, c, d, e = self.getfmopbenergy(target)
            return a, b, c, d, e



    def readajf(self, fname: str) -> abinit_io:
        """AJFファイルを読み込み、フラグメント情報をselfに格納する。

        Args:
            fname: AJFファイルのパス。

        Returns:
            self: フラグメント原子数・電荷・BDA・原子情報・結合情報を格納したインスタンス。
        """
        with open(fname, 'r') as _fh:
            lines = _fh.readlines()

        flag = False
        baseline = 1

        datas = []
        fatomnums = []
        fchgs = []
        fbaas = []
        nfcount = 0
        fatminfo = []
        fatminfos = []
        connects = []
        for i in range(len(lines)):
            itemlist = lines[i].split()
            if len(itemlist) == 0:
                continue
            if itemlist[0][:2].upper() == 'NF':
                self.nf = int(itemlist[0].split('=')[-1])
                nf = self.nf
                logger.info("NF = %s", self.nf)
                if self.nf > 10:
                    baseline = math.ceil(self.nf/10)
                    logger.debug("n_line %s", baseline)
            if itemlist[0] == '&FRAGMENT':
                flag = True
                typcount = 0
                lcount = 0
                continue
            # fatom section
            if flag and typcount == 0:
                fatomnums.append(itemlist)
                lcount += 1
                if lcount == baseline:
                    typcount += 1
                    lcount = 0
                    fatomnums = self.flatten(fatomnums)
                    continue
            # chg section
            if flag and typcount == 1:
                fchgs.append(itemlist)
                lcount += 1
                if lcount == baseline:
                    typcount += 1
                    lcount = 0
                    fchgs = self.flatten(fchgs)
                    continue

            # bda section
            if flag and typcount == 2:
                fbaas.append(itemlist)
                lcount += 1
                if lcount == baseline:
                    typcount += 1
                    lcount = 0
                    fbaas = self.flatten(fbaas)
                    continue

            # seginfo section
            if flag and typcount == 3:
                # print('typ', typcount)
                # print(nfcount)
                fatomnum = fatomnums[nfcount]
                base2 = math.ceil(int(fatomnum)/10)
                fatminfo.append(itemlist)
                lcount += 1
                if lcount == base2:
                    nfcount += 1
                    lcount = 0
                    fatminfo = self.flatten(fatminfo)
                    fatminfos.append(fatminfo)
                    fatminfo = []
                    if  nfcount > nf - 1:
                        typcount += 1
                        lcount = 0
                        continue

            # bda-baa atom section
            if flag and itemlist[0] == '/':
                break

            if flag and typcount == 4:
                connects.append(itemlist)
                connects = self.functor(int, connects)
                # datas.append(itemlist)

        self.fatomnums = fatomnums
        self.fchgs = fchgs
        self.fbaas = fbaas
        self.fatminfos = fatminfos
        self.connects = connects
        return self
        # return fatomnums, fchgs, fbaas, fatminfos, connects

# def readfragcpf(file, nf):
#     '''Read the fragment data from the CPF file.
#
#     Args:
#         file (file): The CPF file to read from
#         nf (int): The number of fragments in the CPF file
#
#     Returns:
#         fnatoms (list): The number of atoms in each fragment
#         fchgs (list): The charge of each fragment
#         fbaas (list): The BAA of each fragment
#         fatminfos (list): The atom IDs of each fragment
#         connects (list): The connect (bda baa) atom IDs of each fragment
#
#     Note:
#         The CPF file format is described in the CPF manual.
#
#         fnatoms = [natom1, natom2, ...]
#         fchgs = [chg1, chg2, ...]
#         fbaas = [baa1, baa2, ...]
#         fatminfos = [[atomid1, atomid2, ...], [atomid1, atomid2, ...], ...]
#         connects = [[atomid1, atomid2, ...], [atomid1, atomid2, ...], ...]
#     '''
#
#     # Initialize the data structures
#     flag = True
#     nline = 1  # line number
#     fnatoms = []  # fragment natom section
#     fchgs = []  # fragment charge section
#     fbaas = []  # bda section
#     nfcount = 0  # fragment count
#     fatminfo = []  # fraginfo atomid section
#     fatminfos = []  # fraginfo atomid section
#     connects = []  # bda-baa atom section
#     typcount = 0
#     # type (0: fnatoms, 1: fchgs, # 2: fbaas, 3: fatminfo, 4: connects)
#     lcount = 0  # line count
#
#     while True:
#         itemlist = file.readline().strip().split()
#         if len(itemlist) == 0:
#             continue
#         if nf > 10:
#             nline = math.ceil(nf/10)
#             print('n_line', nline)
#         # fragment natom section
#         if flag is True and typcount == 0:
#             fnatoms.append(itemlist)
#             lcount += 1
#             if lcount == nline:
#                 typcount += 1
#                 lcount = 0
#                 fnatoms = flatten(fnatoms)
#                 continue
#
#         # fragment charge section
#         if flag is True and typcount == 1:
#             fchgs.append(itemlist)
#             lcount += 1
#             if lcount == nline:
#                 typcount += 1
#                 lcount = 0
#                 fchgs = flatten(fchgs)
#                 continue
#
#         # bda section
#         if flag is True and typcount == 2:
#             fbaas.append(itemlist)
#             lcount += 1
#             if lcount == nline:
#                 typcount += 1
#                 lcount = 0
#                 fbaas = flatten(fbaas)
#                 continue
#
#         # fraginfo atomid section
#         if flag is True and typcount == 3:
#             fatomnum = fnatoms[nfcount]
#             base2 = math.ceil(int(fatomnum)/10)
#             fatminfo.append(itemlist)
#             lcount += 1
#             if lcount == base2:
#                 nfcount += 1
#                 lcount = 0
#                 fatminfo = flatten(fatminfo)
#                 fatminfos.append(fatminfo)
#                 fatminfo = []
#                 if nfcount > nf - 1:
#                     typcount += 1
#                     lcount = 0
#                     continue
#
#         # bda-baa atom section
#         count = 0
#         if flag is True and typcount == 4:
#             count += 1
#             connects.append(itemlist)
#             connects = functor(int, connects)
#             if count == sum(fbaas):
#                 break
#
#             # datas.append(itemlist)
#
#     return fnatoms, fchgs, fbaas, fatminfos, connects


    def modifyfragparam(self, totalMol: int, atomnameMol: list[list[str]], molnames: list[str], anummols: list[list[int]], posMol: list[Any], heads: list[Any], labs: list[Any], chains: list[Any], resnums: list[Any], codes: list[Any], occs: list[Any], temps: list[Any], amarks: list[Any], charges: list[Any], fatomnums: list[int], fchgs: list[int], fbaas: list[int], fatminfos: list[list[int]], connects: list[list[int]], bridgeds: list[int], doubles: list[list[int]], nagatoms: list[list[int]], nagmolids: list[int], nagbdas: list[int]) -> tuple[list[int], list[int], list[int], list[list[int]], list[list[int]]]:
        """NAG糖鎖結合に伴うフラグメントパラメータの修正を行う。

        ASN残基に結合するNAG原子をフラグメントから分離し、
        新規フラグメントとして追加する。

        Returns:
            tuple: (fatomnums, fchgs, fbaas, connects, fatminfos) の修正済みフラグメント情報。
        """
        # modify start
        for i in range(len(nagatoms)):
            for j in range(len(fatminfos)):
                # i: nag mol index
                # j: ajf fragment index

                # print(fatminfos[j])
                # print(nagatoms[i][0])
                dendflag = False
                dmidflag = False
                if nagatoms[i][0] in fatminfos[j]: # search (asn + nag)
                    logger.debug("NAG %s frag %s start atom %s", i, j, nagatoms[i][0])

                    # check baa asn atom id
                    tgtid = j
                    for bridged in bridgeds:
                        if tgtid >= bridged:
                            tgtid += 1
                    logger.debug("%s %s", molnames[tgtid], resnums[tgtid][0])
                    logger.debug("atomnameMol %s %s", tgtid, atomnameMol[tgtid])
                    baaidx = atomnameMol[tgtid].index(' ND2')
                    logger.debug("asn_baaidx %s", anummols[tgtid][baaidx])


                    tgtbaaid = sum(fbaas[:j+1])
                    logger.debug("tgtbaaid %s", tgtbaaid)

                    #fatomnums
                    # j: asn frag id
                    fatomnums[j] -= len(nagatoms[i])
                    fatomnums.append(len(nagatoms[i]))

                    logger.debug("doubles %s", doubles)
                    for double in doubles:
                        if double[1] == nagmolids[i]:
                            dendflag = True
                        if double[0] == nagmolids[i]:
                            dmidflag = True

                    if dendflag:
                        baaidx = atomnameMol[nagmolids[i-1]].index(' O4 ')
                        logger.debug("baaidx double %s %s", baaidx, nagmolids[i-1])
                        baaatomid = anummols[nagmolids[i-1]][baaidx]
                        logger.debug("baaaaaaaaaaaaaa %s", baaatomid)

                    # for old(asn) frag
                    if not dendflag:
                        fchgs[j] -= 1
                        fbaas[j] += 1

                    # for new frag
                    if dmidflag:
                        fchgs.append(0)
                        fbaas.append(1)
                    else:
                        fchgs.append(1)
                        fbaas.append(0)

                    #fatminfos
                    delid = []

                    # get delete atom index in mol j
                    for k in range(len(nagatoms[i])):
                        # k: nagatom(in mol) index
                        for l in range(len(fatminfos[j])):
                            # l: atom(in fragment) index
                            if nagatoms[i][k] == fatminfos[j][l]:
                                delid.append(l)
                                break

                    logger.debug("delid %s", delid)
                    dellist = lambda items, indexes: [item for index, item in enumerate(items) if index not in indexes]
                    fatminfos[j] = dellist(fatminfos[j], delid)
                    # atomnameMol[tgtid] = dellist(atomnameMol[j], delid)
                    logger.debug("atom-len-check frag %s : %s -- mol %s : %s", j, len(fatminfos[j]), j, len(atomnameMol[tgtid]))
                    fatminfos.append(nagatoms[i])

                    #connects
                    logger.debug("%s", len(connects))
                    if dendflag:
                        connects.append([int(nagbdas[i]), int(baaatomid)])
                        logger.debug("connects %s %s", nagbdas[i], baaatomid)

                    else:
                        connects.insert(tgtbaaid, [int(nagbdas[i]), int(anummols[tgtid][baaidx])])
                        logger.debug("connects %s %s", nagbdas[i], anummols[tgtid][baaidx])
                    logger.debug(" ")
        return fatomnums, fchgs, fbaas, connects, fatminfos

    def flatten(self, nested_list: list[list[Any]]) -> list[int]:
        """2重のリストをフラットにする関数"""
        return [int(e) for inner_list in nested_list for e in inner_list]

    # def listtoint(nested_list):
    #     """リストの中をintegerに"""
    #     return [int(e) for inner_list in nested_list for e in inner_list]

    def functor(self, f: Callable[[Any], Any], l: Any) -> Any:
        """リストの各要素に関数を再帰的に適用する。"""
        if isinstance(l,list):
            return [self.functor(f,i) for i in l]
        else:
            return f(l)


    def getfragdict(self, fnames: list[str], ofile: str) -> None:
        """複数のAJFファイルからフラグメント辞書データを生成しファイルに書き出す。

        Args:
            fnames: AJFファイルパスのリスト。
            ofile: 出力ファイルパス（segment_data形式）。
        """
        with open(ofile, 'w') as f:

            print("seg_data = [", file=f)
            for fname in fnames:
                head, ext = os.path.splitext(fname)
                # fatomnums, fchgs, fbaas, fatminfos, connects, = self.readajf(fname)
                self = self.readajf(fname)
                print("    {", file=f)
                print("    'name': '" + head.split('/')[-1] + "',", file=f)
                print("    'atom': ", self.fatomnums, ',', file=f)
                print("    'charge': ", self.fchgs, ',', file=f)
                print("    'connect_num': ", self.fbaas, ',', file=f)
                print("    'seg_info': ", self.fatminfos, ',', file=f)
                print("    'connect': ", self.connects, ',', file=f)
                end =  """    'nummol_seg': [1],
    'repeat': [1],
    'pair_file': [],
    'multi_xyz': 'none'
    },"""
                print(end, file=f)

            print(']', file=f)

        '''
                seg_conf = {'name': seg_name,
                            'atom': [seg_atom],
                            'charge': [0],
                            'connect_num': [0],
                            'connect': [],
                            'seg_info': [[i + 1 for i in range(seg_atom)]],
                            'nummol_seg': [1],
                            'repeat': [1],
                            'pair_file': [],
                            'multi_xyz': 'none'}
        '''

    def config_read(self, seg_name: str, seg_atom: int = 0) -> dict[str, Any]:
        """segment_data.datからセグメント設定を読み込む。

        Args:
            seg_name: セグメント名。
            seg_atom: セグメントの原子数（デフォルト: 0）。

        Returns:
            dict: セグメント設定辞書（atom, charge, connect等を含む）。
        """
        fdata = {}
        with open(self.mainpath + "/segment_data.dat", "r") as _f:
            _tree = ast.parse(_f.read())
        for _node in _tree.body:
            if isinstance(_node, ast.Assign) and len(_node.targets) == 1 and isinstance(_node.targets[0], ast.Name):
                fdata[_node.targets[0].id] = ast.literal_eval(_node.value)
        seg_data = fdata['seg_data']
        # print (seg_data)

        # default data(initialize)
        seg_conf = {
                    'name': seg_name,
                    'atom': [seg_atom],
                    'charge': [0],
                    'connect_num': [0],
                    'connect': [],
                    'seg_info': [[i + 1 for i in range(seg_atom)]],
                    'nummol_seg': [1],
                    'repeat': [1],
                    'pair_file': [],
                    'partial_charge': [],
                    'multi_xyz': 'none',
                    'nterm': 2,
                   }

        # search seg_data
        for i in range(len(seg_data)):
            if seg_name == seg_data[i]['name']:
                for key, data in seg_data[i].items():
                    seg_conf[key] = data

            # print "seg_conf =", seg_conf
        return seg_conf

    def getmb_frag_seclists(self, param: list[list[Any]], nameid: list[int]) -> abinit_io:
        """マルチボディ用フラグメントセクションのリストを構築しselfに格納する。

        Args:
            param: [frag_atom, frag_charge, frag_connect_num, frag_connect, seg_info] のリスト。
            nameid: セグメント名IDのリスト。

        Returns:
            self: フラグメント情報を格納したインスタンス。
        """
        frag_atom = param[0]
        frag_charge = param[1]
        frag_connect_num = param[2]
        frag_connect = param[3]
        seg_info = param[4]

        # fragment atom
        count = 0

        frag_atoms = []
        for i in range(len(nameid)):
            for j in range(len(frag_atom[nameid[i]])):
                    frag_atoms.append(frag_atom[nameid[i]][j])
                    count += 1

        # fragment charge
        count = 0
        frag_charges = []
        for i in range(len(nameid)):
            for j in range(len(frag_charge[nameid[i]])):
                frag_charges.append(frag_charge[nameid[i]][j])
                count += 1

        # fragment connect num
        count = 0
        frag_baanums = []
        for i in range(len(nameid)):
            for j in range(len(frag_connect_num[nameid[i]])):
                frag_baanums.append(frag_connect_num[nameid[i]][j])
                count += 1

        # fragment body
        atom_count = 0
        # i mol  j frag k atom
        frag_atmlabss = []
        for i in range(len(nameid)):
            for j in range(len(seg_info[nameid[i]])):
                icount = 0
                frag_atmlabs = []
                for k in range(len(seg_info[nameid[i]][j])):
                    icount += 1
                    frag_atmlabs.append(seg_info[nameid[i]][j][k] + atom_count)
                frag_atmlabss.append(frag_atmlabs)
            atom_count += sum(frag_atom[nameid[i]])

        atom_count = 0
        # connect info
        frag_connects = []
        frag_connectss = []
        for i in range(len(nameid)):
            for j in range(len(frag_connect[nameid[i]])):
                frag_connects = []
                frag_connects.append(frag_connect[nameid[i]][j][0] + atom_count)
                frag_connects.append(frag_connect[nameid[i]][j][1] + atom_count)
                frag_connectss.append(frag_connects)
            atom_count += sum(frag_atom[nameid[i]])

        self.fatomnums = frag_atoms
        self.fchgs = frag_charges
        self.fbaas = frag_baanums
        self.fatminfos = frag_atmlabss
        self.connects = frag_connectss

        logger.debug("frag_atoms(head) %s ...", frag_atoms[:20])
        logger.debug("frag_charges(head) %s ...", frag_charges[:20])
        logger.debug("frag_baanums(head) %s ...", frag_baanums[:20])
        logger.debug("frag_atmlabss(head) %s ...", frag_atmlabss[:5])
        logger.debug("frag_connectss(head) %s ...", frag_connectss[:5])

        return self
        # return frag_atoms, frag_charges, frag_baanums, frag_atmlabss, frag_connectss

    def writemb_frag_section(self, param: list[list[Any]], nameid: list[int]) -> str:
        """マルチボディ用AJFフラグメントセクションの文字列を生成する。

        Args:
            param: [frag_atom, frag_charge, frag_connect_num, frag_connect, seg_info] のリスト。
            nameid: セグメント名IDのリスト。

        Returns:
            str: AJF形式のフラグメントセクション文字列。
        """
        frag_atom = param[0]
        frag_charge = param[1]
        frag_connect_num = param[2]
        frag_connect = param[3]
        seg_info = param[4]

        ajf_fragment = ''
        # fragment atom
        count = 0
        for i in range(len(nameid)):
            for j in range(len(frag_atom[nameid[i]])):
                    ajf_fragment += '%8d' % (frag_atom[nameid[i]][j])
                    count += 1
                    if count % 10 == 0:
                        ajf_fragment += '\n'
        if count % 10 != 0:
            ajf_fragment += '\n'

        # fragment charge
        count = 0
        for i in range(len(nameid)):
            for j in range(len(frag_charge[nameid[i]])):
                ajf_fragment += '%8d' % (frag_charge[nameid[i]][j])
                count += 1
                if count % 10 == 0:
                    ajf_fragment += '\n'
        if count % 10 != 0:
            ajf_fragment += '\n'

        # fragment connect num
        count = 0
        for i in range(len(nameid)):
            for j in range(len(frag_connect_num[nameid[i]])):
                ajf_fragment += '%8d' % (frag_connect_num[nameid[i]][j])
                count += 1
                if count % 10 == 0:
                    ajf_fragment += '\n'
        if count % 10 != 0:
            ajf_fragment += '\n'

        # fragment body
        atom_count = 0
        for i in range(len(nameid)):
            for j in range(len(seg_info[nameid[i]])):
                icount = 0
                for k in range(len(seg_info[nameid[i]][j])):
                    icount += 1
                    ajf_fragment += '%8d' % (
                            seg_info[nameid[i]][j][k] + atom_count)
                    if icount % 10 == 0:
                        icount = 0
                        ajf_fragment += '\n'
                if icount % 10 != 0:
                    ajf_fragment += '\n'
            atom_count += sum(frag_atom[nameid[i]])

        logger.debug("atom_count %s", atom_count)
        atom_count = 0
        # connect info
        logger.debug("frag_atom %s", frag_atom)
        for i in range(len(nameid)):
            for j in range(len(frag_connect[nameid[i]])):
                ajf_fragment += '%8d' % (
                        frag_connect[nameid[i]][j][0] + atom_count)
                ajf_fragment += '%8d' % (
                        frag_connect[nameid[i]][j][1] + atom_count)
                if len(frag_connect[nameid[i]][j]) == 3:
                    ajf_fragment += '%8d' % (frag_connect[nameid[i]][j][2])
                ajf_fragment += '\n'
            atom_count += sum(frag_atom[nameid[i]])

        return ajf_fragment


    def saveajf(self, oname: str | None = None) -> None:
        """AJFファイルを生成して保存する。

        Args:
            oname: 出力ファイル名。Noneの場合は自動生成される。
        """
        if oname is None:
            oname = os.path.splitext(self.readgeom)[0] + '-' + self.ajf_method + '-' + self.ajf_basis_set.replace('*', 'd') + '.ajf'
            ohead = os.path.splitext(self.readgeom)[0] + '-' + self.ajf_method
        else:
            ohead = os.path.splitext(oname)[0]
        ajf_body = self.gen_ajf_bodywrap(ohead)

        logger.info("ajf_oname: %s", oname)
        with open(oname, 'w') as ajf_file:
            print(ajf_body, file=ajf_file)
