pcc

#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# @Time    : 2018-04-11  15:34
# @Author  : PengZhw
# @Software: PyCharm
import os
from typing import List
# from for_simulation.classes import *
from sys import argv
# from for_simulation.classes import Point, Flat
from os import path
import csv
import re
import shutil
from classes import *  # 这样写是因为Linux系统能识别同目录下的模组
import shutil
import time



terminal_columns, terminal_lines = os.get_terminal_size()

name = '=    PCC    ='
print()
print('\n\n')


print(name.center(terminal_columns-1))


print("Version 2021.07.15".center(terminal_columns-1))
print('\n\n')

# print("\n===========\n= 程序pcc =\n===========\n")

# 自身文件名，数据文件路径，石墨烯规格，输出文件名,溶剂原子

if '-f' in argv:
    FILE_PREFIX = argv[1 +  argv.index('-f')]
else:
    raise Exception("缺少-f（文件名前缀）参数")

# 是否是预处理
if '-pretreatment' in argv:
    PRETREATMENT = True
else:
    PRETREATMENT = False

# 不是预处理模式就接收下面的参数
if not PRETREATMENT:
    if '-s' in argv:
        SIZE = argv[1+argv.index('-s')]
    else:
        raise Exception("缺少-s（石墨烯尺寸）参数，３、５、１０")

    if '-o' in argv:
        OUTPUT_NAME = argv[1+argv.index('-o')]
    else:
        raise Exception("缺少-o（输出文件名）参数")

    if '-t' in argv:
        SOLVENT = argv[1+argv.index('-t')]
    else:
        raise Exception("缺少-t（溶剂中心原子在pdb文件中的名称）参数")
    
    print("文件名前缀%s，石墨烯尺寸%s，输出文件名前缀%s，被计算原子%s。"%(FILE_PREFIX, SIZE, OUTPUT_NAME, SOLVENT))

self = argv[0]
# FILE_PATH = argv[1]
# FILE_PREFIX = argv[1]
# SIZE = argv[2]
# OUTPUT_NAME = argv[3]

FILE_PATH = os.getcwd()

# 列表：所有待处理的文件，用re来获取
FILE_LIST = []
for file in os.listdir(FILE_PATH):

    match = re.match(FILE_PREFIX + '[0-9]*.pdb', file)
    if str(match) != 'None':
        FILE_LIST.append(os.path.join(FILE_PATH, match.group()))
        # print(file)

re_digits = re.compile(r'(\d+)')

def emb_numbers(s):
    pieces = re_digits.split(s)
    pieces[1::2] = map(int, pieces[1::2])
    return pieces


def sort_strings_with_emb_numbers(alist):
    aux = [(emb_numbers(s), s) for s in alist]
    aux.sort()
    return [s for __, s in aux]


def sort_strings_with_emb_numbers2(alist):
    return sorted(alist, key=emb_numbers)


FILE_LIST = sort_strings_with_emb_numbers(FILE_LIST)


print("文件列表读入完毕。")

# 测试
# main()

# 定义石墨烯片层列表
GRAPHENES = []

# 定义石墨烯片层间隙列表
GRABOXES = []

# 定义碳原子列表
GRACS = []

# 定义二氧化碳列表
CO2CS = []

# 记录溶剂分子的个数
SOLV_NUM = []


def read_data(file_path):
        
    with open(file_path, 'r') as FILE_IN:
        file_lines = FILE_IN.readlines()

    s = 1
    sheep_c_list = []

    for line in file_lines:
        sl = line.split()  # sl = sl，切割后的一行，是一个list

        # 将石墨烯碳原子的信息存入list里面
        if (sl[0] == 'ATOM') and (sl[2] == 'C'):
            if int(sl[4]) == s:
                sheep_c_list.append(
                    GrapheneC(float(sl[5]), float(sl[7]), float(sl[6]), s, float(sl[1])))  # x, z, y, 第几层, 名字（pdb里面的总序号）

            elif int(sl[4]) > s:
                s = int(sl[4])
                GRACS.append(sheep_c_list)
                sheep_c_list = []
        elif (sl[0] == 'ATOM') and (sl[2] != 'C'):
            GRACS.append(sheep_c_list)
            break

    for line in file_lines:
        sl = line.split()  # sl = sl，切割后的一行，是一个list

        # 将溶剂中指定原子的信息存入list里面
        if (sl[0] == 'ATOM') and (sl[2] == str(SOLVENT)):
            try:
                CO2CS.append(CO2C(float(sl[5]), float(sl[7]), float(sl[6]), float(sl[1]), int(sl[4])))
            except TypeError:
                print("数据类型错误，这通常意味着第五列不是原子序号而是被其他字符占领了。这一行是\n\
                %s\n请确保文件无误。" % line)
            
    '''
    print("石墨烯片层数：", len(GRACS))
    print("第一片石墨烯上第一个碳原子坐标[x, z, y] =", GRACS[0][0].coordinate)

    print("溶剂分子(二氧化碳)的个数：", len(CO2CS))
    print("第二个溶剂分子中中心原子"+str(SOLVENT)+"坐标[x, z, y] =", CO2CS[1].coordinate)
    '''

# 石墨烯片层中六边形的边长为1.451A
L = 1.451


# print(type(GRACS[1][1]))


def corner_3(cl):
    """
    接收单层石墨烯的碳原子列表，计算得到四个Point对象
    :param cl: 单层石墨烯的碳原子列表，相当于前面的{list}sheep_c_list、{list}GRACS[X]
    :return: 四个Point对象的list
    """
    k = len(cl)  # 这一层有k个碳原子
    c3 = cl[2]  # 第3个碳原子
    c8 = cl[7]  # 第8个碳原子
    cn3 = cl[k - 3]  # 倒数第3个碳原子
    cn8 = cl[k - 8]  # 倒数第8个碳原子
    pa = Point(c3.x - (c8.x - c3.x) / 2, c3.z - (c8.z - c3.z) / 2, c3.y - (c8.y - c3.y) / 2)
    pb = Point(c8.x - (c3.x - c8.x) / 8, c8.z - (c3.z - c8.z) / 8, c8.y - (c3.y - c8.y) / 8)
    pc = Point(cn3.x - (cn8.x - cn3.x) / 2, cn3.z - (cn8.z - cn3.z) / 2, cn3.y - (cn8.y - cn3.y) / 2)
    pd = Point(cn8.x - (cn3.x - cn8.x) / 8, cn8.z - (cn3.z - cn8.z) / 8, cn8.y - (cn3.y - cn8.y) / 8)
    point_list = [pa, pb, pc, pd]  # 用来盛放四个角点并返回
    return point_list


def corner_5(cl) -> list:
    """
    接收单层石墨烯的碳原子列表，计算得到四个Point对象
    :param cl: 单层石墨烯的碳原子列表，相当于前面的{list}sheep_c_list、{list}GRACS[X]
    :return: 四个Point对象的list
    """
    k = len(cl)  # 这一层有k个碳原子
    c3 = cl[2]  # 第3个碳原子
    c16 = cl[15]  # 第16个碳原子
    cn3 = cl[k - 3]  # 倒数第3个碳原子
    cn16 = cl[k - 16]  # 倒数第16个碳原子
    pa = Point(c3.x - (c16.x - c3.x) / 5, c3.z - (c16.z - c3.z) / 5, c3.y - (c16.y - c3.y) / 5)
    pb = Point(c16.x - (c3.x - c16.x) / 20, c16.z - (c3.z - c16.z) / 20, c16.y - (c3.y - c16.y) / 20)
    pc = Point(cn3.x - (cn16.x - cn3.x) / 5, cn3.z - (cn16.z - cn3.z) / 5, cn3.y - (cn16.y - cn3.y) / 5)
    pd = Point(cn16.x - (cn3.x - cn16.x) / 20, cn16.z - (cn3.z - cn16.z) / 20, cn16.y - (cn3.y - cn16.y) / 20)
    point_list = [pa, pb, pc, pd]  # 用来盛放四个角点并返回
    return point_list

def corner_7(cl)-> list:
    """
    接收单层石墨烯的碳原子列表，计算得到四个Point对象
    :param cl: 单层石墨烯的碳原子列表，相当于前面的{list}sheep_c_list、{list}GRACS[X]
    :return: 四个Point对象的list
    """
    k = len(cl)  # 这一层有k个碳原子
    c0 = cl[0]
    c3 = cl[2]  # 第3个碳原子
    c24 = cl[23]  # 第24个碳原子

    cn0= cl[k-1] #倒数第1个C原子
    cn3 = cl[k - 3]  # 倒数第3个碳原子
    cn24 = cl[k - 24]  # 倒数第24个碳原子

    # index = 0 的点的x，=2的点的z, 两者y的平均值
    pa = Point(c0.x, c3.z, (c3.y + c0.y)/2 )
    pb = Point(c24.x, c24.z, c24.y)
    pc = Point(cn0.x , cn3.z, (cn3.y + cn0.y)/2)
    pd = Point(cn24.x, cn24.z, cn24.y)
    point_list = [pa, pb, pc, pd]  # 用来盛放四个角点并返回
    return point_list

def corner_10(cl) -> list:
    """
    接收单层石墨烯的碳原子列表，计算得到四个Point对象
    :param cl: 单层石墨烯的碳原子列表，相当于前面的{list}sheep_c_list、{list}GRACS[X]
    :return: 四个Point对象的list
    """
    k = len(cl)  # 这一层有k个碳原子
    c3 = cl[2]  # 第3个碳原子
    c36 = cl[35]  # 第36个碳原子
    cn3 = cl[k - 3]  # 倒数第3个碳原子
    cn36 = cl[k - 36]  # 倒数第36个碳原子
    pa = Point(c3.x - (c36.x - c3.x) / 12.5, c3.z - (c36.z - c3.z) / 12.5, c3.y - (c36.y - c3.y) / 12.5)
    pb = Point(c36.x - (c3.x - c36.x) / 50, c36.z - (c3.z - c36.z) / 50, c36.y - (c3.y - c36.y) / 50)
    pc = Point(cn3.x - (cn36.x - cn3.x) / 12.5, cn3.z - (cn36.z - cn3.z) / 12.5, cn3.y - (cn36.y - cn3.y) / 12.5)
    pd = Point(cn36.x - (cn3.x - cn36.x) / 50, cn36.z - (cn3.z - cn36.z) / 50, cn36.y - (cn3.y - cn36.y) / 50)
    point_list = [pa, pb, pc, pd]  # 用来盛放四个角点并返回
    return point_list

def corner_x(cl, x: int):
    """
        接收单层石墨烯的碳原子列表，计算得到四个Point对象
        :param cl: 单层石墨烯的碳原子列表，相当于前面的{list}sheep_c_list、{list}GRACS[X]
        :param x: 从原子序号0-10这条边有多少个环，相当于以前的 3/5/7/10
        :return: 四个Point对象的list
        """
    # 定义按原子index的第一条边上，除第3个原子(index=2)之外的第二个定点的序号
    # 实际上是等于4(x-1)  如 3--8,5--16,7--24,10--36
    n = 4*(x-1)

    
    k = len(cl)  # 这一层有k个碳原子
    corner_1 = cl[2]  # 第3个碳原子
    corner_2 = cl[n-1]  # 第n个碳原子
    corner_4 = cl[k - 3]  # 倒数第3个碳原子
    corner_3 = cl[k - n]  # 倒数第n个碳原子
    pa = Point(corner_1.x - (corner_2.x - corner_1.x) / 12.5, corner_1.z - (corner_2.z - corner_1.z) / 12.5, corner_1.y - (corner_2.y - corner_1.y) / 12.5)
    pb = Point(corner_2.x - (corner_1.x - corner_2.x) / 50, corner_2.z - (corner_1.z - corner_2.z) / 50, corner_2.y - (corner_1.y - corner_2.y) / 50)
    pc = Point(corner_4.x - (corner_3.x - corner_4.x) / 12.5, corner_4.z - (corner_3.z - corner_4.z) / 12.5, corner_4.y - (corner_3.y - corner_4.y) / 12.5)
    pd = Point(corner_3.x - (corner_4.x - corner_3.x) / 50, corner_3.z - (corner_4.z - corner_3.z) / 50, corner_3.y - (corner_4.y - corner_3.y) / 50)
    point_list = [pa, pb, pc, pd]  # 用来盛放四个角点并返回
    return point_list

corner_list = []  # {list}corner_list是二维list，第一维是几个石墨烯片层，第二维是每一层的四个角点Point对象


def read_3():
    for i in range(len(GRACS)):
        l = corner_3(GRACS[i])
        corner_list.append(l)
    # print('a')


def read_5():
    for i in range(len(GRACS)):
        l = corner_5(GRACS[i])
        corner_list.append(l)
    # print('a')

def read_7():
    for i in range(len(GRACS)):
        l = corner_7(GRACS[i])
        corner_list.append(l)
    # print('a')


def read_10():
    for i in range(len(GRACS)):
        l = corner_10(GRACS[i])
        corner_list.append(l)
    # print('a')


def read_x(x: int):
    for i in range(len(GRACS)):
        l = corner_x(GRACS[i])
        corner_list.append(l)


# read_data('npt02-d-55-6-6.pdb')





def final_cal():
    """
    if SIZE == '3':
        read_3()
    elif SIZE == '5':
        read_5()
    elif SIZE == '7':
        read_7()
    elif SIZE == '10':
        read_10()
    else:
        print("错误: 没有正确输入石墨烯尺寸参数。")

    # print("第一层石墨烯有%d个碳原子" % len(GRACS[0]))
    """
    read_x(int(SIZE))
    """
    for i in range(len(corner_list)):
        for j in range(4):
            print(corner_list[i][j].coordinate)
    """
    for i in range(len(corner_list)):
        GRAPHENES.append(Flat(corner_list[i][0], corner_list[i][1], corner_list[i][2], corner_list[i][3]))

    # print(len(GRAPHENES))
    # GRAPHENES[0].tell_corner()

    outlines = []

    for i in range(len(GRAPHENES) - 1):
        GRABOXES.append(Box(GRAPHENES[i], GRAPHENES[i + 1]))
        print("\r第%d、%d层石墨烯构成了一个盒子。" % (i + 1, i + 2))
        outlines.append("\r\r\r[  Box%d  ]" % (i + 1))
        outlines.append("\r#第%d、%d层石墨烯构成的一个盒子。" % (i + 1, i + 2))

        outlines.append("\r%12s%12s%12s%12s%12s"
                        % ('an', 'rn', 'x', 'y', 'z'))
        n = 0
        for j in range(len(CO2CS)):
            if GRABOXES[i].if_inside(CO2CS[j]):
                n += 1
                '''测试代码块
                print("原子序号%d，残基序号%d，坐标[x, z, y] = %s。"
                      % (CO2CS[j].name,
                         CO2CS[j].mol_n,
                         str(CO2CS[j].coordinate)))
                '''
                outlines.append("\r%12d%12d%12.3f%12.3f%12.3f"
                                % (CO2CS[j].name,
                                   CO2CS[j].mol_n,
                                   CO2CS[j].x,
                                   CO2CS[j].y,
                                   CO2CS[j].z))
        outlines[-n - 2] = '\r#总溶剂分子数%12s' % str(n) + outlines[-n - 2]
        SOLV_NUM.append(n)

    with open(OUTPUT_NAME + ".txt", "w+") as FILE_OUT:
        if path.isfile(OUTPUT_NAME+".txt"):
            shutil.copyfile(OUTPUT_NAME + ".txt", '#'+OUTPUT_NAME + ".txt#")
            print("已备份文件%s -> %s"%(OUTPUT_NAME + ".txt", '#'+OUTPUT_NAME + ".txt#"), "备份文件被覆盖后将永久丢失，若希望保留备份文件，请在下次运行本程序之前将其改名。")
        FILE_OUT.writelines(outlines)
        print("成功输出文件%s" % OUTPUT_NAME+".txt")

    '''测试代码块
    with open(OUTPUT_NAME, 'r') as f:
        print(f.read(), os.getcwd())    
    print(SOLV_NUM)    
    #with open(r'test.csv', 'a') as f:
    # with open(OUTPUT_NAME, 'r') as f:
    #    print(f.read(), os.getcwd())
    '''

    
    # print(SOLV_NUM)
    with open(OUTPUT_NAME + '.csv', 'a') as f:
        writer = csv.writer(f)
        writer.writerow(SOLV_NUM)

    
def cal():

    """
    # print("选择读取方法。")
    if SIZE == '3':
        read_3()
    elif SIZE == '5':
        read_5()
    elif SIZE == '7':
        read_7()
    elif SIZE == '10':
        read_10()
    else:
        print("错误: 没有正确输入石墨烯尺寸参数。")
    """

    # print("第一层石墨烯有%d个碳原子" % len(GRACS[0]))

    read_x(int(SIZE))

    """
    for i in range(len(corner_list)):
        for j in range(4):
            print(corner_list[i][j].coordinate)
    """
    for i in range(len(corner_list)):
        GRAPHENES.append(Flat(corner_list[i][0], corner_list[i][1], corner_list[i][2], corner_list[i][3]))

    # print(len(GRAPHENES))
    # GRAPHENES[0].tell_corner()

    outlines = []

    for i in range(len(GRAPHENES) - 1):
        GRABOXES.append(Box(GRAPHENES[i], GRAPHENES[i + 1]))
        print("\r第%d、%d层石墨烯构成了一个盒子。" % (i + 1, i + 2))
        outlines.append("\r\r\r[  Box%d  ]" % (i + 1))
        outlines.append("\r#第%d、%d层石墨烯构成的一个盒子。" % (i + 1, i + 2))

        outlines.append("\r%12s%12s%12s%12s%12s"
                        % ('an', 'rn', 'x', 'y', 'z'))
        n = 0
        
        for j in range(len(CO2CS)):
            if GRABOXES[i].if_inside(CO2CS[j]):
                n += 1

                # print("原子序号%d，残基序号%d，坐标[x, z, y] = %s。"
                #       % (CO2CS[j].name,
                #          CO2CS[j].mol_n,
                #          str(CO2CS[j].coordinate)))
                
                outlines.append("\r%12d%12d%12.3f%12.3f%12.3f"
                                % (CO2CS[j].name,
                                   CO2CS[j].mol_n,
                                   CO2CS[j].x,
                                   CO2CS[j].y,
                                   CO2CS[j].z))
        outlines[-n - 2] = '\r#总溶剂分子数%12s' % str(n) + outlines[-n - 2]
        SOLV_NUM.append(n)


    # print(SOLV_NUM)
    with open(OUTPUT_NAME + '.csv', 'a') as f:
        writer = csv.writer(f)
        writer.writerow(SOLV_NUM)


# re_complie = re.compile(r'^((?:.+\s+){4})(\w+)(.+)$')
re_complie = re.compile(r'^((?:[0-9a-zA-Z]+\s+){4})([a-zA-Z]+)(.+)$')
# 预处理函数
def pretreatment(file_path):
    with open(file_path, 'r') as FILE_IN:
        file_lines = FILE_IN.readlines()
    print('读入文件...', file_path)
    for n, line in enumerate(file_lines):     # mr = match_result
        sl = line.split()
        mr = re.match(re_complie, line)
        if str(mr) != 'None' and sl[0] == 'ATOM':
            if '-print' in argv:
                print(file_lines[n]+'替换为\n'+mr.group(1) + " "*len(mr.group(2)) + mr.group(3))
            # print(mr.group(1))
            # print(mr.group(2))
            # print(mr.group(3))
            file_lines[n] = mr.group(1) + " "*len(mr.group(2)) + mr.group(3) + '\n'
            # quit()

    # 备份
    path, basename = os.path.split(file_path)
    short_name, extension = os.path.splitext(basename)

    time_str = time.strftime("#%Y%m%d_%H-%M-%S#",time.localtime(time.time()))
    shutil.copy(file_path, os.path.join(path, time_str+short_name+extension))
    print('*' * 10 + "\n备份文件" + basename + "至" + os.path.join(path, time_str+short_name+extension))

    # 写入
    with open(file_path, 'w+')as fw:
        fw.writelines(file_lines)


if PRETREATMENT:
    print("\n预处理……\n")
    for n, file_path in enumerate(FILE_LIST):
        pretreatment(file_path)


else:
    print("计算……")
    if path.isfile(OUTPUT_NAME+".csv"):
            shutil.copyfile(OUTPUT_NAME + ".csv", '#'+OUTPUT_NAME + ".csv#")
            print("已备份文件%s -> %s"%(OUTPUT_NAME + ".csv", '#'+OUTPUT_NAME + ".csv#"), "备份文件被覆盖后将永久丢失，若希望保留备份文件，请在下次运行本程序之前将其改名。")

    with open(OUTPUT_NAME + '.csv', 'w') as f:
        writer = csv.writer(f)
        writer.writerow(["Mol Num in Boxes:"])

    # print(len(FILE_LIST))

    for n, file_path in enumerate(FILE_LIST):
        read_data(file_path)
        print('读入文件...', n, file_path)
        if n < len(FILE_LIST) - 1:
            cal()
            # print("文件'test.csv'已写入。")
        elif n == len(FILE_LIST) - 1:
            final_cal()
        
        
        # 初始化石墨烯片层列表
        GRAPHENES = []

        # 初始化石墨烯片层间隙列表
        GRABOXES = []

        # 初始化碳原子列表
        GRACS = []

        # 初始化二氧化碳列表
        CO2CS = []

        # 初始化溶剂分子的个数
        SOLV_NUM = []

        # 初始化corner_list
        corner_list = []  # {list}corner_list是二维list，第一维是几个石墨烯片层，第二维是每一层的四个角点Point对象