#!/usr/bin/env python
import spss,spssaux
""" Notes:
    Each model is currently limited to 4 vars plus an interaction term
        for instances: var1 gender race age var1*gender*race*age
    scripts creates ageAdjusted.SPS in same folder as data file. This syntax may be used to
        rerun the analysis or fine tune the syntax.
    script also works when SET OLANG is not ENGLISH

    TO DO:
    test when age bands of data and reference pop are not compatible

Here is a sample syntax file that analyses 4 models using ageAdjusted.py:

BEGIN PROGRAM.
#!/usr/bin/env python
import ageAdjusted as aa
aa.master(dataFile=r'c:\Mes Documents\CDC\Age_Adj Exercise NHANES_labeled.sav',
           planFile='mec4yr_plan.csaplan',
           referencePop=r'c:\Mes Documents\CDC\US Census pop 2000.sav',
           depVar='hbpx',
           domainVar='sel(1)',
           factors=['race age_grp',
                  'age_grp',
                  'riagendr race age_grp',
                  'riagendr age_grp'])
END PROGRAM.
"""

__author__  =  'Raynald Levesque raynald@spsstools.net'
__version__ =  '1.11'
version = __version__

def _getPopFiles(dataFile, referencePop, ageVar):
    import re
    spss.Submit(r"""CD "%(dataFile)s".
                SET MPRINT=YES /PRINTBACK=YES.""" % locals())
    #Get age bands from reference population file
    cmd=r"""
    DATASET CLOSE ALL .
    GET FILE='%(referencePop)s.'
    """ % locals()
    spss.Submit(cmd)
    varDict=spssaux.VariableDict()
    ageRange=varDict[0] #get variable name of age ranges
    pop=varDict[1]      #get variable name containing population within range

    cmd=r"""
    DATASET CLOSE ALL .
    GET FILE="%(referencePop)s".
    COMPUTE idx=INDEX(%(ageRange)s,"-+",1).
    DO IF idx>0.
      COMPUTE beg=NUMBER(SUBSTR(%(ageRange)s,1,idx-1),F3).
      COMPUTE end=NUMBER(SUBSTR(%(ageRange)s,idx+1),F3).
    ELSE IF idx=0.
      COMPUTE beg=NUMBER(%(ageRange)s,F3).
      COMPUTE end=beg .
    END IF.
    IF SYSMIS(end) end=999.
    SAVE OUTFILE="_temp.sav" /DROP=idx.
    """ % locals()
    spss.Submit(cmd)

    #Get age bands from data file
    cmd=r"""
    DATASET CLOSE ALL .
    GET FILE='%(dataFile)s.'
    """ % locals()
    spss.Submit(cmd)

    varDict=spssaux.VariableDict(pattern=r'\b%(ageVar)s\b' % locals())
    ageLabels=varDict.valueLabels(varDict.variables[0])
    cmd=[]

    for (idx,lab) in enumerate(sorted(ageLabels)):
        print 'ageLabels[lab]',ageLabels[lab]
        beg,sep,end = re.split('^(?P<beg>\d*)(?P<oper>[-|+]{0,1})(?P<end>\d*)$',ageLabels[lab])[1:4]
        if not sep:     #single value is given
            if end: beg=end
            elif beg: end=beg
        else:           #single value with either - or +
            if not end: end='999'
            if not beg: beg=0

        cnt = idx + 1
        print beg,' and ', end
        cmdTmp=[r"""
        GET FILE="_temp.sav".
        COMPUTE cnt = %(cnt)s .
        COMPUTE target1 = %(beg)s.
        COMPUTE target2 = %(end)s.
        COMPUTE keepit = (target1 LE end) AND (target2 GE beg).
        SELECT IF keepit.
        SAVE OUTFILE="_pop%(cnt)s.sav".""" % locals()]
        cmd.extend(cmdTmp)
    cmd='\n'.join(cmd)
    spss.Submit(cmd)
    return cnt

def addFiles(stem, nb, byVar=None):
    """adds up to 50 sav files
    eg addFiles('pop',3) ADD FILES FILE='pop1.sav' /FILE='pop2.sav' FILE='pop3.sav'.
    """
    cmd=['ADD FILES ']
    cmd2=['  /FILE="%s%s.sav" ' % (stem,idx) for idx in range(1,nb+1)]
    if byVar:
        cmd2.extend(' /BY=%s.' % byVar)
    cmd.extend(cmd2)
    cmd.extend('.')
    spss.Submit('\n'.join(cmd))

def _getPopRatios(fileName):
    cmd=r"""
            AGG OUT=*
                /BREAK=cnt
                /nb=SUM(Population)
                /target1 target2 =FIRST(target1 target2) .
            COMPUTE nobreak=1.
            AGG
              /BREAK=nobreak
              /totPop=SUM(nb).
            SAVE OUTFILE='%s'.""" % fileName
    spss.Submit(cmd)

def _getAgeFact(fileName, dataFile, referencePop, ageVar):
    import spssdata
    cnt=_getPopFiles(dataFile, referencePop, ageVar)
    addFiles('_pop',cnt)
    _getPopRatios(fileName)
    spss.Submit(r"""GET FILE='%s'.""" % fileName)
    data=spssdata.Spssdata(indexes=('nb','totPop'))
    mydata=[(row.nb,row.totPop) for row in data]
    ageFact=['%u/%u ' % (mydata[nb][0],mydata[nb][1]) for nb in range(len(mydata))]
    del data
    return cnt,' '.join(ageFact)

def _mult(vector,age):
    """ vector: a list such as ['0','1','0','0']
        age: a string containing the age weights for instance '.3966 .3718 .2316'
        returns '0 0 0 .3966 .3718 .2316 0 0 0 0 0 0'
    """
    tmp=[]
    nb=len(age.split())
    age1=['0 '*nb,age]
    for n in vector:
        tmp.append(age1[int(n)])
    return " ".join(tmp)

def _f2(row1, row2):
    returnVal=[]
    rowzero=['0']*len(row2)
    for n in row1:
        returnVal.extend([rowzero, row2][int(n)])
    return returnVal

def _f3(model,ageFact, nb):
    """To contruct all /CUSTOM subcommands """
    model=model.split()
    customCmd=[]
    print "factors =",nb
    #nb.extend([None]*(4-len(nb)))
    nb.extend([None])
    test=[]
    for n0 in range(nb[0]):
        vec0=['0']*nb[0]
        vec0[n0]='1'
        if nb[1]:
            for n1 in range(nb[1]):
                vec1=['0']*nb[1]
                vec1[n1]='1'
                if nb[2]:
                    for n2 in range(nb[2]):
                        vec2=['0']*nb[2]
                        vec2[n2]='1'
                        if nb[3]:
                            for n3 in range(nb[3]):
                                vec3=['0']*nb[3]
                                vec3[n3]='1'
                                row1=_f2(vec0,vec1)  # combine first 2 vectors
                                row2=_f2(row1,vec2)  # combine result with 3rd vector
                                row3=_f2(row2,vec3)  # combine result with 4th vector
                                cmd=[model[0],' ']
                                cmd.extend([             ' '.join(vec0),' ',
                                            model[1],' ',' '.join(vec1),' ',
                                            model[2],' ',' '.join(vec2),' ',
                                            model[3],' ',' '.join(vec3),' ',
                                            model[4],' ',ageFact,' ',
                                            model[5],' ',_mult(row3, ageFact)]) #create LMATRIX weight
                        else:
                            row1=_f2(vec0,vec1)  # combine first 2 vectors
                            row2=_f2(row1,vec2)  # combine result with 3rd vector
                            cmd=[model[0],' ']
                            cmd.extend([             ' '.join(vec0),' ',
                                        model[1],' ',' '.join(vec1),' ',
                                        model[2],' ',' '.join(vec2),' ',
                                        model[3],' ',ageFact,' ',
                                        model[4],' ',_mult(row2, ageFact)]) #create LMATRIX weight
                        customCmd.append(cmd)
                else:
                    row1=_f2(vec0,vec1)
                    cmd=[model[0],' ']
                    cmd.extend([             ' '.join(vec0),' ',
                                model[1],' ',' '.join(vec1),' ',
                                model[2],' ',ageFact,' ',
                                model[3],' ',_mult(row1, ageFact)])
                    customCmd.append(cmd)
        else:               # there is only 1 factor in model (ignoring age and interaction)
            cmd=[model[0],' ']
            cmd.extend([' '.join(vec0),' ',
                        model[1],' ',ageFact,' ',
                        model[2],' ',_mult(vec0, ageFact)])
            customCmd.append(cmd)
    return customCmd

def deleteFiles(fnames):
    import os
    for fname in fnames:
        try:
            os.remove(fname)
            print '** Deleted file: ',fname
        except OSError:
            pass

def master(dataFile, planFile, referencePop, depVar, domainVar, factors):
    import os.path, time
    myTime=time.localtime()
    pathd=os.path.dirname(dataFile)         #get path of data files
    sFile=open(r"%(pathd)s\ageAdjusted.SPS" % locals(),'w')
    sFile.write("*Created by ageAdjusted.py version %(__version__)s\n" % globals())
    sFile.write("*on %s/%s/%s  %s:%s .\n" % myTime[:5])

    models=factors[:]
    for idx in range(len(factors)):
        if len(factors[idx].split()) > 1:   #add the interaction term
            models[idx]=factors[idx] + ' ' + '*'.join(factors[idx].split())

    ageVar=factors[0].split()[-1]
    cnt, ageFact = _getAgeFact('_pop by age ranges.sav', dataFile, referencePop, ageVar)
    allLabels = []
    cmd=r"""
SET MPRINT=YES /PRINTBACK=YES.
CD "%(pathd)s".
DATASET CLOSE ALL.
GET FILE='%(dataFile)s'.
OMS SELECT TABLES
   /IF SUBTYPES=['IndividualTestResults' 'TestsofModelEffects']
   /DESTINATION FORMAT=SAV OUTFILE="%(pathd)s/_indivTest.sav".
   """ % locals()
    try:
        sFile.write(cmd)
        spss.Submit(cmd)
    except:
        sFile.close()
        raise

    for (factor,model) in zip(factors,models):
        byVar=factor.split()
        if len(byVar) > 1:
            searchL=[r"\b%s\b" % (var) for var in byVar[0:-1]] # no need to search for the age var
            varLab=[]
            for vname in searchL:   # done 1 by 1 to preserve order
                varDict=spssaux.VariableDict(pattern=vname) #case of var names are ignored
                varLab.extend([varDict.valueLabels(vname) for vname in varDict])
            if len(byVar) == 2:
                customNames=[varLab[0][str(val0)] for val0 in range(1,len(varLab[0])+1)]
            elif len(byVar) == 3:
                customNames=[varLab[0][str(val0)] + " " + varLab[1][str(val1)]
                     for val0 in range(1,len(varLab[0])+1)
                     for val1 in range(1,len(varLab[1])+1)]
            elif len(byVar) == 4:
                customNames=[varLab[0][str(val0)] + " " + varLab[1][str(val1)] + " " + varLab[2][str(val2)]
                     for val0 in range(1,len(varLab[0])+1)
                     for val1 in range(1,len(varLab[1])+1)
                     for val2 in range(1,len(varLab[2])+1)]
            customCmd=_f3(model,ageFact,[len(a) for a in varLab])
        else:
            customNames=["Total"]
            customCmd=['%(model)s %(ageFact)s' % locals()]

        cmd=[]
        for [idx,name] in enumerate(customNames):
            str2=''.join(customCmd[idx])
            cmd.append("""   /CUSTOM LABEL="%s"\n    LMATRIX=%s """ % (name,str2))
        byVarStr=' '.join(byVar)
        cmdStr='\n'.join(cmd)
        cmd1=r"""
CSGLM  %(depVar)s BY %(byVarStr)s
   /PLAN FILE='%(planFile)s'
   /MODEL %(model)s
   /INTERCEPT INCLUDE=NO SHOW=YES
%(cmdStr)s
   /PRINT SUMMARY VARIABLEINFO SAMPLEINFO
   /STATISTICS PARAMETER SE TTEST
   /DOMAIN VARIABLE= %(domainVar)s
   /TEST TYPE=F PADJUST=LSD
   /MISSING CLASSMISSING=EXCLUDE
   /CRITERIA CILEVEL=95.
          """ % locals()
        try:
            sFile.write(cmd1)
            spss.Submit(cmd1)
        except:
            print "error writing or running syntax:\n",cmd1
            sFile.close()
            raise

        allLabels.extend(['(Model)'])
        allLabels.extend(model.split())
        allLabels.extend(customNames)

    defineLabels = ['VALUE LABELS customID ']
    defineLabels.extend(['   %s "%s" ' % (n+1,lab) for (n,lab) in enumerate(allLabels)])
    defineLabels.extend(['.'])
    defineCmd = '\n '.join(defineLabels)
    cmd=r"""
        OMSEND.
        CD "%(pathd)s".
        GET FILE='_indivTest.sav'.
        COMPUTE customID=$CASENUM.
        EXECUTE /* This EXECUTE is required */.""" % locals()
    try:
        spss.Submit(cmd)
    except:
        sFile.close()
        raise

    cmd=r"""
OMSEND.
GET FILE='_indivTest.sav'.
COMPUTE customID=$CASENUM.
    """
    sFile.write(cmd)

    nbVars=spss.GetVariableCount()  #Get variable names (they vary with SET OLANG command used)
    varList=[spss.GetVariableName(i) for i in range(nbVars)]
    varNames = [varList[i] for i in [2,12]]
    varNames.extend([varList[i] for i in range(3,12)])
    df1=varNames[3]
    df2=varNames[4]
    WaldF =varNames[5]
    ContrastEstimate=varNames[7]
    StdError = varNames[10]
    sig=varNames[6]

    cmd=r"""
%(defineCmd)s
FORMATS %(ContrastEstimate)s %(WaldF)s(F20.3) %(StdError)s %(sig)s(F20.8)%(df1)s %(df2)s(F4) .
SUMMARIZE
   /TABLES=Label_ customID %(ContrastEstimate)s %(StdError)s %(df1)s %(df2)s %(WaldF)s %(sig)s
   /FORMAT=VALIDLIST NOCASENUM TOTAL
   /TITLE='Summary of Tests of Model Effects and of Individual Test Results'
   /MISSING=VARIABLE
   /CELLS=NONE.
NEW FILE.
   """ % locals()
    sFile.write(cmd)
    sFile.close()
    spss.Submit(cmd)

    # reach this point without problems --> delete temporary files
    # try to delete up to 20 files (or current number of files if larger)
    # in case temporary files from earlier runs are still present
    filesToDelete=[r'%s/_pop%s.sav' % (pathd,n) for n in range(1,max(cnt + 1,21))]
    filesToDelete.extend([r'%s/_pop by age ranges.sav' % pathd,
                          r'%s/_indivTest.sav' % pathd,
                          r'%s/_temp.sav' % pathd])
    deleteFiles(filesToDelete)


def main():
    master(dataFile=r'c:\Mes Documents\CDC\Age_Adj Exercise NHANES_labeled.sav',
           planFile='mec4yr_plan.csaplan',
           referencePop=r'c:\Mes Documents\CDC\US Census pop 2000.sav',
           depVar='hbpx',
           domainVar='sel(1)',
           factors=['race age',
                  'age',
                  'riagendr race age',
                  'riagendr age'])

def main2():
    """ temporary function used for testing purposes """
    master(dataFile=r'c:\Mes Documents\CDC\Problem2\2000_2003 NHIS3.sav',
           planFile='survey.csaplan',
           referencePop=r'c:\Mes Documents\CDC\US Census pop 2000.sav',
           depVar='HRT_DIS',
           domainVar='sel2(1)',
           factors=['sex hsp_nhsp rec_age4',
                    'rec_age4'])
if __name__ == "__main__":
    main2()
