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
andmestik = open('andmestik.txt', 'a')

#Tunnusenimetuste kirjutamine faili
andmestik.write('failinimi,autor,keskm_sonapikkus,2_tahte,3_tahte,4_tahte,5_tahte,6_tahte,7_tahte,8_tahte,9_tahte,10_tahte,a_osak,b_osak,'
	'd_osak,e_osak,f_osak,g_osak,h_osak,i_osak,j_osak,k_osak,l_osak,m_osak,n_osak,o_osak, p_osak,r_osak,s_osak,t_osak,u_osak,v_osak,otilde_osak,'
	'aum_osak,oum_osak,uum_osak,omadussonad,maarsonad,hyydsonad,sidesonad,kaassonad,asesonad,nimisonad,tegusonad,nominaalsus,leksvar\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'):
			#teksti sisselugemine, suurtähtede kaotamine, läbipaistvate sümbolite eemaldamine teksti algusest ja lõpust
			luuletekst = open('luule/'+folder+'/'+filename).read().lower().strip()


			# Tähtede osakaalude arvutamine

			#Määran tähed, mille osakaalu soovin võrrelda
			tahestik = ['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
			tahed = len(re.findall('[\w]', luuletekst))

			#Leian kõigi tähestiku loendi tähtede osakaalu.
			#Kaks loendit on seotud indeksi kaudu, nt a-tähe osakaalu indeks on 0.
			taheosakaalud = []
			for taht in tahestik:
				sagedus = len(re.findall(taht, luuletekst))
				osakaal = sagedus / tahed
				taheosakaalud.append(osakaal)

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

			doc = nlp(open('luule/'+folder+'/'+filename).read())
			valjund = open('luule/'+folder+'_parsitud/'+filename, 'a')
			valjund.write('sona\tlemma\tsonaliik\n')

			#Kuna luuletuses on lausepiire raske määrata, siis pole sõnade arv lauses oluline.
			#Soovin väljundifaili ainult sõnavormi, algvormi ja sõnaliiki.
			for sent in doc.sentences:
				for word in sent.words:
					word.text = word.text.lower()
					margendid = '\n'.join([f'{word.text}\t{word.lemma}\t{word.xpos}'])
					valjund.write(margendid+'\n')

			valjund.close()

			# Sõnapikkuste ja sõnaliikide osakaalude arvutamine

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

			#Keskmine sõnade pikkus luuletuses
			kesksona = luuletus['sonapikkus'].mean()

			#Leian eri pikkusega sõnade osakaalu
			sonapikkused = luuletus.groupby('sonapikkus').sonapikkus.count().to_frame()
			sonapikkused.rename(columns={'sonapikkus':'sagedus'}, inplace=True)
			sonapikkused['sonapikkus'] = sonapikkused.index
			sonapikkused['osakaal'] = sonapikkused['sagedus'] / sonad

			#Talletan kahe- kuni kümnetäheliste sõnade osakaalu loendisse, kus indeksiga 0 on tähistatud kahetähelised sõnad, indeksiga 1 kolmetähelised sõnad jne.
			pikkuseosakaalud = []
			for i in range(2, 11):
				osak = sonapikkused[sonapikkused.sonapikkus == i].osakaal.sum()
				pikkuseosakaalud.append(osak)

			#Leian sõnaliikide osakaalud
			sonaliigid = luuletus.groupby('sonaliik').sonaliik.count().to_frame()
			sonaliigid.rename(columns={'sonaliik':'sagedus'}, inplace=True)
			sonaliigid['sonaliik'] = sonaliigid.index
			sonaliigid['osakaal'] = sonaliigid['sagedus'] / sonad
			
			#Talletan sõnaliikide osakaalud loendisse, kus indeksiga 0 on tähistatud omadussõnad ('A'), indeksiga 1 määrsõnad ('D') jne.
			liigilist = ['A', 'D', 'I', 'J', 'K', 'P', 'S', 'V']
			liigiosakaalud = []
			for liik in liigilist:
				osak = sonaliigid[sonaliigid.sonaliik == liik].osakaal.sum()
				liigiosakaalud.append(osak)
			
			#Leian luuletuse nimisõnalisuse
			SV = sonaliigid[sonaliigid.sonaliik == 'S'].sagedus.sum() / sonaliigid[sonaliigid.sonaliik == 'V'].sagedus.sum()
			
			#Erinevate sõnade ja kõigi tekstisõnade suhtarvu leidmine
			
			lemmad = luuletus.groupby('lemma').lemma.count().to_frame() #andmetabel sõnade grupeerimiseks algvormi e lemma järgi
			lemmad.rename(columns={'lemma':'sagedus'}, inplace=True)
			lemmad['lemma'] = lemmad.index
			erilemmad = lemmad.lemma.count()
			leksvar = erilemmad / sonad
			
			#Salvestan konkreetse teksti andmed andmestikku - NB! Autori nime tähis!
			andmestik.write(str(filename)+','+folder+','+str(kesksona)+','+str(pikkuseosakaalud[0])+','+str(pikkuseosakaalud[1])+','+str(pikkuseosakaalud[2])+','+str(pikkuseosakaalud[3])+','+str(pikkuseosakaalud[4])+
				','+str(pikkuseosakaalud[5])+','+str(pikkuseosakaalud[6])+','+str(pikkuseosakaalud[7])+','+str(pikkuseosakaalud[8])+','+str(taheosakaalud[0])+','+str(taheosakaalud[1])+','+str(taheosakaalud[2])+
				','+str(taheosakaalud[3])+','+str(taheosakaalud[4])+','+str(taheosakaalud[5])+','+str(taheosakaalud[6])+','+str(taheosakaalud[7])+','+str(taheosakaalud[8])+','+str(taheosakaalud[9])+
				','+str(taheosakaalud[10])+','+str(taheosakaalud[11])+','+str(taheosakaalud[12])+','+str(taheosakaalud[13])+','+str(taheosakaalud[14])+','+str(taheosakaalud[15])+','+str(taheosakaalud[16])+
				','+str(taheosakaalud[17])+','+str(taheosakaalud[18])+','+str(taheosakaalud[19])+','+str(taheosakaalud[20])+','+str(taheosakaalud[21])+','+str(taheosakaalud[22])+','+str(taheosakaalud[23])+
				','+str(liigiosakaalud[0])+','+str(liigiosakaalud[1])+','+str(liigiosakaalud[2])+','+str(liigiosakaalud[3])+','+str(liigiosakaalud[4])+','+str(liigiosakaalud[5])+','+str(liigiosakaalud[6])+
				','+str(liigiosakaalud[7])+','+str(SV)+','+str(leksvar)+'\n')

andmestik.close()