# modified may 2011 to name components (map/ped) as RgeneticsData to align with default base_name

# otherwise downstream tools fail

# modified march  2011 to remove post execution hook  

# pedigree data faker

# specifically designed for scalability testing of

# Shaun Purcel's PLINK package

# derived from John Ziniti's original suggestion

# allele frequency spectrum and random mating added

# ross lazarus me fecit january 13 2007

# copyright ross lazarus 2007

# without psyco

# generates about 10k snp genotypes in 2k subjects (666 trios) per minute or so.

# so 500k (a billion genotypes), at about 4 trios/min will a couple of hours to generate

# psyco makes it literally twice as quick!!

# all rights reserved except as granted under the terms of the LGPL

# see http://www.gnu.org/licenses/lgpl.html 

# for a copy of the license you receive with this software

# and for your rights and obligations

# especially if you wish to modify or redistribute this code

# january 19 added random missingness inducer

# currently about 15M genos/minute without psyco, 30M/minute with

# so a billion genos should take about 40 minutes with psyco or 80 without...

# added mendel error generator jan 23 rml





import random,sys,time,os,string



from optparse import OptionParser



defbasename="RgeneticsData"    

width = 500000

ALLELES = ['1','2','3','4']

prog = os.path.split(sys.argv[0])[-1]

debug = 0



"""Natural-order sorting, supporting embedded numbers.

# found at http://lists.canonical.org/pipermail/kragen-hacks/2005-October/000419.html

note test code there removed to conserve brain space

foo9bar2 < foo10bar2 < foo10bar10



"""

import random, re, sys



def natsort_key(item): 

    chunks = re.split('(\d+(?:\.\d+)?)', item)

    for ii in range(len(chunks)):

        if chunks[ii] and chunks[ii][0] in '0123456789':

            if '.' in chunks[ii]: numtype = float

            else: numtype = int

            # wrap in tuple with '0' to explicitly specify numbers come first

            chunks[ii] = (0, numtype(chunks[ii]))

        else:

            chunks[ii] = (1, chunks[ii])

    return (chunks, item)



def natsort(seq):

    "Sort a sequence of text strings in a reasonable order."

    alist = [item for item in seq]

    alist.sort(key=natsort_key)

    return alist





def makeUniformMAFdist(low=0.02, high=0.5):

    """Fake a non-uniform maf distribution to make the data

    more interesting. Provide uniform 0.02-0.5 distribution"""

    MAFdistribution = []

    for i in xrange(int(100*low),int(100*high)+1):

       freq = i/100.0 # uniform

       MAFdistribution.append(freq)

    return MAFdistribution



def makeTriangularMAFdist(low=0.02, high=0.5, beta=5):

    """Fake a non-uniform maf distribution to make the data

    more interesting - more rare alleles """

    MAFdistribution = []

    for i in xrange(int(100*low),int(100*high)+1):

       freq = (51 - i)/100.0 # large numbers of small allele freqs

       for j in range(beta*i): # or i*i for crude exponential distribution 

            MAFdistribution.append(freq)

    return MAFdistribution



def makeFbathead(rslist=[], chromlist=[], poslist=[], width=100000):

    """header row

    """

    res = ['%s_%s_%s' % (chromlist[x], poslist[x], rslist[x]) for x in range(len(rslist))]

    return ' '.join(res)



def makeMap( width=500000, MAFdistribution=[], useGP=False):

    """make snp allele and frequency tables for consistent generation"""

    usegp = 1

    snpdb = 'snp126'

    hgdb = 'hg18'

    alleles = []

    freqs = []

    rslist = []

    chromlist = []

    poslist = []

    for snp in range(width):

        random.shuffle(ALLELES)

        alleles.append(ALLELES[0:2]) # need two DIFFERENT alleles!

        freqs.append(random.choice(MAFdistribution)) # more rare alleles

    if useGP:

        try:

            import MySQLdb

            genome = MySQLdb.Connect('localhost', 'hg18', 'G3gn0m3')

            curs = genome.cursor() # use default cursor

        except:

            if debug:

                print 'cannot connect to local copy of golden path'

            usegp = 0

    if usegp and useGP: # urrrgghh getting snps into chrom offset order is complicated....

        curs.execute('use %s' % hgdb)

        print 'Collecting %d real rs numbers - this may take a while' % width

        # get a random draw of enough reasonable (hapmap) snps with frequency data

        s = '''select distinct chrom,chromEnd, name from %s where avHet > 0 and chrom not like '%%random'

        group by name order by rand() limit %d''' % (snpdb,width)

        curs.execute(s)

        reslist = curs.fetchall()

        reslist = ['%s\t%09d\t%s' % (x[3:],y,z) for x,y,z in reslist] # get rid of chr

        reslist = natsort(reslist)

        for s in reslist:

            chrom,pos,rs = s.split('\t')

            rslist.append(rs)

            chromlist.append(chrom)

            poslist.append(pos)

    else:

        chrom = '1'

        for snp in range(width):

            pos = '%d' % (1000*snp)

            rs = 'rs00%d' % snp

            rslist.append(rs)

            chromlist.append(chrom)

            poslist.append(pos)

    return alleles,freqs, rslist, chromlist, poslist



def writeMap(fprefix = '', fpath='./', rslist=[], chromlist=[], poslist=[], width = 500000):

    """make a faked plink compatible map file - fbat files

    have the map written as a header line"""

    outf = '%s.map'% (fprefix)

    outf = os.path.join(fpath,outf)

    amap = open(outf, 'w')

    res = ['%s\t%s\t0\t%s' % (chromlist[x],rslist[x],poslist[x]) for x in range(len(rslist))]

    res.append('')

    amap.write('\n'.join(res))

    amap.close()



def makeMissing(genos=[], missrate = 0.03, missval = '0'):

    """impose some random missingness"""

    nsnps = len(genos)

    for snp in range(nsnps): # ignore first 6 columns

        if random.random() <= missrate:

            genos[snp] = '%s %s' % (missval,missval)

    return genos



def makeTriomissing(genos=[], missrate = 0.03, missval = '0'):

    """impose some random missingness on a trio - moth eaten like real data"""

    for person in (0,1):

        nsnps = len(genos[person])

        for snp in range(nsnps):

            for person in [0,1,2]:

                if random.random() <= missrate:

                    genos[person][snp] = '%s %s' % (missval,missval)

    return genos





def makeTriomendel(p1g=(0,0),p2g=(0,0), kiddip = (0,0)):

    """impose some random mendels on a trio

    there are 8 of the 9 mating types we can simulate reasonable errors for

    Note, since random mating dual het parents can produce any genotype we can't generate an interesting

    error for them, so the overall mendel rate will be lower than mendrate, depending on

    allele frequency..."""

    if p1g[0] <> p1g[1] and p2g[0] <> p2g[1]: # both parents het

            return kiddip # cannot simulate a mendel error - anything is legal!

    elif (p1g[0] <> p1g[1]): # p1 is het parent so p2 must be hom

        if p2g[0] == 0: # - make child p2 opposite hom for error

            kiddip = (1,1)

        else:

            kiddip = (0,0)

    elif (p2g[0] <> p2g[1]): # p2 is het parent so p1 must be hom

        if p1g[0] == 0: # - make child p1 opposite hom for error

            kiddip = (1,1)

        else:

            kiddip = (0,0)

    elif (p1g[0] == p1g[1]): # p1 is hom parent and if we get here p2 must also be hom

        if p1g[0] == p2g[0]: # both parents are same hom - make child either het or opposite hom for error

            if random.random() <= 0.5:

                kiddip = (0,1)

            else:

                if p1g[0] == 0:

                    kiddip = (1,1)

                else:

                    kiddip = (0,0)

        else: # parents are opposite hom - return any hom as an error

            if random.random() <= 0.5:

                kiddip = (0,0)

            else:

                kiddip = (1,1)

    return kiddip

            

            





def makeFam(width=100, freqs={}, alleles={}, trio=1, missrate=0.03, missval='0', mendrate=0.0):

    """this family is a simple trio, constructed by random mating two random genotypes

    TODO: why not generate from chromosomes - eg hapmap

    set each haplotype locus according to the conditional

    probability implied by the surrounding loci - eg use both neighboring pairs, triplets

    and quads as observed in hapmap ceu"""

    dadped = '%d 1 0 0 1 1 %s'

    mumped = '%d 2 0 0 2 1 %s' # a mother is a mum where I come from :)

    kidped = '%d 3 1 2 %d %d %s'

    family = [] # result accumulator

    sex = random.choice((1,2)) # for the kid

    affected = random.choice((1,2))

    genos = [[],[],[]] # dad, mum, kid - 0/1 for common,rare initially, then xform to alleles

    # parent1...kidn lists of 0/1 for common,rare initially, then xformed to alleles

    for snp in xrange(width):

        f = freqs[snp]           

        for i in range(2): # do dad and mum

            p = random.random()

            a1 = a2 = 0

            if p <= f: # a rare allele

               a1 = 1

            p = random.random()

            if p <= f: # a rare allele

               a2 = 1

            if a1 > a2:

                a1,a2 = a2,a1 # so ordering consistent - 00,01,11

            dip = (a1,a2)

            genos[i].append(dip) # tuples of 0,1

        a1 = random.choice(genos[0][snp]) # dad gamete  

        a2 = random.choice(genos[1][snp]) # mum gamete

        if a1 > a2:

            a1,a2 = a2,a1 # so ordering consistent - 00,01,11

        kiddip = (a1,a2) # NSFW mating!

        genos[2].append(kiddip)

        if mendrate > 0:

            if random.random() <= mendrate:

                genos[2][snp] = makeTriomendel(genos[0][snp],genos[1][snp], kiddip)

        achoice = alleles[snp]

        for g in genos: # now convert to alleles using allele dict

          a1 = achoice[g[snp][0]] # get allele letter

          a2 = achoice[g[snp][1]]              

          g[snp] = '%s %s' % (a1,a2)

    if missrate > 0:

        genos = makeTriomissing(genos=genos,missrate=missrate, missval=missval)

    family.append(dadped % (trio,' '.join(genos[0]))) # create a row for each member of trio

    family.append(mumped % (trio,' '.join(genos[1])))

    family.append(kidped % (trio,sex,affected,' '.join(genos[2])))

    return family



def makePerson(width=100, aff=1, freqs={}, alleles={}, id=1, missrate = 0.03, missval='0'):

    """make an entire genotype vector for an independent subject"""

    sex = random.choice((1,2))

    if not aff:

        aff = random.choice((1,2))

    genos = [] #0/1 for common,rare initially, then xform to alleles

    family = []

    personped = '%d 1 0 0 %d %d %s'

    poly = (0,1)

    for snp in xrange(width):

        achoice = alleles[snp]

        f = freqs[snp]

        p = random.random()

        a1 = a2 = 0

        if p <= f: # a rare allele

           a1 = 1

        p = random.random()

        if p <= f: # a rare allele

           a2 = 1

        if a1 > a2:

            a1,a2 = a2,a1 # so ordering consistent - 00,01,11

        a1 = achoice[a1] # get allele letter

        a2 = achoice[a2]

        g = '%s %s' % (a1,a2)

        genos.append(g)

    if missrate > 0.0:

        genos = makeMissing(genos=genos,missrate=missrate, missval=missval)

    family.append(personped % (id,sex,aff,' '.join(genos)))

    return family



def makeHapmap(fprefix= 'fakebigped',width=100, aff=[], freqs={},

               alleles={}, nsubj = 2000, trios = True, mendrate=0.03, missrate = 0.03, missval='0'):

    """ fake a hapmap file and a pedigree file for eg haploview

    this is arranged as the transpose of a ped file - cols are subjects, rows are markers

    so we need to generate differently since we can't do the transpose in ram reliably for

    a few billion genotypes...

    """

    outheadprefix = 'rs# alleles chrom pos strand assembly# center protLSID assayLSID panelLSID QCcode %s'

    cfake5 = ["illumina","urn:LSID:illumina.hapmap.org:Protocol:Golden_Gate_1.0.0:1", 

"urn:LSID:illumina.hapmap.org:Assay:27741:1","urn:lsid:dcc.hapmap.org:Panel:CEPH-30-trios:1","QC+"]

    yfake5 = ["illumina","urn:LSID:illumina.hapmap.org:Protocol:Golden_Gate_1.0.0:1", 

"urn:LSID:illumina.hapmap.org:Assay:27741:1","urn:LSID:dcc.hapmap.org:Panel:Yoruba-30-trios:1","QC+"]

    sampids = ids

    if trios:

        ts = '%d trios' % int(nsubj/3.)

    else:

        ts = '%d unrelated subjects' % nsubj

    res = ['#%s fake hapmap file %d snps and %s, faked by %s' % (timenow(), width, ts, prog),]

    res.append('# ross lazarus me fecit')

    res.append(outheadprefix % ' '.join(sampids)) # make a header compatible with hapmap extracts

    outf = open('%s.hmap' % (fprefix), 'w')

    started = time.time()

    if trios:

        ntrios = int(nsubj/3.)

        for n in ntrios: # each is a dict

            row = copy.copy(cfake5) # get first fields

            row = map(str,row)

            if race == "YRI":

                row += yfake5

            elif race == 'CEU':

                row += cfake5

            else:

                row += ['NA' for x in range(5)] # 5 dummy fields = center protLSID assayLSID panelLSID QCcode

            row += [''.join(sorted(line[x])) for x in sampids] # the genotypes in header (sorted) sample id order

            res.append(' '.join(row))

    res.append('')

    outfname = '%s_%s_%s_%dkb.geno' % (gene,probeid,race,2*flank/1000)

    f = file(outfname,'w')

    f.write('\n'.join(res))

    f.close()

    print '### %s: Wrote %d lines to %s' % (timenow(), len(res),outfname)

     



def makePed(fprefix= 'fakebigped', fpath='./',

            width=500000, nsubj=2000, MAFdistribution=[],alleles={},

            freqs={}, fbatstyle=True, mendrate = 0.0, missrate = 0.03, missval='0',fbathead=''):

    """fake trios with mendel consistent random mating genotypes in offspring

    with consistent alleles and MAFs for the sample"""

    res = []

    if fbatstyle: # add a header row with the marker names

        res.append(fbathead) # header row for fbat

    outfname = '%s.ped'% (fprefix)

    outfname = os.path.join(fpath,outfname)

    outf = open(outfname,'w')

    ntrios = int(nsubj/3.)

    outf = open(outfile, 'w')

    started = time.time()

    for trio in xrange(ntrios):

        family = makeFam(width=width, freqs=freqs, alleles=alleles, trio=trio,

                         missrate = missrate, mendrate=mendrate, missval=missval)

        res += family

        if (trio + 1) % 10 == 0: # write out to keep ram requirements reasonable

            if (trio + 1) % 50 == 0: # show progress

                dur = time.time() - started

                if dur == 0:

                    dur = 1.0

                print 'Trio: %d, %4.1f genos/sec at %6.1f sec' % (trio + 1, width*trio*3/dur, dur)

            outf.write('\n'.join(res))

            outf.write('\n')

            res = []

    if len(res) > 0: # some left

        outf.write('\n'.join(res))

    outf.write('\n')

    outf.close()

    if debug:

        print '##makeped : %6.1f seconds total runtime' % (time.time() - started)



def makeIndep(fprefix = 'fakebigped', fpath='./',

              width=500000, Nunaff=1000, Naff=1000, MAFdistribution=[],

              alleles={}, freqs={}, fbatstyle=True, missrate = 0.03, missval='0',fbathead=''):

    """fake a random sample from a random mating sample

    with consistent alleles and MAFs"""

    res = []

    Ntot = Nunaff + Naff

    status = [1,]*Nunaff

    status += [2,]*Nunaff

    outf = '%s.ped' % (fprefix)

    outf = os.path.join(fpath,outf)

    outf = open(outf, 'w')

    started = time.time()

    #sample = personMaker(width=width, affs=status, freqs=freqs, alleles=alleles, Ntomake=Ntot)

    if fbatstyle: # add a header row with the marker names

        res.append(fbathead) # header row for fbat

    for id in xrange(Ntot):

        if id < Nunaff:

            aff = 1

        else:

            aff = 2

        family = makePerson(width=width, aff=aff, freqs=freqs, alleles=alleles, id=id+1)

        res += family

        if (id % 50 == 0): # write out to keep ram requirements reasonable

            if (id % 200 == 0): # show progress

                dur = time.time() - started

                if dur == 0:

                    dur = 1.0

                print 'Id: %d, %4.1f genos/sec at %6.1f sec' % (id, width*id/dur, dur)

            outf.write('\n'.join(res))

            outf.write('\n')

            res = []

    if len(res) > 0: # some left

        outf.write('\n'.join(res))

    outf.write('\n')

    outf.close()

    print '## makeindep: %6.1f seconds total runtime' % (time.time() - started)



u = """

Generate either trios or independent subjects with a prespecified

number of random alleles and a uniform or triangular MAF distribution for

stress testing. No LD is simulated - alleles are random. Offspring for

trios are generated by random mating the random parental alleles so there are

no Mendelian errors unless the -M option is used. Mendelian errors are generated

randomly according to the possible errors given the parental mating type although

this is fresh code and not guaranteed to work quite right yet - comments welcomed



Enquiries to ross.lazarus@gmail.com



eg to generate 700 trios with 500k snps, use:

fakebigped.py -n 2100 -s 500000

or to generate 500 independent cases and 500 controls with 100k snps and 0.02 missingness (MCAR), use:

fakebigped.py -c 500 -n 1000 -s 100000 -m 0.02



fakebigped.py -o myfake -m 0.05 -s 100000 -n 2000

will make fbat compatible myfake.ped with 100k markers in

666 trios (total close to 2000 subjects), a uniform MAF distribution and about 5% MCAR missing



fakebigped.py -o myfake -m 0.05 -s 100000 -n 2000 -M 0.05

will make fbat compatible myfake.ped with 100k markers in

666 trios (total close to 2000 subjects), a uniform MAF distribution,

about 5% Mendelian errors and about 5% MCAR missing





fakebigped.py  -o myfakecc -m 0.05 -s 100000 -n 2000 -c 1000 -l

will make plink compatible myfakecc.ped and myfakecc.map (that's what the -l option does),

with 100k markers in 1000 cases and 1000 controls (affection status 2 and 1 respectively),

a triangular MAF distribution (more rare alleles) and about 5% MCAR missing



You should see about 1/4 million genotypes/second so about an hour for a

500k snps in 2k subjects and about a 4GB ped file - these are BIG!!



"""



import sys, os, glob



galhtmlprefix = """<?xml version="1.0" encoding="utf-8" ?>

<!DOCTYPE html PUBLIC "-//W3C//DTD XHTML 1.0 Transitional//EN" "http://www.w3.org/TR/xhtml1/DTD/xhtml1-transitional.dtd">

<html xmlns="http://www.w3.org/1999/xhtml" xml:lang="en" lang="en">

<head>

<meta http-equiv="Content-Type" content="text/html; charset=utf-8" />

<meta name="generator" content="Galaxy %s tool output - see http://g2.trac.bx.psu.edu/" />

<title></title>

<link rel="stylesheet" href="/static/style/base.css" type="text/css" />

</head>

<body>

<div class="document">

"""





def doImport(outfile=None,outpath=None):

    """ import into one of the new html composite data types for Rgenetics

        Dan Blankenberg with mods by Ross Lazarus 

        October 2007

    """

    flist = glob.glob(os.path.join(outpath,'*'))

    outf = open(outfile,'w')

    outf.write(galhtmlprefix % prog)

    for i, data in enumerate( flist ):

        outf.write('<li><a href="%s">%s</a></li>\n' % (os.path.split(data)[-1],os.path.split(data)[-1]))

    outf.write('<br><h3>This is simulated null genotype data generated by Rgenetics!</h3>')

    outf.write('%s called with command line:<br><pre>' % prog)

    outf.write(' '.join(sys.argv))

    outf.write('\n</pre>\n')

    outf.write("</div></body></html>")

    outf.close()







if __name__ == "__main__":

    """

    """

    parser = OptionParser(usage=u, version="%prog 0.01")

    a = parser.add_option

    a("-n","--nsubjects",type="int",dest="Ntot",

      help="nsubj: total number of subjects",default=2000)

    a("-t","--title",dest="title",

      help="title: file basename for outputs",default='fakeped')

    a("-c","--cases",type="int",dest="Naff",

      help="number of cases: independent subjects with status set to 2 (ie cases). If not set, NTOT/3 trios will be generated", default = 0)

    a("-s","--snps",dest="width",type="int",

      help="snps: total number of snps per subject", default=1000)

    a("-d","--distribution",dest="MAFdist",default="Uniform",

      help="MAF distribution - default is Uniform, can be Triangular")

    a("-o","--outf",dest="outf",

      help="Output file", default = 'fakeped')

    a("-p","--outpath",dest="outpath",

      help="Path for output files", default = './')

    a("-l","--pLink",dest="outstyle", default='L',

      help="Ped files as for Plink - no header, separate Map file - default is Plink style")

    a("-w","--loWmaf", type="float", dest="lowmaf", default=0.01, help="Lower limit for SNP MAF (minor allele freq)")

    a("-m","--missing",dest="missrate",type="float",

      help="missing: probability of missing MCAR - default 0.0", default=0.0)

    a("-v","--valmiss",dest="missval",

      help="missing character: Missing allele code - usually 0 or N - default 0", default="0")

    a("-M","--Mendelrate",dest="mendrate",type="float",

      help="Mendelian error rate: probability of a mendel error per trio, default=0.0", default=0.0)   

    a("-H","--noHGRS",dest="useHG",type="int",

      help="Use local copy of UCSC snp126 database to generate real rs numbers", default=True)

    (options,args) = parser.parse_args()

    low = options.lowmaf

    try:

        os.makedirs(options.outpath)

    except:

        pass

    if options.MAFdist.upper() == 'U':

        mafDist = makeUniformMAFdist(low=low, high=0.5)

    else:

        mafDist = makeTriangularMAFdist(low=low, high=0.5, beta=5)

    alleles,freqs, rslist, chromlist, poslist = makeMap(width=int(options.width),

                                        MAFdistribution=mafDist, useGP=False)

    fbathead = []

    s = string.whitespace+string.punctuation

    trantab = string.maketrans(s,'_'*len(s))

    title = string.translate(options.title,trantab)

    

    if options.outstyle == 'F':

        fbatstyle = True

        fbathead = makeFbathead(rslist=rslist, chromlist=chromlist, poslist=poslist, width=options.width)

    else:

        fbatstyle = False

        writeMap(fprefix=defbasename, rslist=rslist, fpath=options.outpath,

                 chromlist=chromlist, poslist=poslist, width=options.width)

    if options.Naff > 0: # make case control data

        makeIndep(fprefix = defbasename, fpath=options.outpath,

                  width=options.width, Nunaff=options.Ntot-options.Naff,

                  Naff=options.Naff, MAFdistribution=mafDist,alleles=alleles, freqs=freqs,

                  fbatstyle=fbatstyle, missrate=options.missrate, missval=options.missval,

                  fbathead=fbathead)

    else:

        makePed(fprefix=defbasename, fpath=options.fpath,

            width=options.width, MAFdistribution=mafDist, nsubj=options.Ntot,

            alleles=alleles, freqs=freqs, fbatstyle=fbatstyle, missrate=options.missrate,

            mendrate=options.mendrate, missval=options.missval,

                  fbathead=fbathead)

    doImport(outfile=options.outf,outpath=options.outpath)