# Representación de 'Markov Chain'
import random

# Definicion del graph como vertices
class Vertex:
    def __init__(self, value):  # la palabra será el valor
        self.value = value
        self.adjacent = {}  # nodos que contienen extremos en este vertices
        self.neighbors = []
        self.neighbors_weights = []

    def add_edge_to(self, vertex, weight=0):
        self.adjacent[vertex] = weight

    def increment_edge(self, vertex):
        self.adjacent[vertex] = self.adjacent.get(vertex, 0) + 1

    def get_probability_map(self):
        for (vertex, weight) in self.adjacent.items():
            self.neighbors.append(vertex)
            self.neighbors_weights.append(weight)

    def next_word(self):
        """
        Selecciona proxima palabra basado en 'pesos' (weights)
        """
        return random.choices(self.neighbors, weights=self.neighbors_weights)[0]


class Graph:
    def __init__(self):
        """
        mapeo de strings en Vertex
        """
        self.vertices = {}

    def get_vertex_values(self):
        """
        Retorna todas las posibles palabras
        """
        return set(self.vertices.keys())

    def add_vertex(self, value):
        self.vertices[value] = Vertex(value)

    def get_vertex(self, value):
        if value not in self.vertices:
            self.add_vertex(value)
        return self.vertices[value]     # retorna el objeto Vertex

    def get_next_word(self, current_vertex):
        return self.vertices[current_vertex.value].next_word()

    def generate_probability_mappings(self):
        for vertex in self.vertices.values():
            vertex.get_probability_map()