#!/bin/python
# -*- coding:utf-8 -*-
#Convert XDATCAR to PDB and extract energy & temperature profile for AIMD simulations
# -h for help;-b for frame started;-e for frame ended;;-p for PDB conversion  
#by nxu tamas@zju.edu.cn
#version 1.2
#date 2019.4.9

import os
import shutil
import numpy as np
import matplotlib as mpl
import math
mpl.use('Agg') #silent mode
from matplotlib import pyplot as plt
from optparse import OptionParser  
import sys
from copy import deepcopy

class debug(object):
    def __init__(self,info='debug'):
        self.info=info
    def __call__(self,func):
        def wrapper(*args,**kwargs):
            print('[{info}] Now entering function {function}.....' \
                .format(info=self.info,function=getattr(func,"__name__")))
            return func(*args,**kwargs)
        return wrapper

class Energy_Temp(object):
    def __init__(self):
        'Read vasp MD energies and temperature.'
        self.temp=[]
        self.energy=[]
        self.energy_extract()
            
    def energy_extract(self):
        print('Now reading vasp MD energies and temperature.')        
        if os.path.exists("OSZICAR"):
            oszicar=open("OSZICAR",'r')
            for index,line in enumerate(oszicar):
                if "E0=" in line:
                    self.temp.append(float(line.split()[2]))
                    self.energy.append(float(line.split()[4]))
            oszicar.close()
            #return (self.temp,self.energy)                         
        else:
            raise IOError('OSZICAR does not exist!')

class XDATCAR(Energy_Temp):
    def __init__(self):
        'Read lattice information and atomic coordinate.'
        super(XDATCAR,self).__init__()
        print('Now reading vasp XDATCAR.')       
        self.lattice=np.zeros((3,3))
        self.NPT=False
        self.frames=0
        self._timestep=1  # timestep 1fs
        self.alpha=0.0
        self.beta=0.0
        self.gamma=0.0
        self.current_frame=0
        
        self._lattice1=0.0
        self._lattice2=0.0
        self._lattice3=0.0
        self._lattice=np.zeros(3)
        self.format_trans=False

        @property 
        def timestep(self):
            return self._timestep
        @timestep.setter
        def timestep(self,new_time_step):
            self._timestep=new_time_step
            
        if os.path.exists("XDATCAR"):
            self.XDATCAR=open("XDATCAR",'r')
        else:
            raise IOError('XDATCAR does not exist!')
        title=self.XDATCAR.readline().strip()
        for index,line in enumerate(self.XDATCAR):
            if "Direct" in line:
                self.frames+=1
            # elif title in line:
            #     self.NPT=True
        self.XDATCAR.seek(0)
        self.lattice_read()
        self.XDATCAR.readline()
        for i in range(self.total_atom):
            self.XDATCAR.readline()

        if "Direct configuration" not in self.XDATCAR.readline():
        	self.NPT=True	
        #self.frames=len(['Direct' for line in self.XDATCAR if "Direct" in line])
        print('Total frames {0}, NpT is {1}'.format(self.frames,self.NPT))
        # assert len(self.energy)==self.frames, \
        #     'Number of XDATCAR frames does not equal to that of Energy terms in OSZICAR'   
        self.XDATCAR.seek(0)
        self.lowrange=0;self.uprange=self.frames-1
        if self.NPT == False: self.lattice_read()
        
    def step_select(self,selected_step): # 't>100ps and t < 1000ps'
        'eg. t > 100 and t < 1000'
        assert isinstance(selected_step,str),'Selected timestep must be in a "string"'
        if 'and' in selected_step:
            conditions=selected_step.split("and")
        else:
            conditions=[selected_step]
        for condition in conditions:
            try:
                if '>=' in condition:
                    ranges=int(condition.split('>=')[1].strip())
                    self.lowrange=ranges-1
                elif '<=' in condition:
                    ranges=int(condition.split('<=')[1].strip())
                    self.uprange=ranges-1
                elif '>' in condition:
                    ranges=int(condition.split('>')[1].strip())
                    self.lowrange=ranges
                elif '<' in condition:
                    ranges=int(condition.split('<')[1].strip())
                    self.uprange=ranges-2
                else:
                    print('Selected timestep is invaid!');continue
            except ValueError:
                print('Selected timestep is invaid!') 
                sys.exit(0)
            if (self.lowrange >= self.frames-1) or (self.uprange < 0):
                raise ValueError('Selected timestep is wrong!')
            if self.lowrange < 0: self.lowrange= 0
            if self.uprange > self.frames-1: self.uprange= self.frames-1          
                
    def lattice_read(self):    
        self.title=self.XDATCAR.readline().rstrip('\r\n').rstrip('\n')
        self.scaling_factor=float(self.XDATCAR.readline())
        for i in range(3):
            self.lattice[i]=np.array([float(j) for j in self.XDATCAR.readline().split()])
        self.lattice*=self.scaling_factor
        self._lattice1=np.sqrt(np.sum(np.square(self.lattice[0])))
        self._lattice2=np.sqrt(np.sum(np.square(self.lattice[1])))
        self._lattice3=np.sqrt(np.sum(np.square(self.lattice[2])))
        self._lattice[0]=self._lattice1
        self._lattice[1]=self._lattice2
        self._lattice[2]=self._lattice3
        self.alpha=math.acos(np.dot(self.lattice[1],self.lattice[2]) \
            /float((self._lattice2*self._lattice3)))/np.pi*180.0
        self.beta=math.acos(np.dot(self.lattice[0],self.lattice[2]) \
            /float((self._lattice1*self._lattice3)))/np.pi*180.0
        self.gamma=math.acos(np.dot(self.lattice[0],self.lattice[1]) \
            /float((self._lattice1*self._lattice2)))/np.pi*180.0
        self.element_list=[j for j in self.XDATCAR.readline().split()]
        try:
            self.element_amount=[int(j) for j in self.XDATCAR.readline().split()]
        except ValueError:
            raise ValueError('VASP 5.x XDATCAR is needed!')
        self.total_elements=[]
        for i in range(len(self.element_amount)):
            self.total_elements.extend([self.element_list[i]]*self.element_amount[i])
        #print(self.total_elements)
        self.total_atom=sum(self.element_amount)
        self.atomic_position=np.zeros((self.total_atom,3))

    def __iter__(self):
        return self

    def __next__(self):
        self.next()

    def skiplines_(self):
        if self.NPT == True: self.lattice_read() 
        self.XDATCAR.readline()
        for i in range(self.total_atom):
            self.XDATCAR.readline()

    def next(self):
        if self.NPT == True: self.lattice_read() 
        self.XDATCAR.readline()
        for i in range(self.total_atom):
            line_tmp=self.XDATCAR.readline()
            self.atomic_position[i]=np.array([float(j) for j in line_tmp.split()[0:3]])         
        self.cartesian_position=np.dot(self.atomic_position,self.lattice)
        self.current_frame+=1
        #
        #self.cartesian_position*=self.scaling_factor
        return self.cartesian_position

    def writepdb(self,pdb_frame):
        tobewriten=[]
        tobewriten.append("MODEL         %r" %(pdb_frame))
        tobewriten.append("REMARK   Converted from XDATCAR file")
        tobewriten.append("REMARK   Converted using VASPKIT")
        tobewriten.append('CRYST1{0:9.3f}{1:9.3f}{2:9.3f}{3:7.2f}{4:7.2f}{5:7.2f}' .format(self._lattice1,\
            self._lattice2,self._lattice3,self.alpha,self.beta,self.gamma))
        #print(len(self.total_elements))
        for i in range(len(self.total_elements)):
            tobewriten.append('%4s%7d%4s%5s%6d%4s%8.3f%8.3f%8.3f%6.2f%6.2f%12s' \
                %("ATOM",i+1,self.total_elements[i],"MOL",1,'    ',self.cartesian_position[i][0],\
                self.cartesian_position[i][1],self.cartesian_position[i][2],1.0,0.0,self.total_elements[i]))
        tobewriten.append('TER')
        tobewriten.append('ENDMDL')
        with open('XDATCAR.pdb','a+') as pdbwriter:
            tobewriten=[i+'\n' for i in tobewriten]
            pdbwriter.writelines(tobewriten)

    def unswrapPBC(self,prev_atomic_cartesian):
        diff= self.cartesian_position-prev_atomic_cartesian
        prev_atomic_cartesian=deepcopy(self.cartesian_position) 
        xx=np.where(diff[:,0]>(self._lattice1/2),diff[:,0]-self._lattice1\
            ,np.where(diff[:,0]<-(self._lattice1/2),diff[:,0]+self._lattice1\
                ,diff[:,0]))
        yy=np.where(diff[:,1]>(self._lattice2/2),diff[:,1]-self._lattice2\
            ,np.where(diff[:,1]<-(self._lattice2/2),diff[:,1]+self._lattice2\
                ,diff[:,1]))
        zz=np.where(diff[:,2]>(self._lattice3/2),diff[:,2]-self._lattice3\
            ,np.where(diff[:,2]<-(self._lattice3/2),diff[:,2]+self._lattice3\
                ,diff[:,2]))
        xx=xx.reshape(-1,1);yy=yy.reshape(-1,1);zz=zz.reshape(-1,1)
        return (prev_atomic_cartesian,np.concatenate([xx,yy,zz],axis=1))

    def reset_cartesian(self,real_atomic_cartesian,center_atom):
        if center_atom > len(real_atomic_cartesian)-1:
            raise SystemError("Selected atom does not exist!")
        for i in range(0,len(real_atomic_cartesian)):
            for j in range(3):
                if (real_atomic_cartesian[i][j]-real_atomic_cartesian[center_atom][j])>self._lattice[j]/2:
                    real_atomic_cartesian[i][j]-=self._lattice[j]
                if (real_atomic_cartesian[i][j]-real_atomic_cartesian[center_atom][j])<-self._lattice[j]/2:
                    real_atomic_cartesian[i][j]+=self._lattice[j]                
        return real_atomic_cartesian
                    
    def __call__(self,selected_step):
        self.step_select(selected_step)
    def __str__(self):
        return ('Read lattice information and atomic coordinate.')
    __repr__=__str__
    
class plot(object):
    'Plot MD temperature and energy profile'
    def __init__(self,lwd,font,dpi,figsize,XDATCAR_inst=None):
        self.XDATCAR_inst=XDATCAR_inst;self.timestep=self.XDATCAR_inst.timestep
        self.time_range=(self.XDATCAR_inst.lowrange,self.XDATCAR_inst.uprange+1)
        self.lwd=lwd;self.font=font;self.dpi=dpi;self.figsize=figsize

    #@debug(info='debug')
    def plotfigure(self,title='MD temperature and energy profile'):
        from matplotlib import pyplot as plt
        self.newtemp=self.XDATCAR_inst.temp;self.newenergy=self.XDATCAR_inst.energy
        xdata=np.arange(self.time_range[0],self.time_range[1])*self.timestep
        axe = plt.subplot(121)
        self.newtemp=self.newtemp[self.time_range[0]:self.time_range[1]]
        axe.plot(xdata,self.newtemp, \
                 color='black', lw=self.lwd, linestyle='-', alpha=1)
        with open("Temperature.dat",'w') as writer:
            writer.write("Time(fs)   Temperature(K)\n")
            for line in range(len(xdata)):
                writer.write("{0:f}  {1:f}\n" .format(xdata[line],self.newtemp[line]))
        axe.set_xlabel(r'$Time$ (fs)',fontdict=self.font)
        axe.set_ylabel(r'$Temperature$ (K)',fontdict=self.font)
        axe.set_xlim((self.time_range[0]*self.timestep, self.time_range[1]*self.timestep))
        axe.set_title('MD temperature profile')
        
        axe1 = plt.subplot(122)
        self.newenergy=self.newenergy[self.time_range[0]:self.time_range[1]]
        axe1.plot(xdata,self.newenergy, \
                 color='black', lw=self.lwd, linestyle='-', alpha=1)
        with open("Energy.dat",'w') as writer:
            writer.write("Time(fs)   Energy(ev)\n")
            for line in range(len(xdata)):
                writer.write("{0:f}  {1:f}\n" .format(xdata[line],self.newenergy[line]))        
        axe1.set_xlabel(r'$Time$ (fs)',fontdict=self.font)
        axe1.set_ylabel(r'$Energy$ (ev)',fontdict=self.font)
        axe1.set_xlim((self.time_range[0]*self.timestep, self.time_range[1]*self.timestep))
        axe1.set_title('MD energy profile')
        #plt.suptitle(title)
        fig = plt.gcf()
        fig.set_size_inches(self.figsize)
        plt.tight_layout() 
        plt.savefig('ENERGY.png',bbox_inches='tight',pad_inches=0.1,dpi=self.dpi)
        
        
if __name__ == "__main__":
    parser = OptionParser()  
    parser.add_option("-b", "--begin",  
                       dest="begin", default=1,  
                      help="frames begin") 

    parser.add_option("-e", "--end",  
                      dest="end", default='false',  
                      help="frames end") 

    parser.add_option("-t", "--timestep",  
                      dest="timestep", default=1.0,  
                      help="timestep per frame") 

    parser.add_option("-p", "--pdb",  
                      action="store_true", dest="format_trans", default=False,  
                      help="choose whether to convert XDATCAR to PDB!") 

    parser.add_option("--pbc",  
                      action="store_true", dest="periodic", default=False,  
                      help="choose whether to swrap PBC images!") 

    parser.add_option("-i", "--index", 
                      dest="index", default=-1,  
                      help="choose which atom to center whole molecule!") 

    parser.add_option("--interval", 
                      dest="interval", default=1,  
                      help="extract frames interval!") 

    (options,args) = parser.parse_args()
    try:
        float(options.timestep)
        int(options.index)
        _interval=int(options.interval)
        int(options.begin)
        if options.end != 'false': int(options.end)
    except:
        raise ValueError('wrong arguments')
    XDATCAR_inst=XDATCAR()
    XDATCAR_iter=iter(XDATCAR_inst)
    if options.format_trans:
        XDATCAR_inst.format_trans=True
        if os.path.exists("XDATCAR.pdb"):
            shutil.copyfile("XDATCAR.pdb","XDATCAR-bak.pdb")
            os.remove("XDATCAR.pdb")
    else:
        XDATCAR_inst.format_trans=False
    XDATCAR_inst.timestep=float(options.timestep)   #timestep 1fs
    if options.end == 'false':
        XDATCAR_inst('t>= %r' %(int(options.begin)))
    else:
        XDATCAR_inst('t>= %r and t <= %r' %(int(options.begin),int(options.end))) # frame 10~300  corresponding to 20~600fs
    if _interval >(XDATCAR_inst.uprange-XDATCAR_inst.lowrange):
        raise SystemError('Trajectory interval exceed selected time range!')
    elif _interval<=1:
        _interval=1
    count=0 
    current_pdb=1  
    for i in range(XDATCAR_inst.uprange+1): 
        if (i>=XDATCAR_inst.lowrange):
            cartesian_position=XDATCAR_iter.next()
            if count % _interval == 0:
                if options.format_trans == True: 
                    if options.periodic == True: 
                        if i == XDATCAR_inst.lowrange:
                            real_atomic_cartesian=deepcopy(cartesian_position)
                            if int(options.index) != -1:
                                real_atomic_cartesian=XDATCAR_inst.reset_cartesian(real_atomic_cartesian,int(options.index)-1)
                            XDATCAR_inst.cartesian_position=real_atomic_cartesian
                            prev_atomic_cartesian=deepcopy(cartesian_position)
                        else:
                            prev_atomic_cartesian,diffs=XDATCAR_inst.unswrapPBC(prev_atomic_cartesian)
                            real_atomic_cartesian+=diffs
                            XDATCAR_inst.cartesian_position=real_atomic_cartesian              
                    XDATCAR_inst.writepdb(current_pdb)
                    current_pdb+=1
            count+=1
        else:
            XDATCAR_iter.skiplines_()

    newtemp=XDATCAR_inst.temp
    newenergy=XDATCAR_inst.energy
    print('Finish reading XDATCAR.')

    timestep=XDATCAR_inst.timestep
    print('Selected time-range:{0}~{1}fs'.format((XDATCAR_inst.lowrange)*timestep,\
                        (XDATCAR_inst.uprange)*timestep))
    XDATCAR_inst.XDATCAR.close()
    print('Timestep for new PDB trajectory is :{0}fs'.format(timestep*_interval))  
    lwd = 0.2  # Control width of line
    dpi=300          # figure_resolution
    figsize=(5,4)   #figure_inches
    font = {'family' : 'arial', 
        'color'  : 'black',
        'weight' : 'normal',
        'size' : 13.0,
        }       #FONT_setup
    plots=plot(lwd,font,dpi,figsize,XDATCAR_inst)
    plots.plotfigure()

(base) [storm@X16 vaspkit]$ python 1.py
Now reading vasp MD energies and temperature.
Now reading vasp XDATCAR.
Total frames 15000, NpT is False
Finish reading XDATCAR.
Selected time-range:0.0~14999.0fs
Timestep for new PDB trajectory is :1.0fs
findfont: Font family 'arial' not found.
findfont: Font family ['arial'] not found. Falling back to DejaVu Sans.
findfont: Font family 'arial' not found.
findfont: Font family 'arial' not found.
findfont: Font family 'arial' not found.

AttributeError: `np.mat` was removed in the NumPy 2.0 release. Use `np.asmatrix` instead

pip install numpy==1.26.4

from decimal import Decimal, getcontext
import re
from sympy import symbols, Matrix, linsolve, Rational
from collections import defaultdict

# Improve Decimal precision
getcontext().prec = 100

class FormulaParser:
    @staticmethod
    def parse(formula: str) -> dict:
        elements = {}
        pattern = re.compile(r"([A-Z][a-z]?)([\d.]*)?")
        idx = 0
        while idx < len(formula):
            m = pattern.match(formula, idx)
            if not m:
                raise ValueError(f"Invalid formula at: {formula[idx:]}")
            elem, num = m.group(1), m.group(2)
            coeff = Decimal(num) if num else Decimal(1)
            elements[elem] = elements.get(elem, Decimal(0)) + coeff
            idx = m.end()
        return elements

class AtomicMassReader:
    @staticmethod
    def read(filename="atomic_masses.txt") -> dict:
        masses = {}
        try:
            with open(filename, 'r') as f:
                for line in f:
                    line = line.strip()
                    if not line or line.startswith('#'):
                        continue
                    parts = line.split()
                    if len(parts) != 2:
                        raise ValueError(f"Line format error: {line}")
                    masses[parts[0]] = Decimal(parts[1])
        except FileNotFoundError:
            raise FileNotFoundError(f"File {filename} not found")
        if not masses:
            raise ValueError("Atomic masses file empty")
        return masses

class SynthesisCalculator:
    def __init__(self, atomic_masses: dict):
        self.atomic_masses = atomic_masses
        self.target = {}
        self.target_mm = Decimal(0)
        self.reagents = []
        self.reagent_formulas = []
        self.molar_masses = []

    def input_target(self):
        f = input("Enter target formula: ").strip()
        self.target = FormulaParser.parse(f)
        for e, c in self.target.items():
            if e not in self.atomic_masses:
                raise ValueError(f"Mass of element {e} not defined")
            self.target_mm += c * self.atomic_masses[e]
        print(f"Parsed: {self.target}")
        print(f"Target molar mass = {self.target_mm.normalize()} g/mol")

    def input_reagents(self):
        n = int(input("Enter number of reagents: "))
        for i in range(n):
            f = input(f"Reagent {i+1} formula: ").strip()
            self.reagent_formulas.append(f)
            parsed = FormulaParser.parse(f)
            mm = sum(parsed[e] * self.atomic_masses[e] for e in parsed)
            self.reagents.append(parsed)
            self.molar_masses.append(mm)
            print(f"{f} molar mass = {mm.normalize()} g/mol")

    def validate(self):
        tgt = set(self.target)
        prov = set().union(*(r.keys() for r in self.reagents))
        miss = tgt - prov
        extra = prov - tgt
        if miss:
            print(f"Error: Missing elements: {miss}")
            exit()
        if extra:
            print(f"Warning: Extra elements: {extra}")

    def solve_basis(self):
        # Construct A, b corresponding to 1 mol 
        elems = sorted(set(self.target) | set().union(*(r.keys() for r in self.reagents)))
        A = [[r.get(e, Decimal(0)) for r in self.reagents] for e in elems]
        b = [self.target.get(e, Decimal(0)) for e in elems]
        # Convert to rational number to avoid floating point error
        A_rat = [[Rational(str(val)) for val in row] for row in A]
        b_rat = [Rational(str(val)) for val in b]
        vars = symbols(f'x0:{len(self.reagents)}')
        sol = linsolve((Matrix(A_rat), Matrix(b_rat)), vars)
        if not sol:
            print("Error: No exact solution for 1 mol target.")
            exit()
        tup = next(iter(sol))
        # Convert rational to decimal exact representation
        base = {}
        for i, r in enumerate(tup):
            base[vars[i]] = Decimal(r.p) / Decimal(r.q)
        return base, elems

    def calculate(self):
        t = input("Input type (0 mass g, 1 moles): ").strip()
        if t not in ['0','1']:
            print("Invalid type")
            exit()
        v = Decimal(input("Amount: ").strip())
        n = v / self.target_mm if t=='0' else v
        print(f"Target moles = {n.normalize()}")

        base, elems = self.solve_basis()
        print("\nReagent requirements:\n" + "="*40)
        total_mass = Decimal(0)
        for i, var in enumerate(sorted(base.keys(), key=lambda x: int(str(x)[1:]))):
            moles = base[var] * n
            mass = moles * self.molar_masses[i]
            total_mass += mass
            print(f"Reagent {i+1} ({self.reagent_formulas[i]}):")
            print(f"  Moles: {moles.normalize():.10f} mol")
            print(f"  Mass:  {mass.normalize():.10f} g\n")
        print(f"Total mass = {total_mass.normalize():.10f} g")

        #Verify element balance
        print("Verification:")
        for e in elems:
            actual = sum((base[symbols(f'x{i}')] * n) * r.get(e, Decimal(0))
                         for i, r in enumerate(self.reagents))
            target_amt = self.target.get(e, Decimal(0)) * n
            print(f"{e}: target {target_amt.normalize():.10f}, actual {actual.normalize():.10f}")


def main():
    print("High-Precision Compound Synthesis Calculator")
    print("="*60)
    try:
        masses = AtomicMassReader.read()
        print(f"Loaded elements: {', '.join(masses.keys())}")
    except Exception as e:
        print(e)
        exit()
    calc = SynthesisCalculator(masses)
    calc.input_target()
    calc.input_reagents()
    calc.validate()
    calc.calculate()

if __name__ == "__main__":
    main()