from __future__ import annotations
"""XYZファイルおよびPDB形式の分子構造データの読み書きを行うモジュール。"""

import os
import re
import logging

logger = logging.getLogger(__name__)


class mol_io():
    """XYZ形式やPDB形式の分子構造ファイルの入出力を行うクラス。"""

    def __init__(self) -> None:
        super().__init__()
        # print('## load mol io')

    def read_mol_name(self, fname: str) -> str:
        """ファイル名から拡張子を除いた分子名を取得する。

        Args:
            fname: ファイル名。

        Returns:
            拡張子を除いたファイル名（分子名）。
        """
        molname, ext = os.path.splitext(fname)
        # lines = [l.rstrip() for l in open(filename, "U")]
        # text = lines[1]
        # molname = re.sub(' {1,}', '', text).split('@@')
        # molname = molname[1]
        return molname

    def read_xyz(self, filename: str) -> list[list[str] | list[list[float]]]:
        """XYZファイルを読み込み、原子種と座標を返す。

        Args:
            filename: XYZファイルのパス。

        Returns:
            [原子種リスト, 座標リスト] の2要素リスト。
        """
        with open(filename) as _fh:
            lines = [l.rstrip() for l in _fh]
        atom = []
        coord = []
        for l in range(2, len(lines)):
            temp = []
            temp.append(re.sub(' {1,}', ' ', lines[l]).split())
            atom.append(temp[0][0])
            temp[0].pop(0)
            temp2 = []
            for ll in range(len(temp[0])):
                temp2.append(float(temp[0][ll]))
            coord.append(temp2)
        return [atom, coord]

    def getatoms(self, fname: str) -> tuple[list[list[list[int | str]]], list[str], list[list[list[str]]]]:
        """複数構造を含むXYZファイルを読み込み、構造ごとの原子情報と座標を返す。

        Args:
            fname: 複数構造XYZファイルのパス。

        Returns:
            (原子情報リスト, 原子数リスト, 座標リスト) のタプル。
        """
        with open(fname, "r", newline="\n") as f:
            text = f.readlines()

        atom = []
        pos = []

        atoms = []
        atomnums = []
        poss = []
        skipflag = False
        for i in range(len(text)):
            if skipflag:
                skipflag = False
                continue

            itemList = text[i][:-1].split()
            if len(itemList) == 1:
                if i != 0:
                    atoms.append(atom)
                    poss.append(pos)
                    atomnums.append(atomnum)
                    atom = []
                    pos = []

                atomnum = itemList[0]
                skipflag = True
            else:
                atom.append([i-1, itemList[0]])
                pos.append(itemList[1:4])
            if i == (len(text) -1):
                atoms.append(atom)
                poss.append(pos)
                atomnums.append(atomnum)


        return atoms, atomnums, poss

    def convert_xyzs_pdb(self, fname: str, dirname: str, tgtnum: int = 0) -> None:
        """複数構造XYZファイルの各構造を個別のPDBファイルに変換する。

        Args:
            fname: 入力XYZファイルのパス。
            dirname: 出力先ディレクトリ。
            tgtnum: 特定構造のみ変換する場合の番号（0で全構造）。
        """
        atoms, atomnums, poss = self.getatoms(fname)
        head, ext = os.path.splitext(fname)
        for i in range(len(atoms)):
            oname = dirname + '/' + str(i+1).zfill(5) + ".pdb"
            # print(atoms[i])
            # print(atomnums[i])
            # print(poss[i])
            if tgtnum != 0:
                if i+1 == tgtnum:
                    # print(oname)
                    self.Exportpospdb(atoms[i], atomnums[i], poss[i], oname)
            else:
                    # print(oname)
                    self.Exportpospdb(atoms[i], atomnums[i], poss[i], oname)

    def getatom(self, fname: str) -> tuple[list[list[int | str]], str, list[list[str]]]:
        """単一構造のXYZファイルを読み込み、原子情報と座標を返す。

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

        Returns:
            (原子情報リスト, 原子数, 座標リスト) のタプル。
        """
        atom = []
        pos = []
        with open(fname, "r", newline="\n") as f:
            text = f.readlines()
        for i in range(len(text)):
            itemList = text[i][:-1].split()
            if i == 0:
                atomnum = itemList[0]
                # print(fname, atomnum)
            if i >= 2:
                atom.append([i-2, itemList[0]])
                pos.append(itemList[1:4])
        return atom, atomnum, pos

    def convert_xyz_pdb(self, fname: str) -> None:
        """単一構造のXYZファイルをPDBファイルに変換する。

        Args:
            fname: 入力XYZファイルのパス。
        """
        atom, atomnum, pos = self.getatom(fname)
        head, ext = os.path.splitext(fname)
        oname = str(head) + ".pdb"
        self.Exportpospdb(atom, atomnum, pos, oname)

    def Exportpospdb(self, atom: list[list[int | str]], atomnum: str, pos: list[list[str]], out_file: str) -> None:
        """原子情報と座標をPDB形式でファイルに出力する。

        Args:
            atom: 原子インデックスと元素名のペアのリスト。
            atomnum: 原子数（文字列）。
            pos: 各原子の座標リスト。
            out_file: 出力PDBファイルのパス。
        """
        # # Export position of mol
        # head, ext = os.path.splitext(str(iname))

        # header
        # print(out_file)
        with open(out_file, "w", newline = "\n") as f:
            print("COMPND    " + out_file, file=f)
            print("AUTHOR    " + "GENERATED BY python script in ABMPTools", file=f)

        # aaa = [0.8855]
        # print '{0[0]:.3f}'.format(aaa)
        # print pos[0]
        with open(out_file, "a+", newline = "\n") as f:
            for i in range(len(pos)):
                for j in range(3):
                    pos[i][j] = float(pos[i][j])

            # print pos[0]

            for i in range(len(atom)):
                l1_head='HETATM'
                l2_lab=str(i+1)
                l3_atom=atom[i][1]
                l3atom2="  "
                l4_alt=" "
                l5_res="UNK"
                l6_chain=" "
                l7_labres="  "
                l8_code=" "
                l9_x='{:.3f}'.format(pos[i][0])
                l10_y='{:.3f}'.format(pos[i][1])
                l11_z='{:.3f}'.format(pos[i][2])
                l12_occ='1.00'
                l13_temp="0.00"
                l14_ele=atom[i][1]
                l15_cha='  '
                list = [l1_head, l2_lab, l3_atom, l3atom2, l4_alt, l5_res, l6_chain, l7_labres,  l8_code, l9_x, l10_y, l11_z, l12_occ, l13_temp, l14_ele, l15_cha]
                print('{0[0]:<6}{0[1]:>5} {0[2]:>2}{0[3]:<2}{0[4]:>1}{0[5]:>3} {0[6]:>1}{0[7]:>4}{0[8]:>1}   {0[9]:>8}{0[10]:>8}{0[11]:>8}{0[12]:>6}{0[13]:>6}          {0[14]:>2}{0[15]:>2}'.format(list), file=f)
            # ATOM      1  H   UNK     1     -12.899  32.293   3.964  1.00  0.00           H
            print("END", file=f)

#1 1 – 6  HETATM
#2 7 – 11 原子の通し番号
# blank 1
#3 13 – 16    原子名
#4 17  Alternate location識別子
#5 18 - 20 残基名
# blank 1
#6 22  鎖名
#7 23 - 26 残基番号
#8 27  残基の挿入コード
# blank 3
#9 31 - 38 原子のX座標の値（Å単位）
#10 39 - 46 原子のY座標の値（Å単位）
#11 47 - 54 原子のZ座標の値（Å単位）
#12 55 - 60 占有率
#13 61 - 66 温度因子
#14 77 - 78 元素記号
#15 79 - 80 原子の電荷


