#!/usr/bin/env python

import os,glob,string,sys
from numpy.distutils.core  import setup, Extension
from numpy.distutils import fcompiler
from distutils.dep_util import newer

## -------- set these
KM = 26 
JM = 1
IM = 1
NC_INC = os.getenv('HOME')+'/include'
NC_LIB = os.getenv('HOME')+'/lib'
##----------------------

Extensions = [
    {'name':'grid',
     'dir':'src/grid'},
    {'name':'thermodyn',
     'dir':'src/thermodyn'},
    {'name':'ozone',
     'dir':'src/radiation/ozone'},
    {'name':'insolation',
     'dir':'src/radiation/insolation'},
    {'name':'ccm3_radiation',
     'dir':'src/radiation/ccm3',
     'cppflags':'-DSUN -DPLON=%i -DPLEV=%i -DPLEVR=%i' % (IM,KM,KM)}
    ]

# figure out which compiler we're goint to use
compiler = fcompiler.get_default_fcompiler()
for i in range(len(sys.argv)):
    if '--fcompiler' in sys.argv[i]:
        compiler = sys.argv.pop(i)
        compiler = compiler[compiler.index('=')+1:]
print 'Using %s compiler' % compiler

# set some fortran compiler-dependent flags
if compiler == 'gnu95':
    f77flags='-ffixed-line-length-132 -fdefault-real-8'
    f90flags='-fdefault-real-8'
elif compiler == 'intel' or compiler == 'intelem':
    f77flags='-132 -r8 -w95 -w90 -mp'
    f90flags='-r8 -w95 -mp'
elif compiler == 'ibm':
    f77flags='-qautodbl=dbl4 -qsuffix=f=f:cpp=F -qfixed=132'
    f90flags='-qautodbl=dbl4 -qsuffix=f=f90:cpp=F90 -qfree=f90'
else:
    print 'Sorry, compiler %s not supported' % compiler
    sys.exit()

for i in range(len(Extensions)):
    Defaults = {'cppflags':'-DIM=%i -DJM=%i -DKM=%i' % (IM,JM,KM),
                'f77flags':f77flags,
                'f90flags':f90flags}
    Defaults.update(Extensions[i])
    Extensions[i] = Defaults

if compiler == 'ibm':
    for ext in Extensions:
        ext['cppflags']='-WF,'+string.join(ext['cppflags'].split(),',')

def getSources(dir):
    #Gets list of source files for extensions
    SrcFile = os.path.join(dir,'sources_in_order_of_compilation')
    if os.path.exists(SrcFile):
        Sources = open(SrcFile).readlines()
        Sources = [os.path.join(dir,s[:-1]) for s in Sources]
    else:
        Sources = []
        for pattern in ['*.f','*.F','*.f90','*.F90']:
            Sources += glob.glob(os.path.join(dir,pattern))
            Sources += glob.glob(os.path.join(dir,'src',pattern))        
    return Sources

def buildNeeded(target,src):
    #Checks if source code is newer than extension, so extension needs to be rebuilt
    target = os.path.join('lib/climt_lite',target)
    if not os.path.exists(target):
        return True
    for file in src:
        if newer(file,target):
            return True
    print 'Extension %s is up to date' % os.path.basename(target)
    return False

def build_ext(name=None, dir=None, cppflags='', f77flags='', f90flags='', \
              lib='', libdir='', incdir=''):
    #Builds an extension
    src = getSources(dir)
    target = '_%s.so' % name
    driver = glob.glob(os.path.join(dir,'Driver.f*'))[0]
    f77flags = '"%s %s"' % (cppflags,f77flags)
    f90flags = '"%s %s"' % (cppflags,f90flags)
    if buildNeeded(target,src):
        print '\n Building %s ... \n' % os.path.basename(target)
        # generate signature file
        os.system('f2py --overwrite-signature %s -m _%s -h _%s.pyf'%(driver,name,name))
        # compile extension
        F2pyCommand = []
        F2pyCommand.append('f2py -c -m _%s' % name)
        F2pyCommand.append('--fcompiler=%s' % compiler)
        F2pyCommand.append('-I%s' % dir)
        F2pyCommand.append('-I%s' % os.path.join(dir,'include'))
        F2pyCommand.append('-I%s' % os.path.join(dir,'src'))
        F2pyCommand.append('-I%s' % os.path.join(dir,'src','include'))
        if incdir is not '': F2pyCommand.append('-I%s' % incdir)
        if libdir is not '': F2pyCommand.append('-L%s' % libdir)
        if lib is not '': F2pyCommand.append('-l%s' % lib)
        F2pyCommand.append('--f77flags=%s' % f77flags)
        F2pyCommand.append('--f90flags=%s' % f90flags)
        F2pyCommand.append('_%s.pyf' % name)
        F2pyCommand.append('%s' % string.join(src))
        F2pyCommand = string.join(F2pyCommand)
        print F2pyCommand
        if os.system(F2pyCommand) > 0:
            print '+++ Compilation failed'
            sys.exit()
        os.system('mv -f _%s.so lib/climt_lite' % name)
        os.system('rm -f _%s.pyf' % name)

# Build all extensions
for ext in Extensions: build_ext(**ext)

setup(name         = "CliMT_lite",
      version      = open('Version').read()[:-1],
      description  = "Climate modelling and diagnostics toolkit",
      author       = "Rodrigo Caballero",
      author_email = "rodrigo.caballero@ucd.ie",
      url          = "http://mathsci.ucd.ie/~rca/climt",
      packages    = ['climt_lite'],
      package_dir = {'':'lib'},
      package_data = {'climt_lite':['*.so']})