C2DB数据库是根据已知二维材料替换元素得到扩充数据路，其包含PBE、HSE、结构稳定性、convexhull以及真空能级等性质，可以作为研究二维材料的数据库之一，这里给出其数据库中的结构提取脚本，方便大家前处理遇到的困苦，加快自己的ideal的推进。方法比较简单粗暴莽夫哲学（亚历山大砍戈耳狄俄斯之结）(https://zh.wikipedia.org/wiki/%E4%BA%9A%E5%8E%86%E5%B1%B1%E5%A4%A7%E5%A4%A7%E5%B8%9D)
）。
最近C2DB做了更新，目前我基于的是其2018年的文章中涉及的数据库,之前还能下，现在也找不到下载的地方了。

'''
Created on 14:35, Mar. 25, 2021

@author: yilinZhang 
@address: Jilin University

These files are needed: c2db.db
'''
import numpy as np
import pandas as pd
import os


def prehandle_c2db():
    import json
    import sqlite3
    # Step1： connect database and get the features
    con = sqlite3.connect("c2db.db")
    df = pd.read_sql_query("SELECT * FROM systems", con)

    # Step2: convert database to csv
    key_value_pairs = []
    for line in df.key_value_pairs:
        key_value_pairs.append(json.loads(line))
    key_value_pairs = pd.DataFrame(key_value_pairs)

    # Step3: rename formula
    formula_values = list(key_value_pairs['folder'].str.split('/'))
    flag = []
    for line in formula_values:
        flag.append(line[-2])
    key_value_pairs['formula'] = flag

    # Step4： sorted the vbm
    key_value_pairs['vbm_to_evac'] = key_value_pairs['vbm'] - key_value_pairs['evac']
    # key_value_pairs['cbm_to_evac'] = key_value_pairs['cbm'] - key_value_pairs['evac']
    hasBandGap = key_value_pairs[key_value_pairs['vbm_to_evac'] <= 0]#肯定要满足价带不能超过静电势能吧
    # hasBandGap = hasBandGap[hasBandGap['gap'] > 0.3]

    hasBandGap = hasBandGap[hasBandGap['ehull'] <= 0]# convexhull也要满足一下
    # hasBandGap = hasBandGap[hasBandGap['dynamic_stability_level'] == 3]
    hasBandGap = hasBandGap.sort_values(by=['vbm_to_evac'], axis=0)
    hasBandGap.to_csv('sorted_by_vbm_to_evac.csv')


def writePOSCAR(atom):

    with open(path + os.sep + filename + '_POSCAR', 'w') as f:
        f.write(self.parameters['SystemName'] + '\n')
        f.write(str(self.parameters['ScalingFactor']) + '\n')

        LatticeMatrix = self.parameters['LatticeMatrix']
        for i in range(len(LatticeMatrix)):
            f.write(float2line(LatticeMatrix[i]))

        AtomsInfo = self.parameters['AtomsInfo']
        string = ['', '']
        for key in AtomsInfo.keys():
            string[0] += key + '    '
            string[1] += str(AtomsInfo[key]) + '    '
        string[0] += '\n'
        string[1] += '\n'
        f.write(string[0])
        f.write(string[1])

        f.write(self.parameters['CoordinateType'] + '\n')

        ElementsPositionMatrix = self.parameters['ElementsPositionMatrix'].copy(
        )
        for key in ElementsPositionMatrix.keys():
            arr = ElementsPositionMatrix[key]
            for i in range(len(arr)):
                f.write(float2line(arr[i]))


def output_structure():
    import ase.db
    from ase.io import write
    dbpath = os.path.join(os.getcwd(), 'c2db.db')
    db = ase.db.connect(dbpath)

    def float2line(mart):
        '''
        将矩阵转换成可输出的字符串形式
        '''
        string = ''
        for line in mart:
            for s in line:
                string += '%21.10f' % float(s)
            string += '\n'
        return string

    def count_element(mart):
        string = ['', '']
        elements = sorted(set(mart), key=mart.index)
        for key in elements:
            string[0] += key + '    '
            string[1] += str(mart.count(key)) + '    '
        string[0] += '\n'
        string[1] += '\n'
        return string

    # with open('bandAlignment_CSV_names', 'r') as f:
    #     lines = f.readlines()

    # for i in range(len(lines)):
    #     os.chdir(lines[i].strip()[:-4])
    csv = pd.read_csv('sorted_by_vbm_to_evac.csv')
    for i in range(len(csv)):
        atom = db.get_atoms(uid=csv.iloc[i]['uid'], add_additional_information=True)
        write('{}.cif'.format(csv.iloc[i]['formula']), atom, format='cif')
        with open(csv.iloc[i]['formula'] + '.vasp', 'w') as f:
            f.write(csv.iloc[i]['uid'] + '\n')
            f.write('1.0\n')
            f.write(float2line(atom.get_cell()))
            atomInfo = count_element(atom.get_chemical_symbols())
            f.write(atomInfo[0])
            f.write(atomInfo[1])
            f.write('Cartesian\n')
            f.write(float2line(atom.get_positions()))


if __name__ == "__main__":
    prehandle_c2db()
    # after_handle()
    output_structure()
    os.system('rm *.cif')