import os
from ase.calculators.calculator import FileIOCalculator, CalculationFailed
from xespresso.xio import read_espresso_asei, write_espresso_asei
import copy
import logging
logger = logging.getLogger(__name__)


class PostCalculation:
    
    package = 'dos'
    package_parameters = {}

    def __init__(self, parent_directory, prefix, queue = False, parallel = '', **kwargs) -> None:
        self.parent_directory = parent_directory
        self.prefix = prefix
        self.queue = queue
        self.parallel = parallel
        self.parameters = kwargs
        self.directory = os.path.join(self.parent_directory, '%s/'%self.package)
        self.set_label(self.directory, self.prefix)
        self.parameters['prefix'] = self.prefix
        self.parameters['outdir'] = '../'
        self.state_info = None

    def set_label(self, label, prefix):
        '''
        set directory and prefix from label
        '''
        self.directory = label
        if not prefix:
            self.prefix = os.path.split(label)[1]
        else:
            self.prefix = prefix
        if not os.path.exists(self.directory):
            os.makedirs(self.directory)
        self.label = os.path.join(self.directory, self.prefix)
        self.asei = os.path.join(self.directory, '%s.asei'%self.prefix)
        self.asei_temp = os.path.join(self.directory, '.%s.asei_temp'%self.prefix)
        self.post_asei = os.path.join(self.directory, '%s.post_asei'%self.prefix)
        self.save_directory = os.path.join(self.directory, '%s.save'%self.prefix)
        logger.debug('Directory: %s'%(self.directory))
        logger.debug('Prefix: %s'%(self.prefix))

    def run(self):
        '''
        todo: 
        '''
        from xespresso.scheduler import set_queue
        
        print('{0:=^60}'.format(self.package))
        if self.check_state() == 0:
            return
        self.write_input()
        set_queue(self, package = self.package, parallel = self.parallel, queue = self.queue)
        self.post_calculate()
        self.post_read_results()

    def check_state(self):
        from xespresso.utils import get_hash
        self.state_info = None
        for wfc in ['wfc1', 'wfcdw1']:
            wfcFile = os.path.join(self.parent_directory, '%s.save/%s.dat'%(self.prefix, wfc))
            if os.path.isfile(wfcFile):
                self.state_info = get_hash(wfcFile)
        output, meg = self.read_convergence_post(self.package)
        if output:
            logger.debug('Previous calculation done.')
            if os.path.isfile(self.post_asei):
                system_changes =  self.check_state_post(self.post_asei, self.package)
                if not system_changes:
                    logger.debug('File and Parameters did not change. Use previous results!')
                    return 0
        return 1

    def check_state_post(self, asei, package):
        old_state_info, old_parameters = read_espresso_asei(asei, package)
        if not self.state_info == old_state_info:
            logger.debug('File in save changed')
            return True
        elif not self.parameters == old_parameters:
            logger.debug('Parameters changed')
            return True
        else:
            return False

    def write_input(self):
        self.write_package_input()
        write_espresso_asei(self.post_asei, self.state_info, self.parameters)

    def write_package_input(self):
        filename = os.path.join(self.directory, '%s.%si' %(self.prefix, self.package))
        with open(filename, 'w') as f:
            for section, parameters in self.package_parameters.items():
                # logger.debug(section)
                if section != 'LINE':
                    f.write('&%s\n'%section)
                    for key, value in self.parameters.items():
                        if key in parameters:
                            if isinstance(value, dict):
                                for subkey, subvalue in value.items():
                                    if isinstance(subvalue, str):
                                        f.write('  %s(%s) = "%s", \n' %(key, subkey, subvalue))
                                    else:
                                        f.write('  %s(%s) = %s, \n' %(key, subkey, subvalue))
                            else:
                                if isinstance(value, str):
                                    f.write('  {0:10s} =  "{1}" \n'.format(key, value))
                                else:
                                    f.write('  {0:10s} =  {1} \n'.format(key, value))
                    f.write('/ \n')
                else:
                    for key, value in self.parameters.items():
                        if key in parameters:
                            f.write('  %s \n' %(value))


    def post_calculate(self):
        import subprocess
        
        command = self.command
        print('Running %s'%self.package)
        try:
            proc = subprocess.Popen(command, shell=True, cwd=self.directory)
        except OSError as err:
            msg = 'Failed to execute "{}"'.format(command)
            raise EnvironmentError(msg) from err

        errorcode = proc.wait()

        if errorcode:
            path = os.path.abspath(self.directory)
            msg = ('Command "{}" failed in '
                   '{} with error code {}'.format(command,
                                                  path, errorcode))
            raise CalculationFailed(msg)
        print('Done: %s' %self.package)
    
    def post_read_results(self):
        '''
        '''
        pass

    def read_convergence_post(self, package = 'pw'):
        '''
        Read the status of the calculation.
        {
        '0': 'Done',
        }
        '''
        
        output = self.label + '.%so'%package
        if not os.path.exists(output):
            # print('%s not exists'%output)
            return False, 'No pwo output file'
        with open(output, 'r') as f:
            lines = f.readlines()
            if len(lines) == 0: return False, 'pwo file has nothing'
            nlines = len(lines)
            n = min([100, nlines])
            for line in lines[-n:-1]:
                if line.rfind('JOB DONE.') > -1:
                    logger.debug('JOB DONE.')
                    return True, line
        return False, line
    