#!/usr/bin/env python

######################################################################
#  Distributed under the terms of the GNU General Public License (GPL)
#                  http://www.gnu.org/licenses/
######################################################################

import os, sys, shutil
from configuration import conf, cp, try_run, edit_in_place


######################################################################
### Skip building ATLAS on specific systems
######################################################################

# On Cygwin we simply require that the system-wide lapack is installed.
# This includes BLAS and is enough to build the rest of Sage.
if conf['CYGWIN?']:
    if not os.path.exists('/usr/lib/libblas.a'):
        print "*"*75
        print "On Cygwin you must install the standard LAPACK Cygwin package"
        print "via the Cygwin setup.exe program in the 'Math' category."
        print "*"*75 
        sys.exit(1)
    sys.exit(0)


if conf['Darwin?']:
    sys.exit(0)


######################################################################
### Use SAGE_ATLAS_LIB='directory' if provided
######################################################################

if os.environ.has_key('SAGE_ATLAS_LIB'):
    ATLAS_LIB=os.environ['SAGE_ATLAS_LIB']
    if not os.path.isdir(ATLAS_LIB):
        print 'No such directory: '+ATLAS_LIB
        sys.exit(2)

    # Check for 4 static libraries. 
    has_atlas = os.path.exists(ATLAS_LIB+'/libatlas.a')
    has_lapack = os.path.exists(ATLAS_LIB+'/liblapack.a')
    has_cblas = os.path.exists(ATLAS_LIB+'/libcblas.a')
    has_f77blas = os.path.exists(ATLAS_LIB+'/libf77blas.a')
    has_atlas_headers = os.path.exists(ATLAS_LIB+'/../include/atlas')
    
    symbol_table = try_run('readelf -s ' +ATLAS_LIB+'/libf77blas.so')
    if not symbol_table is None:
        s_gfortran = 'gfortran' in symbol_table
        s_g95 = 'g95' in symbol_table
        if s_gfortran and not config['fortran'] != 'gfortran':
            print "Symbols in lib77blas indicate it was build with gfortran \n"
            print "However SAGE is using a different fortran compiler \n"
            print "If you wish to use this blas library, make sure SAGE_FORTRAN points \n"
            print "to a fortran compiler compatible with this library. \n"
            sys.exit(2)
        if s_g95 and config['fortran'] != 'g95':
            print "Symbols in lib77blas indicate it was build with g95 \n"
            print "However SAGE is using a different fortran compiler \n"
            print "If you wish to use this blas library, make sure SAGE_FORTRAN points \n"
            print "to a fortran compiler compatible with this library. \n"
            sys.exit(2)

    if has_atlas and has_lapack and has_cblas and  has_f77blas and has_atlas_headers:
        os.system('ln -sf ' + ATLAS_LIB + '/libatlas.a '  + SAGE_LOCAL_LIB+'/libatlas.a')
        os.system('ln -sf ' + ATLAS_LIB + '/libcblas.a '  + SAGE_LOCAL_LIB+ '/libcblas.a')
        os.system('ln -sf ' + ATLAS_LIB + '/liblapack.a '  + SAGE_LOCAL_LIB+'/liblapack.a')
        os.system('ln -sf ' + ATLAS_LIB + '/libf77blas.a '  + SAGE_LOCAL_LIB+'/libf77blas.a')
        # We might as well include the shared libraries libatlas.so and libcblas.so, as these
        # build reliably on both Solaris and Linux, and might reduce the size of executable if 
        # we link to them. 
        os.system('ln -sf ' + ATLAS_LIB + '/libatlas.so '  + SAGE_LOCAL_LIB+'/libatlas.so')
        os.system('ln -sf ' + ATLAS_LIB + '/libcblas.so '  + SAGE_LOCAL_LIB+ '/libcblas.so')
        # Building of liblapack.so and libf77blas.so is not reliable, but we will link to them
        # People will have to make sure if they do build them, that they work. But these changes
        # only impact SAGE_ATLAS_LIB, which currently is broken on Solaris and often broken on Linux. 
        os.system('ln -sf ' + ATLAS_LIB + '/liblapack.so '  + SAGE_LOCAL_LIB+'/liblapack.so')
        os.system('ln -sf ' + ATLAS_LIB + '/libf77blas.so '  + SAGE_LOCAL_LIB+'/libf77blas.so')
        # Link to the ATLAS related header files. 
        os.system('ln -sf ' + ATLAS_LIB + '/../include/atlas '  + SAGE_LOCAL_INCLUDE+'/')
        os.system('ln -sf ' + ATLAS_LIB + '/../include/cblas.h '  + SAGE_LOCAL_INCLUDE+'/cblas.h')
        os.system('ln -sf ' + ATLAS_LIB + '/../include/clapack.h '  + SAGE_LOCAL_INCLUDE+'/clapack.h')
        
    else:
        print 'Unable to find one of liblapack.a, libcblas.a, libatlas.a or libf77blas.a'
        print 'in the directory',ATLAS_LIB
        print 'Set SAGE_ATLAS_LIB to the parent directory of the directory containing'
        print 'liblapack.a, libcblas.a, libatlas.a and libf77blas.a if you wish to use'
        print 'existing ATLAS libraries.'
        print 'So if for example those four libraries are in /usr/local/atlas/lib64'
        print 'then set SAGE_ATLAS_LIB to be /usr/local/atlas/lib64'
        print 'Unset SAGE_ATLAS_LIB to build ATLAS from source.'
        print 'Then type make.'
        sys.exit(2)


######################################################################
### Patch source
######################################################################

# Apply a TEMPORARY fix to allow ATLAS to build with
# gcc 4.4.0 on Solaris. Implemented 16th June 2009, by David Kirkby.
# One would expect to remove this within a couple of months,
# once the underlying issue in ATLAS is resolved. The patch
# forces GuessSmallNB() in src/tune/blas/gemm/mmsearch.c
# to return 28 as suggested by Clint Whaley.
# Changed on July 19th 2009 (see trac 6558) to be more selective and only
# apply the fix on sun4v machines which are based on the Sun T1, T2 and T2+
# processors (codenamed Niagra). This is because there are no known problems
# on other Sun architectures such as sun4u, or any x86 based Solaris system. 
if conf['Solaris?'] and conf['machine'] == 'sun4v':
    cp('patches/mmsearch-with-temp-Solaris-fix.c',
       'src/tune/blas/gemm/mmsearch.c')

# Patching SpewMakeInc.c for FreeBSD-specific build
if conf['FreeBSD?']:
    cp('patches/SpewMakeInc.c', 
       'src/CONFIG/src/SpewMakeInc.c')

# add PPC4 7447 CPU and better Itanium2 detection:
cp('patches/archinfo_linux.c',
   'src/CONFIG/src/backend/archinfo_linux.c')

# add Pentium D and E as well as Core2Duo and Dunnington CPUids
cp('patches/archinfo_x86.c',
   'src/CONFIG/src/backend/archinfo_x86.c')

# work around "-m64" cflag issue on Itanium
cp('patches/probe_comp.c', 
   'src/CONFIG/src/probe_comp.c')

# add dynamic libs make install targets to Make.top
cp('patches/Make.top',
   'src')

# add K7, Pentium M and non-AltiVec G4 profiles
cp('patches/*tgz', 
   'src/CONFIG/ARCHS')



######################################################################
### configure
######################################################################


# constants from src/CONFIG/include/atlconf.h
ATLAS_OSTYPE = ('UNKNOWN', 'Linux', 'SunOS', 'SunOS4', 'OSF1', 
                'IRIX', 'AIX', 'Win9x', 'WinNT', 'HPUX', 'FreeBSD', 'OSX')
ATLAS_MACHTYPE = ('UNKNOWN', 'POWER3', 'POWER4', 'POWER5', 'PPCG4', 'PPCG5',
                  'P5', 'P5MMX', 'PPRO', 'PII', 'PIII', 'PM', 'CoreSolo',
                  'CoreDuo', 'Core2Solo', 'Core2', 'Corei7', 'P4', 'P4E', 'Efficeon', 'K7',
                  'HAMMER', 'AMD64K10h', 'UNKNOWNx86', 'IA64Itan', 'IA64Itan2',
                  'USI', 'USII', 'USIII', 'USIV', 'UnknownUS', 'MIPSR1xK', 'MIPSICE9')
ATLAS_ISAEXT = ('', 'AltiVec', 'SSE3', 'SSE2', 'SSE1', '3DNow')



def configure(arch=None, isa_ext=None):
    """
    Configure for ``arch``.

    INPUT:

    - ``arch`` -- ``None`` or one of ``ATLAS_MACHTYPE``

    - ``isa_ext`` -- ``None`` or a sublist of ``ATLAS_ISAEXT``
    """
    if os.path.isdir('ATLAS-build'):
        shutil.rmtree('ATLAS-build')
    os.mkdir('ATLAS-build')
    os.chdir('ATLAS-build')

    LOCAL = os.environ['SAGE_LOCAL']
    cmd = '../src/configure'
    cmd += ' --prefix=' + LOCAL
    cmd += ' --with-netlib-lapack=' + LOCAL + '/lib/liblapack.a'
    # -Si cputhrchk 0: Ignore/heed CPU throttle probe
    cmd += ' -Si cputhrchk 0'
    # -Fa alg -fPIC: set flags so we can build dynamic libraries
    cmd += ' -Fa alg -fPIC'
    ## -t 0: disable threading
    #cmd += ' -t 0'
    cmd += ' -C if sage_fortran'
    # set bit width
    cmd += ' -b ' + conf['bits'][0:2]

    # set OS type
    system = None
    try:
       system = ATLAS_OSTYPE.index(conf['system'])
    except KeyError:
       if conf['Darwin?']: system = ATLAS_OSTYPE.index('OSX')
       if conf['CYGWIN?']: system = ATLAS_OSTYPE.index('WinNT')
    if not system is None:
       cmd += ' -O '+str(system)

    # set machine architecture
    if not arch is None:
       cmd += ' -A '+str(ATLAS_MACHTYPE.index(arch))

    # set cpu instruction set extensions
    if not isa_ext is None:
        isa_extension = sum(1 << ATLAS_ISAEXT.index(x) for x in isa_ext)
        cmd += ' -V '+str(isa_extension)

    print 'Running', cmd
    rc = os.system(cmd)
    os.chdir('..')
    return rc


def configure_opt_fast():
   if conf['Intel?']:
      print 'Fast configuration on Intel compatible CPUs.'
      arch = 'HAMMER'
      isa_ext = ('SSE3', 'SSE2', 'SSE1')
   elif conf['SPARC?']:
      print 'Fast configuration on SPARC.'
      arch = 'USIV'
      isa_ext = ()
   elif conf['PPC?']:
      print 'Fast configuration on PPC.'
      arch = 'POWER5'
      isa_ext = ('AltiVec', )
   else:
      raise NotImplementedError, 'I don\'t know a "fast" configuration for your cpu.'
   return (arch, isa_ext)


def configure_opt_base():
   if conf['Intel?']:
      print 'Base configuration on Intel compatible CPUs.'
      arch = 'P4'
      isa_ext = ('SSE2', 'SSE1')
   elif conf['SPARC?']:
      print 'Base configuration on SPARC.'
      arch = 'USIII'
      isa_ext = ()
   elif conf['PPC?']:
      print 'Base configuration on PPC.'
      arch = 'POWER4'
      isa_ext = ()
   else:
      raise NotImplementedError, 'I don\'t know a "base" configuration for your cpu.'
   return (arch, isa_ext)


# Figure out architecture (see ATLAS_MACHTYPE) and isa extensions (see
# ATLAS_ISAEXT) from environment variables:
arch = None
isa_ext = None
if os.environ.get('SAGE_FAT_BINARY', 'no') == 'yes' and conf['Intel?']:
   print 'Sage "fat" binary mode set: Building SSE2 only Hammer binary'
   print 'NOTE: This can result in a Sage that is significantly slower at certain numerical'
   print 'linear algebra since full FAT binary support has not been implemented yet.'
   arch = 'HAMMER'
   isa_ext = ('SSE2', 'SSE1')
elif os.environ.has_key('SAGE_ATLAS_ARCH'):
   opts = os.environ['SAGE_ATLAS_ARCH'].split(',')
   arch = opts[0]
   if len(opts)>0:
      isa_ext = opts[1:]

if arch == 'fast':
    arch, isa_ext = configure_opt_fast()
   
if arch == 'base':
    arch, isa_ext = configure_opt_base()

if arch is None:
    print 'First automatic tuning attempt.'
    rc = configure()    
    if rc!=0:
        print 'ATLAS failed to build, possibly because of a loaded system.'
        print 'Waiting 5 minutes...'
        time.sleep(5*60)
        print 'Second automatic tuning attempt.'
        rc = configure()
    if rc!=0:
        try:
            arch, isa_ext = configure_opt_fast()
            print 'Third attempt: use "fast" options.'
            rc = configure(arch, isa_ext)   
        except NotImplementedError:
            pass
    if rc!=0:
        try:
            arch, isa_ext = configure_opt_base()
            print 'Fourth attempt: use "base" options.'
            rc = configure(arch, isa_ext)
        except NotImplementedError:
            pass
else:
    print 'Running configure with arch =', arch, 'and isa extensions', isa_ext
    rc = configure(arch, isa_ext)


assert rc==0, 'Configure failed!'
print 'Finished configuring ATLAS.'


######################################################################
### make
######################################################################

# we need to disable parallel builds
os.environ.pop('MAKE', '')

def make_core():
    os.chdir('ATLAS-build')
    rc = os.system('make')
    os.chdir('..')
    return rc

def make_atlas_library():
    os.chdir('ATLAS-build/lib')

    if conf['Solaris?'] and conf['linker_Solaris?']:
        print "The Makefile generated in ATLAS for building shared libraries"
        print "assumes the linker is the GNU linker, which it not true in"
        print "your setup. (It is generally considered better to use the"
        print "Sun linker in /usr/ccs/bin rather than the GNU linker from binutils)"
        print "The linker flags in `pwd`/Makefile will be changed. "
        print "'-shared' will be changed to '-G'"
        print "'-soname' will be changed to '-h'"
        print "'--whole-archive' will be changed to '-zallextract'"
        print "'--no-whole-archive' will be changed to '-zdefaultextract'"
        edit_in_place('Makefile') \
            .replace('-shared', '-G') \
            .replace('-soname', '-h') \
            .replace('--whole-archive', '-z allextract') \
            .replace('--no-whole-archive', '-z defaultextract') \
            .close()
        if conf['Solaris?'] and conf['Intel?']:
            print "Change ldflag -melf_x86_64 to -64 as needed for Sun ld"
            print "on 64-bit builds of ATLAS on x64 hardware"
            edit_in_place('Make.inc').replace('-melf_x86_64', '-64').close()
        if conf['Solaris?'] and conf['Intel?']:
            print "Remove the linker flag -melf_i386 as needed for Sun ld"
            print "on 32-bit builds of ATLAS on x86/x64 hardware"
            edit_in_place('Make.inc').replace('-melf_i386', '').close()
        
    rc = os.system('make shared cshared')
    os.chdir('../..')
    return rc


def make_lapack_library():
    old_cwd = os.getcwd()
    os.chdir(os.environ['SAGE_LOCAL']+'/lib')
    sage_local_lib_dir = ' -L' + os.environ['SAGE_LOCAL'] + '/lib'

    libraries = []

    if conf['Linux?'] or conf['FreeBSD?']:
        if conf['fortran_g95?']:
            fortran_dir = '-L'+conf['fortran_g95_dir']
            fortran_lib = '-lf95'
        else:
            fortran_dir = ''
            fortran_lib = '-lgfortran'
        libraries = ['liblapack', 'libf77blas']
        cmd = 'gcc ' + fortran_dir + sage_local_lib_dir + \
            ' -shared -Wl,-soname,{0}.so -o {0}.so -lc -lm ' + fortran_lib

    if conf['Solaris?']:
        if conf['64bit?']:
            linker_flag_64 = '-64'
        else:
            linker_flag_64 = ''
        libraries = ['libatlas', 'libf77blas', 'libcblas']
        cmd = '/usr/ccs/bin/ld ' + linker_flag_64 + sage_local_lib_dir + \
            ' -G -h {0}.so -o {0}.so  -zallextract {0}.a -zdefaultextract -lc -lm -lgfortran'
        
    rc = 0
    for LIB in libraries:
        cmd_LIB = cmd.format(LIB)
        print 'Running', cmd_LIB
        rc_LIB = os.system(cmd_LIB)
        if rc_LIB!=0:
            print 'Failed to build ATLAS library '+LIB+'.so'
        rc = max(rc, abs(rc_LIB))
        
    # liblapack.so causes problems with R on Solaris. 
    try:
        os.remove('liblapack.so')
    except OSError:
        pass

    os.chdir(old_cwd)
    return rc



rc = make_core()
assert rc==0, 'Failed to build ATLAS.'
print 'Finished building ATLAS core.'

rc = make_atlas_library()
assert rc==0, 'Building shared ATLAS library failed.'
print 'Finished building shared ATLAS library.'

rc = make_lapack_library()
assert rc==0, 'Building LAPACK+ATLAS library failed.'
print 'Finished building LAPACK+ATLAS library.'


######################################################################
### install
######################################################################

def install():
    os.chdir('ATLAS-build')
    rc = os.system('make install')
    os.chdir('..')
    return rc

rc = install()
assert rc==0, 'Make install for ATLAS failed.'
print 'Finished installing ATLAS.'