import PETSc.package

class Configure(PETSc.package.NewPackage):
  def __init__(self, framework):
    PETSc.package.NewPackage.__init__(self, framework)
    self.download     = ['http://icl.cs.utk.edu/projectsfiles/magma/downloads/magma-1.4.0.tar.gz']
    self.functions    = ['magma_zgetrf']
    self.includes     = ['magma.h']
    self.liblist      = [['libmagma.a'],['libmagma.a','libmagmablas.a']]
    self.double       = 0
    self.complex      = 1

    self.worksonWindows   = 0
    self.downloadonWindows= 0
    return

  def setupDependencies(self, framework):
    PETSc.package.NewPackage.setupDependencies(self, framework)
    PETSc.package.NewPackage.setupDependencies(self, framework)
    self.blasLapack = self.framework.require('config.packages.BlasLapack',self)
    self.cuda       = self.framework.require('PETSc.packages.cuda',self)
    self.deps       = [self.blasLapack, self.cuda]
    return

  def Install(self):
    import os
    import re
    import sys

    # set blas name mangling
    if self.blasLapack.mangling == 'underscore':
      self.mangling   = '-DADD_'
    elif self.blasLapack.mangling == 'caps':
      self.mangling   = '-DUPCASE'
    else:
      self.mangling   = '-DNOCHANGE'

    
    g = open(os.path.join(self.packageDir,'make.inc'),'w')

    g.write('LIB          = -lstdc++ -lcudart -lcublas '+self.libraries.toString(self.blasLapack.dlib)+'\n')
    g.write('ARCH         = '+self.setCompilers.AR+'\n')
    g.write('ARCHFLAGS    = '+self.setCompilers.AR_FLAGS+'\n')
    g.write('RANLIB       = '+self.setCompilers.RANLIB+'\n')
    self.setCompilers.pushLanguage('C')
    g.write('CC           = '+self.setCompilers.getCompiler()+'\n')
    g.write('OPTS         = '+self.setCompilers.getCompilerFlags()+' '+self.mangling+' -DMAGMA_SETAFFINITY\n')
    self.setCompilers.popLanguage()
    self.setCompilers.pushLanguage('FC')
    g.write('FORT         = '+self.setCompilers.getCompiler()+'\n')
    g.write('F77OPTS      = '+self.setCompilers.getCompilerFlags()+' '+self.mangling+' \n')
    g.write('FOPTS        = '+self.setCompilers.getCompilerFlags()+' '+self.mangling+' -x f95-cpp-input\n')
    self.setCompilers.popLanguage()
    self.setCompilers.pushLanguage('CUDA')
    g.write('NVCC         = '+self.setCompilers.getCompiler()+'\n')
    
    
    #  Set the GPU_TARGET for MAGMA depending on the GPU detected by PETSc
    self.CUDAFLAGS=self.setCompilers.getCompilerFlags()
    if '-arch=sm_13' in self.CUDAFLAGS:
      g.write('GPU_TARGET   = Tesla\n')
      self.CUDAFLAGS=self.CUDAFLAGS.replace('-arch=sm_13','')
    elif '-arch=sm_20' in self.CUDAFLAGS:
      g.write('GPU_TARGET   = Fermi\n')
      self.CUDAFLAGS=self.CUDAFLAGS.replace('-arch=sm_20','')
    elif  '-arch=sm_30' in self.CUDAFLAGS:
      g.write('GPU_TARGET   = Kepler\n')
      self.CUDAFLAGS=self.CUDAFLAGS.replace('-arch=sm_30','')
    elif  '-arch=sm_35' in self.CUDAFLAGS:
      g.write('GPU_TARGET   = Kepler\n')
      self.CUDAFLAGS=self.CUDAFLAGS.replace('-arch=sm_35','')
    else:
      raise RuntimeError('MAGMA error: GPU_TARGET must be one of Tesla, Fermi, or Kepler.')
    
    g.write('NVOPTS       = '+self.CUDAFLAGS+' '+self.mangling+' -Xcompiler -fno-strict-aliasing\n')
    g.write('LDOPTS       = -fopenmp\n')
    self.setCompilers.popLanguage()

    g.write('CUDADIR      = '+self.cuda.directory+'\n')
    g.write('INC          = -I'+self.cuda.directory+'/include\n')
    g.write('LIBDIR       = -L'+self.cuda.directory+'/lib\n')

    g.close()
    if self.installNeeded('make.inc'):

      try:
        self.logPrintBox('Compiling MAGMA; this may take several minutes')
        output,err,ret = PETSc.package.NewPackage.executeShellCommand('cd '+self.packageDir+' && make clean && make -j8', timeout=2500, log = self.framework.log)
        libDir     = os.path.join(self.installDir, self.libdir)
        includeDir = os.path.join(self.installDir, self.includedir)

        g = open(os.path.join(self.packageDir+'/include','magmaconfig.h'),'w')
        g.write('#define HAVE_CUBLAS 1\n')
        g.close()
        
        g = open(os.path.join(self.packageDir+'/include','magma.h'),'r+')
        old=g.read()
        g.seek(0)
        g.write('#include "magmaconfig.h"\n\n'+old)
        g.close()
 
        output,err,ret = PETSc.package.NewPackage.executeShellCommand('cd '+self.packageDir+' && mv -f lib/*.* '+libDir+'/. && cp -f include/*.* '+includeDir+'/.', timeout=2500, log = self.framework.log)
      except RuntimeError, e:
        raise RuntimeError('Error running make on MAGMA: '+str(e))
      self.postInstall(output+err,'make.inc')
    return self.installDir