import numpy as np
from socket import gethostname
from subprocess import call
from IssmConfig import IssmConfig
from issmdir import issmdir
from pairoptions import pairoptions
from issmssh import issmssh
from issmscpin import issmscpin
from issmscpout import issmscpout
from MatlabFuncs import ispc
try:
    from generic_settings import generic_settings
except ImportError:
    print('Warning generic_settings.py not found, default will be used')


class generic(object):
    """
    GENERIC cluster class definition

       Usage:
          cluster = generic('name', 'astrid', 'np', 3)
          cluster = generic('name', gethostname(), 'np', 3, 'login', 'username')
    """

    def __init__(self, *args):  # {{{
        self.name = ''
        self.login = ''
        self.np = 1
        self.port = 0
        self.interactive = 1
        self.codepath = IssmConfig('ISSM_PREFIX')[0] + '/bin'
        self.executionpath = issmdir() + '/execution'
        self.valgrind = issmdir() + '/externalpackages/valgrind/install/bin/valgrind'
        self.valgrindlib = issmdir() + '/externalpackages/valgrind/install/lib/libmpidebug.so'
        self.valgrindsup = [issmdir() + '/externalpackages/valgrind/issm.supp']  # add any .supp in list form as needed
        self.verbose = 1
        self.shell = '/bin/sh'

        #use provided options to change fields
        options = pairoptions(*args)

        #get name
        self.name = gethostname()

        #initialize cluster using user settings if provided

        try:
            self = generic_settings(self)
        except NameError:
            print("generic_settings.py not found, using default settings")
        # else:
        #     raise

        #OK get other fields
        self = options.AssignObjectFields(self)
    # }}}

    def __repr__(self):  # {{{
        #  display the object
        s = "class '%s' object '%s' = \n" % (type(self), 'self')
        s += "    name: %s\n" % self.name
        s += "    login: %s\n" % self.login
        s += "    np: %i\n" % self.np
        s += "    port: %i\n" % self.port
        s += "    codepath: %s\n" % self.codepath
        s += "    executionpath: %s\n" % self.executionpath
        s += "    valgrind: %s\n" % self.valgrind
        s += "    valgrindlib: %s\n" % self.valgrindlib
        s += "    valgrindsup: %s\n" % self.valgrindsup
        s += "    verbose: %s\n" % self.verbose
        s += "    shell: %s\n" % self.shell
        return s
    # }}}

    def checkconsistency(self, md, solution, analyses):  # {{{
        if self.np < 1:
            md.checkmessage('number of processors should be at least 1')
        if np.isnan(self.np):
            md.checkmessage('number of processors should not be NaN!')

        return md
    # }}}

    def BuildQueueScript(self, dirname, modelname, solution, io_gather, isvalgrind, isgprof, isdakota, isoceancoupling):  # {{{
        # Which executable are we calling?
        executable = 'issm.exe'  # default

        if isdakota:
            version = IssmConfig('_DAKOTA_VERSION_')
            version = float(version[0])
            if version >= 6:
                executable = 'issm_dakota.exe'
        if isoceancoupling:
            executable = 'issm_ocean.exe'

        # Write queueing script
        if not ispc():
            fid = open(modelname + '.queue', 'w')
            fid.write('#!/bin/sh\n')
            if not isvalgrind:
                if self.interactive:
                    if IssmConfig('_HAVE_MPI_')[0]:
                        fid.write('mpiexec -np {} {}/{} {} {}/{} {}'.format(self.np, self.codepath, executable, solution, self.executionpath, dirname, modelname))
                    else:
                        fid.write('{}/{} {} {}/{} {} '.format(self.codepath, executable, solution, self.executionpath, dirname, modelname))
                else:
                    if IssmConfig('_HAVE_MPI_')[0]:
                        fid.write('mpiexec -np {} {}/{} {} {}/{} {} 2>{}.errlog>{}.outlog'.
                                  format(self.np, self.codepath, executable, solution, self.executionpath, dirname, modelname, modelname, modelname))
                    else:
                        fid.write('{}/{} {} {}/{} {} 2>{}.errlog>{}.outlog '.
                                  format(self.codepath, executable, solution, self.executionpath, dirname, modelname, modelname, modelname))
            elif isgprof:
                fid.write('\n gprof {}/{} gmon.out > {}.performance'.format(self.codepath, executable, modelname))
            else:
                #Add --gen -suppressions = all to get suppression lines
                #fid.write('LD_PRELOAD={} \\\n'.format(self.valgrindlib)) it could be deleted
                supstring = ''
                for supfile in self.valgrindsup:
                    supstring += ' --suppressions=' + supfile

                if IssmConfig('_HAVE_MPI_')[0]:
                    fid.write('mpiexec -np {} {} --leak-check=full {} {}/{} {} {}/{} {} 2>{}.errlog>{}.outlog '.
                              format(self.np, self.valgrind, supstring, self.codepath, executable, solution, self.executionpath, dirname, modelname, modelname, modelname))
                else:
                    fid.write('{} --leak-check=full {} {}/{} {} {}/{} {} 2>{}.errlog>{}.outlog '.
                              format(self.valgrind, supstring, self.codepath, executable, solution, self.executionpath, dirname, modelname, modelname, modelname))

            if not io_gather:  #concatenate the output files:
                fid.write('\ncat {}.outbin .*>{}.outbin'.format(modelname, modelname))
            fid.close()

        else:  # Windows
            fid = open(modelname + '.bat', 'w')
            fid.write('@echo off\n')
            if self.interactive:
                fid.write('"{}/{}" {} "{}/{}" {} '.format(self.codepath, executable, solution, self.executionpath, dirname, modelname))
            else:
                fid.write('"{}/{}" {} "{}/{}" {} 2>{}.errlog>{}.outlog'.
                          format(self.codepath, executable, solution, self.executionpath, dirname, modelname, modelname, modelname))
            fid.close()

            #in interactive mode, create a run file, and errlog and outlog file
        if self.interactive:
            fid = open(modelname + '.errlog', 'w')
            fid.close()
            fid = open(modelname + '.outlog', 'w')
            fid.close()
    # }}}

    def BuildKrigingQueueScript(self, modelname, solution, io_gather, isvalgrind, isgprof):  # {{{
        #write queuing script
        if not ispc():
            fid = open(modelname + '.queue', 'w')
            fid.write('#!/bin/sh\n')
            if not isvalgrind:
                if self.interactive:
                    fid.write('mpiexec -np {} {}/kriging.exe {}/{} {} '.format(self.np, self.codepath, self.executionpath, modelname, modelname))
                else:
                    fid.write('mpiexec -np {} {}/kriging.exe {}/{} {} 2>{}.errlog>{}.outlog '.format(self.np, self.codepath, self.executionpath, modelname, modelname, modelname, modelname))
            elif isgprof:
                fid.write('\n gprof {}/kriging.exe gmon.out>{}.performance'.format(self.codepath, modelname))
            else:
                #Add - -    gen - suppressions = all to get suppression lines
                #fid.write('LD_PRELOAD={} \\\n'.format(self.valgrindlib))
                fid.write('mpiexec -np {} {} --leak -check=full --suppressions={} {}/kriging.exe {}/{} {} 2 > {}.errlog > {}.outlog ' .format(self.np, self.valgrind, self.valgrindsup, self.codepath, self.executionpath, modelname, modelname, modelname, modelname))
            if not io_gather:    #concatenate the output files:
                fid.write('\ncat {}.outbin. *>{}.outbin'.format(modelname, modelname))
            fid.close()

        else:    # Windows
            fid = open(modelname + '.bat', 'w')
            fid.write('@echo off\n')
            if self.interactive:
                fid.write('"{}/issm.exe" {} "{}/{}" {} '.format(self.codepath, solution, self.executionpath, modelname, modelname))
            else:
                fid.write('"{}/issm.exe" {} "{}/{}" {} 2>{}.errlog>{}.outlog'.format
                          (self.codepath, solution, self.executionpath, modelname, modelname, modelname, modelname))
            fid.close()

        #in interactive mode, create a run file, and errlog and outlog file
        if self.interactive:
            fid = open(modelname + '.errlog', 'w')
            fid.close()
            fid = open(modelname + '.outlog', 'w')
            fid.close()
    # }}}

    def UploadQueueJob(self, modelname, dirname, filelist):  # {{{
        #compress the files into one zip.
        compressstring = 'tar -zcf {}.tar.gz '.format(dirname)
        for file in filelist:
            compressstring += ' {}'.format(file)
        if self.interactive:
            compressstring += ' {}.errlog {}.outlog '.format(modelname, modelname)
        call(compressstring, shell=True)

        print('uploading input file and queueing script')
        issmscpout(self.name, self.executionpath, self.login, self.port, [dirname + '.tar.gz'])

    # }}}

    def LaunchQueueJob(self, modelname, dirname, filelist, restart, batch):  # {{{
        print('launching solution sequence on remote cluster')
        if restart:
            launchcommand = 'cd {} && cd {} chmod 755 {}.queue && ./{}.queue'.format(self.executionpath, dirname, modelname, modelname)
        else:
            if batch:
                launchcommand = 'cd {} && rm -rf ./{} && mkdir {} && cd {} && mv ../{}.tar.gz ./&& tar -zxf {}.tar.gz'.format(self.executionpath, dirname, dirname, dirname, dirname, dirname)
            else:
                launchcommand = 'cd {} && rm -rf ./{} && mkdir {} && cd {} && mv ../{}.tar.gz ./&& tar -zxf {}.tar.gz  && chmod 755 {}.queue && ./{}.queue'.format(self.executionpath, dirname, dirname, dirname, dirname, dirname, modelname, modelname)
        issmssh(self.name, self.login, self.port, launchcommand)
    # }}}

    def Download(self, dirname, filelist):  # {{{
        if ispc():
            #do nothing
            return
        #copy files from cluster to current directory
        directory = '{}/{}/'.format(self.executionpath, dirname)
        issmscpin(self.name, self.login, self.port, directory, filelist)
    # }}}
