import os
import re
import pandas as pd
import stanza
nlp = stanza.Pipeline(lang='et', processors='tokenize,pos,lemma')

#Loon dokumendi, kuhu allpool arvutatud andmed salvestada
data = open('luuleandmestik.txt', 'a')

#Tunnusenimetuste kirjutamine faili
data.write('filename,author,meanWordLength,2_LetterWords,3_LetterWords,4_LetterWords,5_LetterWords,6_LetterWords,7_LetterWords,8_LetterWords,'
	'9_LetterWords,10_LetterWords,a_Freq,b_Freq,d_Freq,e_Freq,f_Freq,g_Freq,h_Freq,i_Freq,j_Freq,k_Freq,l_Freq,m_Freq,n_Freq,o_Freq,p_Freq,'
	'r_Freq,s_Freq,t_Freq,u_Freq,v_Freq,otilde_Freq,aum_Freq,oum_Freq,uum_Freq,adj_Freq,adv_Freq,interj_Freq,conj_Freq,adpos_Freq,pron_Freq,'
	'noun_Freq,verb_Freq,nounVerbRatio,lexVar\n')

#Määran alamkaustad, mille dokumente tahan analüüsida
folders = ['Viiding', 'Ristikivi', 'Talvik', 'Kareva']

for folder in folders:

	#Määran töökausta ja arvutan soovitud tunnused välja kõigi sealsete txt-dokumentide kohta
	directory = 'luule/'+folder
	for filename in os.listdir(directory):
		if filename.endswith('.txt'):
			poemText = open('luule/'+folder+'/'+filename).read().lower().strip()


			# Tähtede osakaalude arvutamine

			#Määran tähed, mille osakaalu soovin võrrelda
			alphabet = ['a', 'b', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 
				'm', 'n', 'o', 'p', 'r', 's', 't', 'u', 'v', 'õ', 'ä', 'ö', 'ü']

			#Leian kõigi tähtede arvu luuletuses
			letterCount = len(re.findall('[\w]', poemText))

			#Leian kõigi tähestiku loendi tähtede osakaalu.
			letterFreqs = []
			for letter in alphabet:
				absfreq = len(re.findall(letter, poemText))
				relFreq = absfreq / letterCount
				letterFreqs.append(relFreq)

			# Teksti märgendamine sõnapiiride, sõnade algvormide ja sõnaliikide tuvastamiseks

			doc = nlp(open('luule/'+folder+'/'+filename).read())
			output = open('luule/'+folder+'_parsitud/'+filename, 'a')
			output.write('word\tlemma\tpos\n')

			#Stanza nummerdab sõnad lausete kaupa, aga luuletuses pole lausepiirid olulised.
			#Kirjutan märgendamise väljundifaili ainult info sõnavormi, algvormi ja sõnaliigi kohta.
			for sent in doc.sentences:
				for word in sent.words:
					word.text = word.text.lower()
					tags = '\n'.join([f'{word.text}\t{word.lemma}\t{word.xpos}'])
					output.write(tags+'\n')

			output.close()

			# Sõnapikkuste ja sõnaliikide osakaalude arvutamine

			poem = pd.read_csv('luule/'+folder+'_parsitud/'+filename, sep = '\t')
			poem = poem[poem.pos != 'Z'] #kirjavahemärkide kõrvalejätmine
			wordCount = poem.word.count() #sõnade arv luuletuses
			poem['wordlength'] = [len(word) for word in poem.word] #sõnapikkuse tulba loomine

			#Keskmine sõnade pikkus luuletuses
			meanWordLength = poem['wordlength'].mean()

			#Leian eri pikkusega sõnade osakaalu
			wordLengths = poem.groupby('wordlength').wordlength.count().to_frame()
			wordLengths.rename(columns={'wordlength':'frequency'}, inplace=True)
			wordLengths['wordlength'] = wordLengths.index
			wordLengths['relfreq'] = wordLengths['frequency'] / wordCount

			#Talletan kahe- kuni kümnetäheliste sõnade osakaalu uude loendisse.
			lengthFreqs = []
			for i in range(2, 11):
				relfreq = wordLengths[wordLengths.wordlength == i].relfreq.sum()
				lengthFreqs.append(relfreq)

			#Leian sõnaliikide osakaalud
			partsOfSpeech = poem.groupby('pos').pos.count().to_frame()
			partsOfSpeech.rename(columns={'pos':'frequency'}, inplace=True)
			partsOfSpeech['pos'] = partsOfSpeech.index
			partsOfSpeech['relfreq'] = partsOfSpeech['frequency'] / wordCount
			
			#Talletan sõnaliikide osakaalud loendisse, mis on indeksi kaudu seotud sõnaliigitähiste loendiga.
			posList = ['A', 'D', 'I', 'J', 'K', 'P', 'S', 'V']
			posFreqs = []
			for item in posList:
				relfreq = partsOfSpeech[partsOfSpeech.pos == item].relfreq.sum()
				posFreqs.append(relfreq)
			
			#Leian luuletuse nimisõnalisuse (näitab sisulist tihedust)
			nounVerbRatio = partsOfSpeech[partsOfSpeech.pos == 'S'].frequency.sum() / partsOfSpeech[partsOfSpeech.pos == 'V'].frequency.sum()
			
			#Erinevate sõnade ja kõigi tekstisõnade suhtarvu leidmine
			
			lemmas = poem.groupby('lemma').lemma.count().to_frame() #andmetabel sõnade grupeerimiseks algvormi e lemma järgi
			lemmas.rename(columns={'lemma':'frequency'}, inplace=True)
			lemmas['lemma'] = lemmas.index
			uniqueLemmas = lemmas.lemma.count()
			lexVar = uniqueLemmas / wordCount
			
			#Salvestan konkreetse teksti andmed andmestikku
			data.write(str(filename)+','+folder+','+str(meanWordLength)+','+str(lengthFreqs[0])+','+str(lengthFreqs[1])+','+str(lengthFreqs[2])+
				','+str(lengthFreqs[3])+','+str(lengthFreqs[4])+','+str(lengthFreqs[5])+','+str(lengthFreqs[6])+','+str(lengthFreqs[7])+
				','+str(lengthFreqs[8])+','+str(letterFreqs[0])+','+str(letterFreqs[1])+','+str(letterFreqs[2])+','+str(letterFreqs[3])+
				','+str(letterFreqs[4])+','+str(letterFreqs[5])+','+str(letterFreqs[6])+','+str(letterFreqs[7])+','+str(letterFreqs[8])+
				','+str(letterFreqs[9])+','+str(letterFreqs[10])+','+str(letterFreqs[11])+','+str(letterFreqs[12])+','+str(letterFreqs[13])+
				','+str(letterFreqs[14])+','+str(letterFreqs[15])+','+str(letterFreqs[16])+','+str(letterFreqs[17])+','+str(letterFreqs[18])+
				','+str(letterFreqs[19])+','+str(letterFreqs[20])+','+str(letterFreqs[21])+','+str(letterFreqs[22])+','+str(letterFreqs[23])+
				','+str(posFreqs[0])+','+str(posFreqs[1])+','+str(posFreqs[2])+','+str(posFreqs[3])+','+str(posFreqs[4])+','+str(posFreqs[5])+
				','+str(posFreqs[6])+','+str(posFreqs[7])+','+str(nounVerbRatio)+','+str(lexVar)+'\n')

data.close()