options(stringsAsFactors=FALSE)

library(rtracklayer)
session <- browserSession("UCSC")
genome(session) <- "hg18"
#trackNames(session) ## list the track names
q <- ucscTableQuery(session, "knownGene")
print(date())
knownGene <- getTable(q)
print(date())


# merge TSSs with identical starts
knownGene$position <- ifelse(knownGene$strand=="+", knownGene$txStart, knownGene$txEnd)
key <- paste(knownGene$chrom, knownGene$position, sep=":")
ukey <- unique(key)
m <- match(ukey, key)
mkg <- knownGene[m,c("chrom","strand","txStart","txEnd","name","position")]  # mkg = merged known genes

# create merged IDs
rownames(mkg) <- ukey
ids <- split(knownGene$name, key)
collapsedids <- sapply(ids, paste, collapse=",")
mkg$newID <- ""
mkg[names(collapsedids),"newID"] <- collapsedids

# merge nearby TSSs 
tssTol <- 500  # if distance b/w TSSs is less than this, they are merged
splitTol <- 250000

lockey <- paste(mkg$chrom, mkg$strand, sep=":")
ind <- seq_len(nrow(mkg))
inds <- split(ind, lockey)

clustno <- lapply(inds, FUN=function(u) {
  w <- diff(mkg$position[u]) > splitTol
  z <- cumsum(c(0,w))
  s <- split(u, z)
  cat(".")
  clusts <- lapply(s, FUN=function(uu) {
     if(length(uu)==1) return(1)
     d <- dist(mkg$position[uu])
     h <- hclust(d,"ave")
     cutree(h,h=tssTol)
  })
  paste( rep(names(s), sapply(s,length)), unlist(clusts, use.names=FALSE), sep=".")
})

clustno <- unsplit( clustno, lockey )
clustkey <- paste( mkg$chrom, mkg$strand, clustno, sep=":" )

uckey <- unique(clustkey)
mc <- match(uckey, clustkey)
mkg1 <- mkg[mc,]

rownames(mkg1) <- uckey
ids <- split(mkg$newID, clustkey)
collapsedids <- sapply(ids, paste, collapse=";")
mkg1$newID <- ""
mkg1[names(collapsedids),"newID"] <- collapsedids

rownames(mkg1) <- NULL


anno <- data.frame(chr=mkg1$chrom, strand=mkg1$strand, start=mkg1$txStart, end=mkg1$txEnd, name=mkg1$name, allIDs=mkg1$newID)

print(date())