import random


label = input("Please choose your label between pop and traditional:\n")
ngram_list = ["1gram", "2gram", "3gram"]
ngram_type = ""
while(ngram_type not in ngram_list):
    ngram_type = input("Please Enter ngram type between 1gram-2gram-3gram:\n")
seed = input("Now Enter seed number:\n")
n = input("How many sentences you like to be generated by this ngram?\n")

seed = 9

def unigram_generator(path, destPath):
    unigram_text = ''

    txt_file = open(path, 'r', encoding="utf-8")
    for line in txt_file.readlines():
        line = line.rstrip()
        line += '|'
        unigram_text += line
    unigram_list = unigram_text.split("|")
    unigram_dict = {}
    for i in range(0, len(unigram_list)-1, 2):
        unigram_dict[unigram_list[i]] = float(unigram_list[i+1])



    txt_file = open(destPath, 'w', encoding="utf-8")

    for sentences in range(int(n)):
        counter = 0
        r_words_number = 1
        while(counter <= r_words_number):
            random_probability = random.random()
            temp = str(random_probability)
            r_words_number = int(temp[2])
            if r_words_number == 0:
                r_words_number = 6
            # print(random_probability)
            sum = 0
            for i in unigram_dict:
                sum += unigram_dict[i]
                if sum > random_probability:
                    # print(unigram_dict[i])
                    txt_file.write(i+" ")
                    break
            counter += 1
        txt_file.write("\n")


def biagram_generator(path, destPath):
    unigram_text = ''
    txt_file = open(path, 'r', encoding="utf-8")
    for line in txt_file.readlines():
        line = line.rstrip()
        line += '|'
        unigram_text += line
    unigram_list = unigram_text.split("|")
    unigram_dict = {}
    # print(unigram_list)
    for i in range(0, len(unigram_list)-1, 3):
        unigram_dict[(unigram_list[i], unigram_list[i+1])] = float(unigram_list[i+2])

    # print(unigram_dict)


    txt_file = open(destPath, 'w', encoding="utf-8")

    for sentences in range(int(n)):
        counter = 0
        end = " "
        flag = True
        while(end != "</s>"):
            random_probability = random.random()
            # print(random_probability)
            sum = 0
            for i in unigram_dict:
                sum += unigram_dict[i]
                if sum > random_probability:
                    if flag == False:
                        if (i[0] != "<s>" and i[0] == end) or counter >= 30:
                            counter = 0
                            # print(i)
                            txt_file.write(i[1] +" ")
                            end = i[1]
                            break
    
                    else:
                        if i[0] == "<s>":
                            # print(unigram_dict[i])
                            txt_file.write(i[0]+" "+i[1] +" ")
                            end = i[1]
                            flag = False
                            break
                        

            counter += 1
        txt_file.write("\n")


def trigram_generator(path, destPath):
    unigram_text = ''
    txt_file = open(path, 'r', encoding="utf-8")
    for line in txt_file.readlines():
        line = line.rstrip()
        line += '|'
        unigram_text += line
    unigram_list = unigram_text.split("|")
    unigram_dict = {}
    # print(unigram_list)
    for i in range(0, len(unigram_list)-1, 4):
        unigram_dict[(unigram_list[i], unigram_list[i+1], unigram_list[i+2])] = float(unigram_list[i+3])

    # print(unigram_dict)


    txt_file = open(destPath, 'w', encoding="utf-8")

    for sentences in range(int(n)):
        end = " "
        flag = True
        counter = 0
        while(end != "</s>"):
            random_probability = random.random()
            sum = 0
            for i in unigram_dict:          
                sum += unigram_dict[i]
                if sum > random_probability:
                    if flag == False:
                        if (i[0] != "<s>" and i[0] == end) or counter >len(unigram_dict)/2:
                            counter = 0
                            txt_file.write(i[1]+ " "+ i[2]+" ")
                            end = i[2]
                            break
                    else:
                        if i[0] == "<s>":
                            txt_file.write(i[0] + " "+ i[1]+ " "+ i[2]+" ")
                            end = i[2]
                            flag = False
                            break
            counter += 1
        txt_file.write("\n")

if label == 'pop':
    if ngram_type == "1gram":
        unigram_generator("../../Model/pop/1gram.lm", "../pop.1gram.gen")
    elif ngram_type == "2gram":
        biagram_generator("../../Model/pop/2gram.lm", "../pop.2gram.gen")
    elif ngram_type == "3gram":
        trigram_generator("../../Model/pop/3gram.lm", "../pop.3gram.gen")
elif label == 'traditional':
    if ngram_type == "1gram":
        unigram_generator("../../Model/pop/1gram.lm", "../traditional.1gram.gen")
    elif ngram_type == "2gram":
        biagram_generator("../../Model/pop/2gram.lm", "../traditional.2gram.gen")
    elif ngram_type == "3gram":
        trigram_generator("../../Model/pop/3gram.lm", "../traditional.3gram.gen")


