#!/usr/bin/env python
# -*- coding: utf-8 -*-

#PBS -N job_name
#PBS -o job.log
#PBS -e job.err
#PBS -q gold5120
#PBS -l nodes=1:ppn=8
#PBS -l walltime=24:00:00

# Abacus from ase
from ase.calculators.abacus.abacus_out import Abacus

# Cspd
from ase.calculators.abacus.cspd import atomic_structure_generator

# Spap
from ase.calculators.abacus.spap import run_spap

# User
import os
import numpy as np
from math import ceil, pi
import signal
# from ase.io import read

# Please change the path below
# WorkPath = os.getcwd() +'/' # Set working path
WorkPath = '/home/shenzx/project/python_20190718/v20200729/ase-abacus/example/highflux/'
DbFile = '/home/shenzx/software/cspd/CSPD.db'  # Set cspd database
# Set abacus potential path(tc46001)
PotentialPath = '/home/shenzx/software/abacus/SG15_ONCV_PBE_1.0'
# Set abacus orbit path(tc46001)
OrbitPath = '/home/shenzx/software/abacus/Orb_DZP_E100_Standard_v1.0'
# Set abacus running path
os.environ['ASE_ABACUS_COMMAND'] = "mpirun -np 8 /home/shenzx/software/abacus/abacus_v2.0/bin/ABACUS.mpi.2.0>>PREFIX.log"
os.chdir(WorkPath)


def AddText(FileName="running.log", String="\n"):
    # Add log to file
    f = open(FileName, 'a')
    f.write(str(String) + "\n")
    f.close()


def GetKpoints(Stru=None, Kspacing=0.5):
    # Get Kpoints according to Stru and Kspacing
    Kpoints = [ceil(2 * pi / Kspacing *
                    np.linalg.norm(Stru.get_reciprocal_cell()[i])) for
               i in range(3)]
    return Kpoints


def TimeOut(long_function_call, time=1):
    # Time out set by user
    def signal_handler(signum, frame):
        raise Exception("First time out")
    signal.signal(signal.SIGALRM, signal_handler)
    signal.alarm(time)
    try:
        long_function_call()
    except Exception:
        return 1


# Create  structures
# images = list()
# for root, dirs, files in os.walk("structure_folder", topdown = False):
#     files.sort()
#     for name in files:
#         images.append(read(os.path.join(root, name), format = 'cif'))
images = atomic_structure_generator(
    symbols='VSe2',
    lw=True,
    format='cif',
    fu=[1, 1],
    cspd_file=DbFile,
    volume=60.0,
    to_primitive=True,
    maxatomn=60,
)
AddText("running.log", 'CSPD Create images successfully')


# Calculation
for image in images[:5]:
    # Set Abacus running path
    # count =  files[images.index(image)].split('_')[0]
    count = images.index(image)
    AseRunDir = WorkPath + "file" + str(count)
    AddText("running.log", "Structure {} running in path: \n{}".format(
        count, AseRunDir))
    if (not os.path.exists(AseRunDir)):
        os.mkdir(AseRunDir)
    os.chdir(AseRunDir)
    Kpoints1 = GetKpoints(image, 0.5)  # Get Kpoints
    # Let all images use an abacus calculator:
    image.set_calculator(Abacus(
        atoms=image,
        pseudo_dir=PotentialPath,
        potential_name="PotSG15",
        basis_dir=OrbitPath,
        basis_name="SG15act",
        ntype=2,
        ecutwfc=30,
        niter=100,
        nstep=400,
        dr2="1.0e-4",
        calculation='cell-relax',
        basis_type='lcao',
        ks_solver='genelpa',
        mixing_type='pulay-kerker',
        mixing_beta=0.4,
        force=1,
        force_thr_ev=0.50,
        stress=1,
        stress_thr=5.0,
        move_method='cg',
        trust_radius_ini=0.5,
        smearing='gaussian',
        sigma=0.2,
        atom_file='STRU',
        out_stru=1,
        out_charge=0,
        gamma_only=0,
        knumber=0,
        kmode='Gamma',
        kpts=[Kpoints1[0], Kpoints1[1], Kpoints1[2], 0, 0, 0]
    ))
    if (TimeOut(image.get_potential_energy, 21600) == 1):
        os.chdir("../")
        AddText("running.log", "Timed out, stop")
        continue

    os.chdir("../")
    AddText("running.log", "Potential Energy: eV")
    # Get potential energy to log file
    AddText("running.log", image.get_potential_energy())

run_spap(
    total_struc=1500,
    l_cif=True,
    l_db=False,
    # r_cut_off=7.5,
    # symprec=0.02,
    # l_poscar=True,
    # threshold=0.045,
    # l_comp=False,
    # ilat=1,
    work_dir='./',
    i_mode=8,
)
exit()