from Part1.LyricModel import LyricModel

class SongGenerator:
    def __init__(self, song_structure, chorus_length, verse_length):
        """
        Creates a song from markov lyric model

        :param chorus_length: Number of lines per chorus
        :param verse_length: Number of lines per verse
        :param song_structure: String of C's and V's referring to structure of song where C is chorus and V is verse
        """
        self.chorus_length = chorus_length
        self.verse_length = verse_length
        self.song_structure = [ch.upper() for ch in song_structure if ch.upper() in ['C', 'V']]

    def create_song(self, model: LyricModel):
        print("Creating Lyrics...")
        chorus = self._create_chorus(model)

        lyrics = "\n"
        verse_count = 1
        for section in self.song_structure:
            if section == 'C':
                lyrics += "[CHORUS]\n"
                lyrics += chorus
            else:
                lyrics += f'[VERSE {verse_count}]\n'
                lyrics += self._create_verse(model)
                verse_count += 1
            lyrics += '\n\n'

        return lyrics

    def _create_chorus(self, model):
        return '\n'.join(model.create_lyrics(self.chorus_length))

    def _create_verse(self, model):
        return '\n'.join(model.create_lyrics(self.verse_length))

