
library(GetoptLong)

options("timeout" = 99999999)

## https://ftp.uniprot.org/pub/databases/uniprot/current_release/knowledgebase/complete/docs/keywlist.txt

ln = readLines(url("https://ftp.uniprot.org/pub/databases/uniprot/current_release/knowledgebase/complete/docs/keywlist.txt"))
i1 = grep("_____________", ln)
i2 = grep("^//", ln)

ln = ln[seq(min(i1)+1, max(i2))]


parse_kw = function(x) {
	x2 = strsplit(x, "   ")
	lt = list()

	CA = NULL
	for(i in seq_along(x2)) {
		code = x2[[i]][1]

		if(code == "CA") next

		if(code == "HI") {
			CA = c(CA, gsub("^(\\w.*?):.*$", "\\1", x2[[i]][2]))
		}

		if(is.null(lt[[code]])) {
			lt[[code]] = x2[[i]][2]
		} else {
			if(code == "DE") {
				lt[[code]] = paste0(lt[[code]], " ", x2[[i]][2])
			} else {
				lt[[code]] = c(lt[[code]], x2[[i]][2])
			}
		}
	}
	if(!is.null(CA)) {
		lt$CA = unique(CA)
	}
	for(nm in names(lt)) {
		if(!nm %in% c("DE")) {
			lt[[nm]] = gsub("(;|\\.)$", "", lt[[nm]])
		}
	}
	if("IC" %in% names(lt)) {
		names(lt)[names(lt) == "IC"] = "ID"
		lt$ID = gsub("\\.$", "", lt$ID)
	}

	name_map = c("ID" = "Identifier",
		         "AC" = "Accession",
		         "DE" = "Description",
		         "SY" = "Synonyms",
		         "GO" = "Gene_ontology",
		         "HI" = "Hierarchy",
		         "WW" = "WWW_site",
		         "CA" = "Category")
	names(lt) = name_map[names(lt)]
	lt
}


k = which(ln == "//")
k = c(0, k)

kw_terms = NULL
for(i in seq_along(k)[-length(k)]) {
	x = ln[seq(k[i]+1, k[i+1] - 1)]

	kw_terms[[i]] = parse_kw(x)
}

names(kw_terms) = sapply(kw_terms, function(x) x$Identifier)


extract_parents = function(hi) {
	if(is.null(hi)) {
		return(NULL)
	}

	x = gsub(":", ";", hi)
	x2 = lapply(strsplit(x, "; "), function(v) {
		v[seq(length(v)-1, length(v))]
	})

	unique(do.call(rbind, x2))
}


keyword_relations = do.call(rbind, lapply(kw_terms, function(x) extract_parents(x$Hierarchy)))
colnames(keyword_relations) = c("parent", "child")

kw_parents = split(keyword_relations[, 1], keyword_relations[, 2])
kw_children = split(keyword_relations[, 2], keyword_relations[, 1])


library(igraph)

g = graph.edgelist(keyword_relations)

# upstream
cm = distances(g, mode = "in")
kw_ancestors = lapply(1:nrow(cm), function(i) {
	x = cm[i, ]
	l = is.finite(x)
	V(g)$name[setdiff(which(l), i)]
})
names(kw_ancestors) = V(g)$name
kw_ancestors = kw_ancestors[sapply(kw_ancestors, length) > 0]


# downstream
cm = distances(g, mode = "out")
kw_offspring = lapply(1:nrow(cm), function(i) {
	x = cm[i, ]
	l = is.finite(x)
	V(g)$name[setdiff(which(l), i)]
})
names(kw_offspring) = V(g)$name
kw_offspring = kw_offspring[sapply(kw_offspring, length) > 0]


save(kw_parents, file = "../../../data/kw_parents.rda", compress = "xz")
save(kw_children, file = "../../../data/kw_children.rda", compress = "xz")
save(kw_ancestors, file = "../../../data/kw_ancestors.rda", compress = "xz")
save(kw_offspring, file = "../../../data/kw_offspring.rda", compress = "xz")
save(kw_terms, file = "../../../data/kw_terms.rda", compress = "xz")

### annotation to genes
## uniprot_keywords.txt is generated by extract.pl

if(file.exists("uniprot_sprot.dat.gz")) file.remove("uniprot_sprot.dat.gz")
download.file("https://ftp.uniprot.org/pub/databases/uniprot/current_release/knowledgebase/complete/uniprot_sprot.dat.gz", dest = "uniprot_sprot.dat.gz")
system("perl extract.pl")

tb = read.table("uniprot_keywords.txt", sep = "\t", header = T)

format_geneset = function(tb2) {
	g = tb2$Entrez_Gene_ID
	g = as.character(g)
	keywords = tb2$Keywords
	keywords = gsub("\\.$", "", keywords)
	keywords = strsplit(keywords, "; ")

	g = rep(g, times = sapply(keywords, length))
	keywords = unlist(keywords)

	gl = split(g, keywords)

	for(nm in names(gl)) {
		offspring = kw_offspring[[nm]]
		if(!is.null(offspring)) {
			gl[[nm]] = unique(unlist(gl[ offspring ]))
		}
	}

	gl

}


for(t in unique(tb$Taxon_ID)) {
	l = tb$Taxon_ID == t
	if(sum(l) >= 1000) {
		tb2 = tb[l, ]
		lt = format_geneset(tb2)
		saveRDS(lt, file = qq("../keywords_anno/keywords_@{t}_gene_annotation.rds"), compress = "xz")
	}
}


# if(file.exists("uniprot_sprot.dat.gz")) file.remove("uniprot_sprot.dat.gz")




