import chemsh

from ..qm import _QM
from ...dl_py2f  import utils
from ...dl_py2f.utils import fileutils
from ...io    import file

import shutil
from os import system, path

class CASTEP(_QM):

    ''' Interface to CASTEP '''
    ''' All castep options are passed as strings, except dimensionless integer
options, e.g., cut_off_energy='200 eV', not cut_off_energy=200 '''
    # Keywords for castep input 
    # __init__ of parent class is called which parse _kwargs to simple
    # attributes
#todo: fine_grid_scale:, remove fix_com

    _kwargs = {
               'basis_de_dloge'    :'',
               'basis_precision'   : '',
               'castep_exe'        : 'castep.serial', #default serial. mpirun -n number castep.mpi can be provided
               'castep_task'       : 'singlepoint',
               'checkpoint'        : 'test.check',
               'continuation'      : 'NULL', #'default',
               'cut_off_energy'    : '',               #125.0
               'elec_convergence_win' : 0,              #3
               'elec_dump_file'    :'NULL',
               'elec_energy_tol'   : '', #10e-05,  
               'electronic_minimizer' : 'CG', # 'SD'
               'energy_unit'       : 'ha',
               'finite_basis_corr' : 'None',
               'finite_basis_npoints' : 0,  #3
               'finite_basis_spacing' : 0,  #5
               'fix_occupancy'     : 'True',
               'force_unit'        : 'hartree/bohr',
               'geom_method'       : '', #'bfgs',
               'grid_scale'        : 2,
               'input'             : ('_castep.param','_castep.cell'),
               'iprint'            : 0, #2,
               'length_unit'       : 'bohr',
               'maxiter'           : 60,
               'metals_method'     : '',
               'mixing_scheme'     : '',
               'nbands'            : 0,
               'num_backup_iter'   : 0, #5,
               'num_dump_cycles'   : '0',
               'opt_strategy'      : 'speed',
               'output'            : '_castep.castep',
               'pseudopotential'   : None,
               'popn_write'        : 'minimal', 
               'rand_seed'         : '',
               'smearing_width'    : 0,
               # Raj: this option is not written explicitly to input, only used for 
               # truth value;no quotes
               'symmetry_generate' : False ,
               'write_checkpoint'  : 'None',
               'write_formatted_density' : 'False',
              }              

    ''' If castep was compiled for serial build, castep.serial will be the
        default command. User must rename the binary in castep/bin/[your arch]
    '''

    _internals = {
                       '_syscmd'        :'castep.serial' , #default system command to run
#        Raj: Not parsing output files when called as library
                       '_parseoutfile'  : True,
                       '_gradients'     : True,
                 }

    @property
    def runargs(self):    
        runargs=(str(fileutils.getBaseName(self.input[0])))
        return runargs

    def setOutKeys(self):
        '''Define {task:keyword} for searching in output file'''
        # each keyword should contain an array of integers specifying:
        #     [ row_0, # starting row
        #       row_n, # ending row
        #       row_i, # row increment (>=1)
        #       col_0, # starting col
        #       col_n, # ending col
        #       col_i, # col increment (>=1)
        #       skip , # skip the previous (skip-1) match,
        #       coeff, # unit conversion coefficient
        #     ]
        # referencing to the line where the key is matched
        self._outkeys = {'energy' : {'Final energy = ' : [0,0,1,3,3,1,1,1.0] },
                         'gradients' : {'Forces' : [6,6+self.frag.natoms-1,1,3,5,1,1,1.0] },  
                       }

    def sort_in_input_order(self, arraytosort):

        from numpy import ndarray, asarray
        from numpy import argsort 

        # Raj: array to be returned
        values = ndarray(arraytosort.shape)
        unsorted = asarray(arraytosort)

        # is in the order given as input
        elements = self.frag.znums
        sorter = argsort(elements)

        for i in range(self.frag.natoms):
            values[sorter[i]] = unsorted[i]
        return values

    def sort_in_atom_order(self, arraytosort):

        from numpy import ndarray, asarray
        from numpy import argsort, append

        values = ndarray(arraytosort.shape)
        unsorted = asarray(arraytosort)
        elements = self.frag.znums
        sorter = argsort(elements)
        values = unsorted[sorter] 

        return values

    def parseMainOutput(self):

        '''Parse main output file '''
        from numpy import negative 

        if(self._parseoutfile):
               self.setOutKeys()
               self._result.energy = file.getArrayFromFile(self.output, 'out',
                                     self._result.energy.shape,regex=self._outkeys['energy'])[0]
               #print('\033[91m result shape',self._result.energy.shape)
        else:
               self._result.energy = self._result.energy

        if self._gradients:
            if(self._parseoutfile):
                try:
                    grads_searched = file.getArrayFromFile(self.output, 'out',
                              self._result.gradients.shape, regex=self._outkeys['gradients'])
                except:
                    print("\n >>> ChemShell/CASTEP: Failed parsing gradients from output:\n")
                    raise
            else:
             
             # Raj: If run as a library, result.gradients is updated with forces from castep
             # sorted in Z order
                grads_searched = self._result.gradients
             
            # Raj: sorts in the order in input , then gradients=negative forces
            self._result.gradients = negative(self.sort_in_input_order(grads_searched))

    def getInpBuff(self, input=input, **kwargs):

        ''' the input files are created from the input kwargs,
         strbuff is updated with all required input-information
        '''
        from numpy import array, append, shape, vstack, transpose, size
        # prevent reading result from old results
        frag   = self.frag
        cell   = frag.cell
        consts = cell.consts
        strbuff = f'Comment  :   Chemshell/CASTEP Interface\n\n'

        if input == self.input[0]:              
            strbuff += f"{'task':20}:{self.castep_task:<20}\n"
            strbuff += f"{'opt_strategy':20}:{str(self.opt_strategy):<20}\n"
            xc = ['LDA', 'LDA-C', 'LDA-X','HF','SHF','EXX', 'PW91',
                 'PBE','PBE0','PBE_X', 'PBE_C', 'RPBE', 'WC', 'PBESOL','PBESOL_X', 'BLYP',
                  'B3LYP','NONE','HF-LDA', 'SX', 'SX-LDA', 'PBE0', 'HSE03', 'HSE06', 'RSCAN',
                  'B88_X','LYP_C']
            if self.functional.upper().strip() not in xc:
                print(f'\n >>> ERROR: \"{self.functional.upper()}\" is not a valid CASTEP exchange-correlation functional')
                print(f'\n Available functionals are:\n {xc}\n')
                exit(112)
            strbuff += f'{"xc_functional":20}:{self.functional.strip():<20}\n'
            strbuff += f'{"max_scf_cycles":20}:{self.maxiter:<20}\n'
            strbuff += f'{"finite_basis_corr":20}:{self.finite_basis_corr:<20}\n'
            strbuff += f'write_formatted_density:{self.write_formatted_density:<10}\n'
            strbuff += f'{"energy_unit":20}:{self.energy_unit:<20}\n'
            strbuff += f'{"force_unit":20}:{self.force_unit:<20}\n'
            strbuff += f'{"length_unit":20}:{self.length_unit:<20}\n'
            strbuff += f'{"popn_write":20}:{self.popn_write:<20}\n'
            strbuff += f'{"write_checkpoint":20}:{self.write_checkpoint:<20}\n'
            strbuff += f'{"elec_dump_file":20}:{self.elec_dump_file:<20}\n'
            strbuff += f'{"continuation":20}:{self.continuation:<20}\n'
            strbuff += f'{"grid_scale":20}:{self.grid_scale:<20}\n'

           # Optional arguments.Note: if not given a value in the input, castep uses its defaults. 
           # So the defaults are not specified in the chemshell interface(given null values)
          
            if self.basis_de_dloge:
               strbuff += f"{'basis_de_dloge':20}:{self.basis_de_dloge:<20}\n"
            if self.basis_precision:
               strbuff += f"{'basis_precision':20}:{self.basis_precision:<20}\n"
            if self.electronic_minimizer:
               strbuff += f'{"electronic_minimizer":20}:{self.electronic_minimizer:<20}\n'
            if self.fix_occupancy:
               strbuff += f"{'fix_occupancy':20}:{str(self.fix_occupancy):<20}\n"
            if self.geom_method:
                strbuff += f'{"geom_method":20}:{self.geom_method:<20}\n'
            if self.elec_energy_tol:
                strbuff += f'{"elec_energy_tol":20}:{self.elec_energy_tol:<25}\n'
            if self.elec_convergence_win:
                strbuff += f'{"elec_convergence_win":20}:{self.elec_convergence_win:<20}\n'
            if  self.cut_off_energy:
                strbuff += f'{"cut_off_energy":20}:{self.cut_off_energy:<25}\n'
            if self.finite_basis_npoints:
                 strbuff += f'{"finite_basis_npoints":20}:{self.finite_basis_npoints:<20}\n'
            if self.finite_basis_spacing:
                 strbuff += f'{"finite_basis_spacing":20}:{self.finite_basis_spacing:<20}\n'
            if self.iprint:
                 strbuff += f'{"iprint":20}:{self.iprint:<20}\n'
            if  self.metals_method:
                 strbuff += f'{"metals_method":20}:{self.metals_method:<20}\n'
            if  self.num_backup_iter:
                 strbuff += f'{"num_backup_iter":20}:{self.num_backup_iter:<20}\n'
            if  self.num_dump_cycles:
                 strbuff += f'{"num_dump_cycles":20}:{self.num_dump_cycles:<20}\n'
            if  self.nbands:
                 strbuff += f'{"nbands":10}:{self.nbands:<10}\n'
            if  self.mixing_scheme:
                 strbuff += f'{"mixing_scheme":20}:{self.mixing_scheme:<10}\n'
            if  self.rand_seed:
                 strbuff += f'{"rand_seed":20}:{self.rand_seed:<20}\n'
            if  self.smearing_width:
                 strbuff += f'{"smearing_width":20}:{self.smearing_width:<20}\n'
            #self.input_flag=False  
            return strbuff 

        elif input == self.input[1]:
            from ...parallel import parallel

            strbuff = ''
            elements = array(self.frag.znums)
            coordinates = array(cell.vectors)
            buff=''
            #Raj: lattice vector components in castep is a row, in chemshell a column
            for i in range(3):
                for j in range(3):
                   buff += f'{coordinates[j][i]:<0.15f} '
                buff += f'\n'
            strbuff += f'%BLOCK  lattice_cart\n'
            strbuff += f'{self.length_unit}\n'
            strbuff += buff
            strbuff += f'%ENDBLOCK  lattice_cart\n\n'

            strbuff += f'%BLOCK    positions_abs\n'
            strbuff += f'{self.length_unit}\n'
            #args of coords2str-indentation, symbols/znums, coords,charges specify precision,
            # Precision of coordinates here affects the results
            strbuff += frag.coords2str(indent=15*' ',
                   atoms='symbols',unit="a.u.", coords=15, charges=None)
            strbuff += f'%ENDBLOCK    positions_abs\n\n'

            buff=frag.coords2str(indent=3*' ',
                atoms='symbols',unit="a.u.", coords=15, charges=None)
            print('new coords\n', buff)

            #if peseudopotential is provided, give filenames as: atomicsymbol_filename. eg. H_pseudo.usp
            if self.pseudopotential is not None:

                #from os.path import exists

                strbuff +=  f'%BLOCK   species_pot\n' 
                pseudofiles = self.pseudopotential.split(',')

                buff=''

                # Only one keyword, use the pseudo library
                # valid:QC5 HARD MS C7 C8 C9 NCP9 C17 NCP17 C18 NCP18 C19 NCP19 SOC19 NCP

                if(len(pseudofiles) == 1 and '_' not in pseudofiles[0]):
                    for f in pseudofiles:
                        buff = f"{f.strip().upper()}\n"

                # Atom specific: use the form element_name, e.g, H_pseudo.usp, Si_ttt
                #else:
                else:
                    for f in pseudofiles:
                        buff += f"{f.split('_')[0].strip()}  {f.strip()}\n"

                strbuff += buff
                strbuff +=  f'%ENDBLOCK   species_pot\n' 


            strbuff +=  f"kpoint_mp_grid  {' '.join([ str(i) for i in self.kpoints ])}\n"

            if self.symmetry_generate:
                strbuff+= f"symmetry_generate \n"

            return strbuff

    def runLib(self, **kwargs):

        '''Run a library call of CASTEP (called by run.runLib())'''

        from ctypes     import byref 
        from ...dl_py2f import py2f

        print('\n >>> Directly-linked CASTEP library will be used...')

        #Rajany: This changes the opened file's name in the second run of castepgeomopt
        # This can't be avoided since we are not calling any intialisation
        # routines again. If not called, cannot backup previous result
        utils.fileutils.archiveFile(self.output)

        #Rajany: Before running castep, the coordinates to be sorted in Z order
        #..in optimisation, coords are directly passed to castep lib.
        # this changes coords in place, but first time this is not done,
        # instead input file is read. cellfile was written before this step
        
        self.frag.coords = self.sort_in_atom_order(self.frag.coords)
        libcastep = utils.modutils.getLinkedLib(self.__class__.__name__, chemsh.CHEMSH_LIB_DIR)

        print(f'\n >>> castep library {libcastep}')

        #Raj: py2f returns address /pointer of the objects passed
        fragRef    = py2f(self.frag   , byref=True)
        optionsRef = py2f(self.options, byref=True)
        resultsRef = py2f(self._result, byref=True)
        cellRef    = py2f(self.frag.cell, byref=True)

        if libcastep.interface_castep(fragRef, optionsRef, resultsRef,cellRef):
            raise errors.ChemShModuleError('CASTEP')

        #Raj: if library is run, use values of energy and forces returned
        self._parseoutfile = False
        return
#################################### END ######################################
