# Path: siman/chg/vasputil_chgarith_module.py
# def chgarith(chgf1, chgf2, op, filename, wcell):
# 
#     # chgf1 = args[0]
#     # chgf2 = args[2]
#     # op = args[1]
# 
#     chg1 = VaspChargeDensity(chgf1)
#     chg2 = VaspChargeDensity(chgf2)
# 
#     # if options.wcell:
#     #     wcell = int(options.wcell)
#     if wcell == 0:
#         chga = chg1
#     elif wcell == 1:
#         chga = chg2
#     #     else:
#     #         print ('Error, invalid argument to -w option')
#     #         sys.exit(1)
#     # else:
#     # chga = chg1
# 
#     if len(chg1.chg) != len(chg2.chg):
#         print ('Number of images in charge density files not equal. Using just   the final images in both files.')
#         print ('len(chg.chg)', len(chg1.chg), len(chg2.chg) )
#         chg1.chg = [chg1.chg[-1]]
#         chg1.atoms = [chg1.atoms[-1]]
#         if chg1.is_spin_polarized():
#             chg1.chgdiff = [chg1.chgdiff[-1]]
#             chg2.chgdiff = [chg2.chgdiff[-1]]
#         chg2.chg = [chg2.chg[-1]]
#         chg2.atoms = [chg2.atoms[-1]]
# 
#     newchg = VaspChargeDensity(None)
# 
#     print ('Start charge manipul')
# 
#     for i, atchg in enumerate(chg1.chg):
#         c1 = atchg
#         c2 = chg2.chg[i]
#         newchg.atoms.append(chga.atoms[i].copy())
#         if op == '+':
#             nc = c1 + c2
#             oplong = '_add_'
#         elif op == '-':
#             nc = c1 - c2
#             oplong = '_sub_'
#         elif op == '*':
#             nc = c1 * c2
#             oplong = '_mult_'
#         elif op == '/':
#             nc = c1 / c2
#             oplong = '_div_'
#         elif op == 'avg':
#             nc = (c1 + c2) / 2
#             oplong = '_avg_'
#         newchg.chg.append(nc)
# 
# 
#     if chg1.is_spin_polarized():
#         print ('Spin polarized')
# 
#         for i, cd in enumerate(chg1.chgdiff):
#             cd2 = chg2.chgdiff[i]
#             if op == '+':
#                 nd = cd + cd2
#             elif op == '-':
#                 nd = cd - cd2
#             elif op == '*':
#                 nd = cd * cd2
#             elif op == '/':
#                 nd = cd / cd2
#             elif op == 'avg':
#                 nd = (cd + cd2) / 2
#             newchg.chgdiff.append(nd)
# 
#     # Screw doing anything fancy with the augmentation charges
#     # Just take them from the same file as the embedded atoms object.
#     newchg.aug = chga.aug
#     newchg.augdiff = chga.augdiff
# 
#     # if options.outfile:
#     #     fname = options.outfile
#     # else:
#     #     from os.path import basename
#     #     fname = basename(chgf1) + oplong + basename(chgf2)
# 
#     newchg.write(filename)
# 
#     return filename


import sys, os
import shutil
import numpy as np
    from ase.calculators.vasp import VaspChargeDensity
from siman.header import runBash, printlog
from siman.chg.vasputil_chgarith_module import chgarith
    from siman.picture_functions import fit_and_plot
    files = []
    for cl in cl1, cl2, cl3:
        if cl is None:
            continue
        if chg == 'CHGCAR':
            file = cl.get_chg_file(nametype  = 'asoutcar')
            if not file:
                printlog('No CHGCAR for cl',cl.id[0], 'trying CHG', imp = 'Y')
                file = cl.get_chg_file('CHG', nametype  = 'asoutcar')
        elif chg == 'PARCHG':
            file = cl.get_file('PARCHG')


        files.append(file)



    file1 = files[0]
    file2 = files[1]


    if file1 == None or file2 == None:
        printlog('Error!, chg not found for one cl:', file1, file2)


    working_dir = cl1.dir

    dendiff_filename = working_dir + (chg+'_'+str(cl1.id[0])+'-'+str(cl2.id[0])).replace('.', '_')

    printlog('Diff =', file1, '-', file2)

# Path: siman/chg/vasputil_chgarith_module.py
# def chgarith(chgf1, chgf2, op, filename, wcell):
# 
#     # chgf1 = args[0]
#     # chgf2 = args[2]
#     # op = args[1]
# 
#     chg1 = VaspChargeDensity(chgf1)
#     chg2 = VaspChargeDensity(chgf2)
# 
#     # if options.wcell:
#     #     wcell = int(options.wcell)
#     if wcell == 0:
#         chga = chg1
#     elif wcell == 1:
#         chga = chg2
#     #     else:
#     #         print ('Error, invalid argument to -w option')
#     #         sys.exit(1)
#     # else:
#     # chga = chg1
# 
#     if len(chg1.chg) != len(chg2.chg):
#         print ('Number of images in charge density files not equal. Using just   the final images in both files.')
#         print ('len(chg.chg)', len(chg1.chg), len(chg2.chg) )
#         chg1.chg = [chg1.chg[-1]]
#         chg1.atoms = [chg1.atoms[-1]]
#         if chg1.is_spin_polarized():
#             chg1.chgdiff = [chg1.chgdiff[-1]]
#             chg2.chgdiff = [chg2.chgdiff[-1]]
#         chg2.chg = [chg2.chg[-1]]
#         chg2.atoms = [chg2.atoms[-1]]
# 
#     newchg = VaspChargeDensity(None)
# 
#     print ('Start charge manipul')
# 
#     for i, atchg in enumerate(chg1.chg):
#         c1 = atchg
#         c2 = chg2.chg[i]
#         newchg.atoms.append(chga.atoms[i].copy())
#         if op == '+':
#             nc = c1 + c2
#             oplong = '_add_'
#         elif op == '-':
#             nc = c1 - c2
#             oplong = '_sub_'
#         elif op == '*':
#             nc = c1 * c2
#             oplong = '_mult_'
#         elif op == '/':
#             nc = c1 / c2
#             oplong = '_div_'
#         elif op == 'avg':
#             nc = (c1 + c2) / 2
#             oplong = '_avg_'
#         newchg.chg.append(nc)
# 
# 
#     if chg1.is_spin_polarized():
#         print ('Spin polarized')
# 
#         for i, cd in enumerate(chg1.chgdiff):
#             cd2 = chg2.chgdiff[i]
#             if op == '+':
#                 nd = cd + cd2
#             elif op == '-':
#                 nd = cd - cd2
#             elif op == '*':
#                 nd = cd * cd2
#             elif op == '/':
#                 nd = cd / cd2
#             elif op == 'avg':
#                 nd = (cd + cd2) / 2
#             newchg.chgdiff.append(nd)
# 
#     # Screw doing anything fancy with the augmentation charges
#     # Just take them from the same file as the embedded atoms object.
#     newchg.aug = chga.aug
#     newchg.augdiff = chga.augdiff
# 
#     # if options.outfile:
#     #     fname = options.outfile
#     # else:
#     #     from os.path import basename
#     #     fname = basename(chgf1) + oplong + basename(chgf2)
# 
#     newchg.write(filename)
# 
#     return filename


# Path: siman/chg/chg_func.py
import sys, os
import shutil
import numpy as np
    from ase.calculators.vasp import VaspChargeDensity
from siman.header import runBash, printlog
from siman.chg.vasputil_chgarith_module import chgarith
    from siman.picture_functions import fit_and_plot

    files = []
    for cl in cl1, cl2, cl3:
        if cl is None:
            continue
        if chg == 'CHGCAR':
            file = cl.get_chg_file(nametype  = 'asoutcar')
            if not file:
                printlog('No CHGCAR for cl',cl.id[0], 'trying CHG', imp = 'Y')
                file = cl.get_chg_file('CHG', nametype  = 'asoutcar')
        elif chg == 'PARCHG':
            file = cl.get_file('PARCHG')


        files.append(file)



    file1 = files[0]
    file2 = files[1]


    if file1 == None or file2 == None:
        printlog('Error!, chg not found for one cl:', file1, file2)


    working_dir = cl1.dir

    dendiff_filename = working_dir + (chg+'_'+str(cl1.id[0])+'-'+str(cl2.id[0])).replace('.', '_')

    printlog('Diff =', file1, '-', file2)

    chgarith(file1, file2, '-', dendiff_filename, wcell)