# -*- coding: utf-8 -*-

'''Kaushik Tandon
   July 21 2018 - August 13 2018
   This program scrapes Wikipedia articles related to Artificial Intelligence. The goal is to build a knowledge graph
   of topics related to AI and categorize them using The Brane's Knowledge Classification System of tags. This program
   successfully creates nodes and links in a CSV file with a false positive rate of less than 10%.
'''

import requests
import csv
from BeautifulSoup import *
import sys
import codecs
import re
import time

reload(sys)
sys.setdefaultencoding('utf8')
sys.stdout = codecs.getwriter('utf8')(sys.stdout)

#Input files
database = 'database.csv' #Entire DB - currently is the August backup
terms_to_collect_file = 'database_terms.csv' #Terms to gather table
avoid_terms_file = 'avoid.txt' #Terms to avoid table
avoid_categories_file = 'avoid_categories.txt' #categories that should not be scraped

#Output files
category_file = 'categories.txt' #Generated file of categories scraped
created_CSV_file = 'scrape.csv' #Created file after scraping

#Create a represesentation of a category in order to easily store and access important data
class Category:
	def __init__(self,name,url,sub_categories,linked_pages,level,category_num):
		self.name = name
		self.url = url;
		self.sub_categories = sub_categories
		self.linked_pages = linked_pages;
		self.level = level
		self.category_num = category_num

def getURL():
	''' Ask the user for which category url to scrape - default is AI category
		Returns:
			URL as string
	'''
	url = raw_input("Enter URL of wikipedia category page to scrape (or enter for default)")
	if len(url) < 2 or 'wikipedia' not in url:
		url = 'https://en.wikipedia.org/wiki/Category:Artificial_intelligence'
		#url = 'https://en.wikipedia.org/wiki/Category:Game_artificial_intelligence'
	return url
def getLinksFromCategoryPage(page):
	''' Given a category page, this method can extract the pages and the subcategories on each page
		Example: https://en.wikipedia.org/wiki/Category:Artificial_intelligence should return 2 arrays, one with 326 pages and one with 37 category titles
		
		Args:
			page: The URL of the page to extract pages/subcategories from		
		Returns:
			Two arrays - one with list of page urls, one with list of category urls
	'''
	#Page must be of form Category:Name
	pages = []
	sub_categories = []
	soup_html = getHTML(page)
	#Extract pages
	a =  soup_html.findAll('div',{'class': 'mw-category-group'})
	for temp in a:
		pageNames = extractPageNames(temp)
		for pageName in pageNames:
			ind1 = pageName.find('(')
			ind2 = pageName.find('P')
			ind3 = pageName.find("C")
			ind4 = pageName.find(')')
			num = bool(re.search(r'\d', pageName)) #Number in pageName
			#Trying to catch pages of type (5 C, 40 P)
			if(num and ind1 >= 0 and ind4 > 0 and (ind2 > 0 or ind3 > 0)):
				continue
			#Remove weird characters
			pageName = ''.join([i if ord(i) < 128 else '' for i in pageName])
			if(len(pageName) > 0):
				pages.append('https://en.wikipedia.org/wiki/' + str(pageName.strip()))
	#Check for additional pages
	c = soup_html.find('div',{'id': 'mw-pages'})
	if(c != None and len(c.findAll('a')) > 2):
		elemToCheck = c.findAll('a')[1]
		if(str(elemToCheck.text).strip().lower() == 'next page'):
			more_page = 'https://en.wikipedia.org' + str(elemToCheck.get('href'))
			additional_pages = extractAdditionalPages(more_page)
			for page in additional_pages:
				pages.append(page)
	#Look for subcategories
	b = soup_html.findAll('a',{'class': 'CategoryTreeLabel  CategoryTreeLabelNs14 CategoryTreeLabelCategory'})
	for sub in b:
		sub = str(sub)
		index = sub.find("Category:")
		name = sub[index:sub.find('"',index)]
		sub_categories.append('https://en.wikipedia.org/wiki/' + name.strip())
	return pages,sub_categories
def extractAdditionalPages(page):
	''' Helper method for getLinksFromCategoryPage() to handle pages which have more than 200 pages linked
		Example url: https://en.wikipedia.org/w/index.php?title=Category:Artificial_intelligence&pagefrom=Leaf+Project%0AThe+Leaf+%28AI%29+Project#mw-pages
		Args:
			page: URL of 'next page' category page being scraped
		Returns:
			List of urls of pages in category on specific page
	'''
	additional_pages = list()
	soup_html = getHTML(page)
	a =  soup_html.findAll('div',{'class': 'mw-category-group'})
	for temp in a:
		pageNames = extractPageNames(temp)
		for pageName in pageNames:
			#Trying to catch pages of type (5 C, 40 P) and ignore them
			ind1 = pageName.find('(')
			ind2 = pageName.find('P')
			ind3 = pageName.find("C")
			ind4 = pageName.find(')')
			num = bool(re.search(r'\d', pageName))
			if(num and ind1 >= 0 and ind4 > 0 and (ind2 > 0 or ind3 > 0)):
				continue
			pageName = ''.join([i if ord(i) < 128 else '' for i in pageName])
			if(len(pageName) > 0):
				additional_pages.append('https://en.wikipedia.org/wiki/' + str(pageName.strip()))
	return additional_pages
def getHTML(url):
	''' Uses BeautifulSoup to get the HTML for a page
		Args:
			url: URL of page to get HTML for
		Returns:
			Beautiful Soup object with HTML
	'''
	try:
		r = requests.get(url)
		return BeautifulSoup(r.text)
	except:
		print("Couldn't get HTML for: " + url)

def getTitle(soup_html):
	''' Uses BeautifulSoup html to get the title of the article
		Args:
			soup_html: Beautiful Soup object with HTML (returned by getHTML())
		Returns:
			title of article or "Error"
	'''	
	if(len(soup_html.findAll("h1", {"id": "firstHeading"})) == 0):
		return "Error"
	txt = soup_html.findAll("h1", {"id": "firstHeading"})[0].getText()
	txt = ''.join([i if ord(i) < 128 else '-' for i in txt])
	return txt

def extractSeeAlso(soup_html):
	''' Uses BeautifulSoup html to get the see also categories from a Wikipedia article
		Args:
			soup_html: Beautiful Soup object with HTML (returned by getHTML())
		Returns:
			list of names of articles in the see also category (or empty list)
	'''
	seeAlso = list()
	section = soup_html.find('span', id='See_also')
	if section != None:
		wrongUL = True
		section = section.parent.findNext('ul')
		count = 0
		while(wrongUL):
			count = count + 1
			for litag in section.findAll('a', href=True):
				if litag.get('href') == None and wrongUL:
					continue
				elif 'wiki' not in str(litag.get('href')) and wrongUL:
					continue
				else:
					wrongUL = False
					name = litag.text
					if name == None:
						continue
					name = str(name).strip()

					if('page does not exist' in name):
						continue
					if name in seeAlso:
						continue
					else:
						seeAlso.append(name)
			if(wrongUL):
				section = section.parent.findNext('ul')
			if(count == 5):
				break
	return seeAlso

def extractCategories(soup_html):
	''' Uses BeautifulSoup html to get the categories of a Wikipedia article
		Args:
			soup_html: Beautiful Soup object with HTML (returned by getHTML())
		Returns:
			list of names of categories (or empty list)
	'''
	categories = []
	a = soup_html.find('div',{'class': 'mw-normal-catlinks'})
	if a != None:
		for litag in a.findAll('li'):
			categories.append(str(litag.text))
	return categories

def extractReferences(soup_html):
	''' Uses BeautifulSoup html to get the references of a Wikipedia article
		Args:
			soup_html: Beautiful Soup object with HTML (returned by getHTML())
		Returns:
			list of names of references (or empty list)
	'''	
	references = []
	a = soup_html.find('ol',{'class': 'references'})
	if a != None:
		for litag in a.findAll('li'):
			references.append(str(litag.text))
	return references

def isStub(soup_html):
	''' Uses BeautifulSoup html to determine whether Wikipedia article is a stub
		Args:
			soup_html: Beautiful Soup object with HTML (returned by getHTML())
		Returns:
			True if article is a stub (should be skipped)
	'''
	a = soup_html.find('table',{'class': 'metadata plainlinks stub'})
	if a != None:
		return True
	return False

def loadAvoidTerms():
	''' Load the terms to avoid table from a predefined text file
		Returns:
			list of terms to avoid (lowercase)
	'''
	with open(avoid_terms_file) as f:
		content = f.readlines()
	content = [x.strip().lower() for x in content] 
	return content

def loadAvoidCategories():
	''' Load the categories to avoid from a predefined text file
		Returns:
			list of categories to avoid (lowercase)
	'''
	with open(avoid_categories_file) as f:
		content = f.readlines()
	content = [x.strip().lower() for x in content] 
	return content

def loadGatherTerms():
	''' Load the terms to gather table from a predefined csv file
		Returns:
			dictionary with key being lower case word (and plural versions) and value being the database ID
	'''	
	terms = dict()
	with open(terms_to_collect_file) as f:
		for line in f.readlines():
			words = line.split(',')
			temp_id = str(words[0])
			for word in words:
				word = word.strip().lower()
				word2 = word + 's'
				word3 = word + 'es'
				if(len(word) > 0 and (word[0] < '0' or word[0] > '9')):
					terms[word] = temp_id
					terms[word2] = temp_id
					terms[word3] = temp_id
	return terms
def splitAndLower(words):
	''' Returns a list of words ignoring parentheses and splitting to lowercase
		Args:
			words: list of words to handle
		Returns:
			list of words where each word is lowercase
	'''
	words = words.replace("("," ")
	words = words.replace(")"," ")
	words = words.lower().strip()
	all_words = words.split(" ")
	for i in range(len(all_words)):
		all_words[i] = all_words[i].strip()
	return all_words
def numCapitalsInTitle(title):
	''' Determine the number of capitals in an article title
		Args:
			title: article title to check
		Returns:
			integer number of capitals in the title
	'''
	title = title.replace("("," ").replace(")"," ").strip()
	all_words = title.split(" ")
	numCap = 0
	for word in all_words:
		if len(word) > 0 and word[0].isupper():
			numCap = numCap + 1
	return numCap

def validArticleTitle(article_title,avoid_terms,gather_terms):
	''' Determine if a Wikipedia article title is valid, or if the article should be skipped
		Args:
			article_title: title to check
			avoid_terms: list of terms to avoid
			gather_terms: dict of terms to gather
		Returns:
			True if article title is valid
	'''
	#check for partial match
	words_in_title = splitAndLower(article_title)
	for word in words_in_title:
		if word.lower() in avoid_terms:
			return False
	#All individual words cannot be capital
	allCapital = True
	individualWords = article_title.split(" ")
	if(len(individualWords) > 1):
		for word in individualWords:
			word = word.strip()
			if word[0].islower():
				allCapital = False
			if word.lower() in gather_terms or word[:-1].lower() in gather_terms or (word+'s').lower() in gather_terms or (word+'es').lower() in gather_terms:
				return True
			if word[0] == '(' and word[1].islower():
				allCapital = False
		if(allCapital):
			return False
	#avoid_terms only contains lower case, so convert article_title to lower case for checking
	article_title = article_title.lower()
	#check for full title
	if article_title in avoid_terms:
		return False
	#check for plural title
	if article_title + 's' in avoid_terms or article_title + 'es' in avoid_terms:
		return False
	return True

def validCategoryName(name,invalidNames):
	''' Determines whether a category name is invalid and should be scraped
		Args:
			name: string name of category
			invalidNames: list of invalid names loaded from file