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('luuleandmestik.txt', 'a')

#Tunnusenimetuste kirjutamine faili on vajalik vaid programmi esmakasutamisel, andmestikku täiendades tuleb see osa välja toimetada:
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 töökausta ja arvutan samad tunnused välja kõigi sealsete txt-dokumentide kohta
directory = 'luule/Kareva'
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/Kareva/'+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/Kareva/'+filename).read())
		valjund = open('luule/Kareva_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/Kareva_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)+',K,'+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()