import os
import string
import random
import re
from graph import Graph, Vertex

def get_words_from_text(text_path):
    with open(text_path, 'r') as f:
        text = f.read()

        # remover [textos en corchetes] regex
        text = re.sub(r'\[(.+)\]', ' ', text)

        text = ' '.join(text.split()) # normalizar espaciado
        text = text.lower()
        # remover puntuaciones
        text = text.translate(str.maketrans('', '', string.punctuation))

    words = text.split()
    return words

def make_graph(words):
    g = Graph()
    previous_word = None
    # por cada palabra, chequear si esta en el grafo, si no, agregarla
    for word in words:
        word_vertex = g.get_vertex(word)
        # si hay una palabra previa, entonces agregar un edge(un borde o extremo)
        # si no existe uno, de otra forma incrementa el weight (peso) en 1
        if previous_word:
            previous_word.increment_edge(word_vertex)
        # setea la palabra a la palabra previa e itera.
        previous_word = word_vertex

    # generear mapeo de probablidades antes de componer
    g.generate_probability_mappings()
    return g

def compose(g, words, length=50):
    composition = []
    word = g.get_vertex(random.choice(words)) # empezar por una palabra aleatoria
    for _ in range(length):
        composition.append(word.value)
        word = g.get_next_word(word)
    return composition

def main(artist):
    # paso 1: obtener palabras del texto
    #words = get_words_from_text('texts/hp_sorcerer_stone.txt')
    words = []
    # para letras de canciones
    for song_file in os.listdir(f'songs/{artist}'):
        song_words = get_words_from_text(f'songs/{artist}/{song_file}')
        words.extend(song_words)

    # paso 2: hacer gráfica(grafo?) usando esas palabras
    g = make_graph(words)

    # paso 3: obtener la próxima palabra para x número de palabras (definidas por el usuario)

    # paso 4: mostrar al usuario
    composition = compose(g, words, 100)
    # retorna un string, todas las palabras separas por un espacio
    return ' '.join(composition)

if __name__ == '__main__':
    print(main('queen'))