import pronouncing
import string
import itertools
import gradio as gr

CONSONANTS = set(string.ascii_uppercase) - set('AEIOU')

def get_phones(word):
    phones_for_word = pronouncing.phones_for_word(word)
    if not phones_for_word:
        return None
    return phones_for_word[0].split()

def _get_rhyming_tail(phones):
    vowels = [phone for phone in phones if phone[0] in 'AEIOU']
    syllable_count = len(vowels)
    if syllable_count < 1:
        return None
    return phones[-syllable_count * 2:]

def get_exact_rhymes(phones):
    rhyming_tail = _get_rhyming_tail(phones)
    if not rhyming_tail:
        return []
    
    rhyming_tail_str = " ".join(rhyming_tail)
    matches = pronouncing.search(rhyming_tail_str + "$")
    # Only include common, recognized words
    exact_rhymes = [match for match in matches if match.isalpha() and match.islower()]
    return exact_rhymes

def find_rhymes_for_phrase(phrase):
    words = phrase.split()
    rhyming_options = []

    for word in words:
        phones = get_phones(word)
        if phones is None:
            rhyming_options.append([[f"{word} (Not recognized)"]])
            continue
        
        exact_rhymes = get_exact_rhymes(phones)
        
        if exact_rhymes:
            rhyming_options.append([[rhyme] for rhyme in exact_rhymes])
        else:
            rhyming_options.append([[f"{word} (No rhymes found)"]])
    
    # Flatten each combination into strings
    combined_results = list(itertools.product(*rhyming_options))
    unique_results = set(" ".join(item[0] for item in combination) for combination in combined_results)
    
    # Return a list of unique rhyming combinations, each wrapped in a list for DataFrame compatibility
    return [[rhyme] for rhyme in unique_results if rhyme.isalpha() and rhyme.islower()]

# Function to update the notepad with a selected rhyme
def add_to_notepad(notepad, selected_rhyme):
    return notepad + "\n" + selected_rhyme

# Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("## Interactive Rhyming Generator with Notepad")
    
    # Inputs for the rhyming generator
    phrase_input = gr.Textbox(label="Enter phrase (space-separated words)")
    
    # Output areas
    rhyme_output = gr.Dataframe(headers=["Rhyming Phrases"], interactive=False)
    notepad = gr.Textbox(label="Notepad", lines=10, placeholder="Write your lyrics here...")
    
    # Button for generating rhymes
    generate_btn = gr.Button("Generate Rhymes")
    
    # Display rhyming results and update notepad on button click
    generate_btn.click(find_rhymes_for_phrase, inputs=[phrase_input], outputs=rhyme_output)
    
    # Functionality to select a rhyme and add it to the notepad
    rhyme_output.select(add_to_notepad, inputs=[notepad, rhyme_output], outputs=notepad)

demo.launch()