from GAML.functions import file_gen_new, file_size_check
from GAML.function_gen_range import function_gen_range

class Charge_gen_range(object):
    """Generate the charge ranges surrounding the mode number"""   

    def __init__(self,*args,**kwargs):
        self.log = {'nice':True,}

        bo = False
        if 'atomnm' in kwargs and kwargs['atomnm'] is not None:
            try:
                self.atomnm = int(kwargs['atomnm'])
                if self.atomnm <= 0:
                    raise ValueError
            except ValueError:
                bo = True
                self.log['info'] = 'Error: the parameter atomnm has to be a positive number'
        else:
            bo = True
            self.log['info'] = 'Error: the parameter atomnm is missing'
        if bo:
            self.log['nice'] = False
            return


        if 'charge_path' in kwargs and kwargs['charge_path'] is not None:
            """ATTENTION! Here may have bugs, if the input charge_file_path is not
               generated by the Gaussian. The follow shows the file format:

               nm  atomtype    charge
               1     C           0.1
               2     H           0.2
               3     O          -0.2

               Only three columns are used. The first column is the input index,
               the second column is the atom-type, the last one is the value"""

            self.charge_path = kwargs['charge_path']
            mcp = file_size_check(self.charge_path,fsize=100)
            if not mcp['nice']:
                self.log['nice'] = False
                self.log['info'] = mcp['info']
                return

            datalist = []
            with open(self.charge_path,mode='rt') as f:
                while True:
                    line = f.readline()
                    if len(line) == 0:
                        break
                    else:
                        lst = line[:line.find('#')].split()
                        if len(lst) == 3 and int(lst[0]) <= self.atomnm:
                            datalist.append(float(lst[2]))

            if len(datalist) % self.atomnm != 0:
                self.log['nice'] = False
                self.log['info'] = 'Error: wrong input_data_file or atom_number\n' + \
                                       'Error: input charge_file_path: ' + self.charge_path + '\n' + \
                                       'Error input atom_number: ' + str(self.atomnm)
                return

            # to avoid using an extra numpy module
            # works similar like;
            #     numpy.array(datalist).reshape(len(datalist)//self.atomnm,self.atomnm)
            self.profile = []
            i = 0
            while i < len(datalist):
                ls = []
                j = 0
                while j < self.atomnm:
                    ls.append(datalist[i])
                    i += 1
                    j += 1
                self.profile.append(ls)
        else:
            self.log['nice'] = False
            self.log['info'] = 'Error: the parameter charge_path is missing'
            return
 
        
        if 'percent' in kwargs and kwargs['percent'] is not None:
            try:
                self.percent = float(kwargs['percent'])
                if self.percent > 1.0 or self.percent <= 0:
                    raise ValueError
            except ValueError:
                self.log['nice'] = False
                self.log['info'] = 'Error: the parameter percent has to be a positive number\n' + \
                                   '     : and its value must fall within in (0,1]'
                return
        else:
            self.percent = 0.8
            
            
        if 'stepsize' in kwargs and kwargs['stepsize'] is not None:
            try:
                self.stepsize = float(kwargs['stepsize'])
                if self.stepsize <= 0:
                    raise ValueError
            except ValueError:
                self.log['nice'] = False
                self.log['info'] = 'Error: the parameter stepsize has to be a positive number'
                return            
        else:
            self.stepsize = 0.01

            
        if 'nmround' in kwargs and kwargs['nmround'] is not None:
            try:
                self.nmround = int(kwargs['nmround'])
                if self.nmround < 0:
                    raise ValueError
            except ValueError:
                self.log['nice'] = False
                self.log['info'] = 'Error: the parameter nmround has to be a positive integer'
                return
        else:
            self.nmround = 3

            
        if 'fname' in kwargs and kwargs['fname'] is not None:
            self.fname = str(kwargs['fname'])
        else:
            self.fname = 'ChargeGenRange'



    def charge_range(self):                       
        atom = 0
        charge_range = []
        while atom < self.atomnm:
            atomcharge = []
            for mol in self.profile:
                atomcharge.append(mol[atom])

            lmin = charmin = min(atomcharge)
            lmax = max(atomcharge)
            
            rlist = []
            while lmin < lmax:          
                lcount = 0
                # Warning! to improve the efficiency, here can be debugged
                for i in atomcharge:
                    if i < lmin + self.stepsize and i >= lmin:
                        lcount += 1
                rlist.append(lcount)
                lmin += self.stepsize

            ndxmin,ndxmax = function_gen_range(rlist,percent=self.percent)

            ltmp = []
            ltmp.append(round((ndxmin * self.stepsize + charmin),self.nmround))
            ltmp.append(round(((ndxmax-1) * self.stepsize + charmin),self.nmround))
            charge_range.append(ltmp)

            atom += 1

        return charge_range


    def file_print(self):
        
        fname = file_gen_new(self.fname,fextend='txt',foriginal=True)

        charge_range = self.charge_range()
        
        with open(fname,mode='wt') as f:
            f.write('# This is the generated charge_range based on the input charge_file \n')
            f.write('# The charge file used is:\n')
            f.write('#     {:s}\n'.format(self.charge_path))
            f.write('# Each line\'s charge_range is corresponded to each atom \n\n')
            j = 1
            for i in charge_range:
                f.write('ATOM {:>6d}    {:>8.4f}    {:>8.4f}\n'.format(j,i[0],i[1]))
                j += 1

