pyLyricFlow/LyricFlow.Core.Backend/LyricFlow.Core/Engine/RhymeEngine.cs

using System.Collections.Generic;
using System.Linq;
using System;
using System.Text.RegularExpressions;
using LyricFlow.Core.Dtos;
using LyricFlow.Core.Services;

namespace LyricFlow.Core.Engine;

public class RhymeEngine
{
    private readonly PhoneticProcessor _processor;
    private readonly WordNetLexicon _wordNet;
    private readonly Dictionary<string, HashSet<string>> _perfectIndex = new();
    private readonly Dictionary<string, HashSet<string>> _slantIndex = new();
    private bool _isIndexed = false;

    public RhymeEngine(PhoneticProcessor processor, WordNetLexicon wordNet)
    {
        _processor = processor;
        _wordNet = wordNet;
    }

    public int CountSyllables(string word)
    {
        var phones = _processor.GetPhonemes(word);
        if (phones.Count > 0)
        {
            return phones[0].Count(p => p.Any(char.IsDigit));
        }

        // Fallback rule-based syllable count
        word = word.ToLower();
        int count = 0;
        string vowels = "aeiouy";
        if (string.IsNullOrEmpty(word)) return 0;
        if (vowels.Contains(word[0])) count++;
        for (int i = 1; i < word.Length; i++)
        {
            if (vowels.Contains(word[i]) && !vowels.Contains(word[i - 1])) count++;
        }
        if (word.EndsWith("e")) count--;
        return Math.Max(1, count);
    }

    public void EnsureIndexed()
    {
        if (_isIndexed) return;

        foreach (var kvp in _processor.Dictionary)
        {
            var word = kvp.Key;
            foreach (var phones in kvp.Value)
            {
                if (phones.Count == 0) continue;

                // Perfect index (last 2 phonemes)
                var suffix = string.Join(" ", phones.Skip(Math.Max(0, phones.Count - 2)));
                if (!_perfectIndex.ContainsKey(suffix)) _perfectIndex[suffix] = new HashSet<string>();
                _perfectIndex[suffix].Add(word);

                // Slant index (last stressed vowel)
                for (int i = phones.Count - 1; i >= 0; i--)
                {
                    if (phones[i].Any(char.IsDigit))
                    {
                        var vowel = phones[i];
                        if (!_slantIndex.ContainsKey(vowel)) _slantIndex[vowel] = new HashSet<string>();
                        _slantIndex[vowel].Add(word);
                        break;
                    }
                }
            }
        }
        _isIndexed = true;
    }

    // MARK: - Calculation Logic
    #region Calculation Logic

    public double CalculateSimilarity(string word1, string word2)
    {
        var phones1 = _processor.GetPhonemes(word1);
        var phones2 = _processor.GetPhonemes(word2);
        if (phones1.Count == 0 || phones2.Count == 0) return 0.0;

        double maxScore = 0.0;
        foreach (var p1 in phones1)
        {
            foreach (var p2 in phones2)
            {
                maxScore = Math.Max(maxScore, PhonMatch(p1, p2));
            }
        }
        return maxScore;
    }

    private double PhonMatch(List<string> first, List<string> second)
    {
        var fRange = first.AsEnumerable().Reverse().ToList();
        var sRange = second.AsEnumerable().Reverse().ToList();
        int limit = Math.Min(fRange.Count, sRange.Count);
        int hits = 0;
        int total = limit;

        for (int i = 0; i < limit; i++)
        {
            if (fRange[i] == sRange[i])
            {
                hits++;
                if (char.IsDigit(fRange[i].Last()))
                {
                    hits++;
                    total++;
                }
            }
            else
            {
                break;
            }
        }

        return total > 0 ? (double)hits / total : 0.0;
    }

    #endregion

    // MARK: - Suggestion Queries
    #region Suggestion Queries

    public SuggestionResponseDto FindSuggestions(string word, int limit = 20)
    {
        EnsureIndexed();
        var normalized = _processor.NormalizeWord(word);
        var phonesList = _processor.GetPhonemes(normalized);
        if (phonesList.Count == 0) return new SuggestionResponseDto([], []);

        var perfect = new HashSet<string>();
        var slant = new HashSet<string>();

        foreach (var phones in phonesList)
        {
            var suffix = string.Join(" ", phones.Skip(Math.Max(0, phones.Count - 2)));
            if (_perfectIndex.TryGetValue(suffix, out var pMatches))
            {
                perfect.UnionWith(pMatches);
            }

            for (int i = phones.Count - 1; i >= 0; i--)
            {
                if (phones[i].Any(char.IsDigit))
                {
                    if (_slantIndex.TryGetValue(phones[i], out var sMatches))
                    {
                        slant.UnionWith(sMatches);
                    }
                    break;
                }
            }
        }

        perfect.Remove(normalized);
        slant.Remove(normalized);
        slant.ExceptWith(perfect);

        return new SuggestionResponseDto(
            perfect.OrderBy(x => x).Take(limit).ToList(),
            slant.OrderBy(x => x).Take(limit).ToList()
        );
    }

    public SynonymResponseDto FindSynonyms(string word, int limit = 15)
    {
        return _wordNet.FindSynonyms(word, limit);
    }

    #endregion

    // MARK: - Rhyme Grouping
    #region Rhyme Grouping

    public List<string[]> GetWordSuffixes(string word)
    {
        return _processor.GetPhonemes(word)
           .Select(phones => phones.Skip(Math.Max(0, phones.Count - 2)).ToArray())
           .ToList();
    }

    public List<RhymeGroupDto> GetRhymeGroups(string text)
    {
        var lines = text.Split('\n');
        var flatWords = new List<(string Orig, string Clean, int Line)>();

        for (int i = 0; i < lines.Length; i++)
        {
            var line = lines[i].Trim();
            if (line.StartsWith("#") || line.StartsWith("@") || line.StartsWith(">")) continue;

            // Remove tags [tag] and find words
            var analysisText = Regex.Replace(line, @"\[.*?\]", "");
            var words = Regex.Matches(analysisText, @"\b\w+\b");
            foreach (Match match in words)
            {
                var clean = _processor.NormalizeWord(match.Value);
                if (!string.IsNullOrEmpty(clean))
                {
                    flatWords.Add((match.Value, clean, i));
                }
            }
        }

        if (flatWords.Count == 0) return [];

        var wordToGroup = new Dictionary<string, int?>();
        var nextGroupId = 0;
        EnsureIndexed();
        for (int i = 0; i < flatWords.Count; i++)
        {
            var current = flatWords[i];
            var suffixes = GetWordSuffixes(current.Clean);
            if (suffixes.Count == 0) continue;

            bool matchFound = false;

            for (int j = Math.Max(0, i - 20); j < i; j++)
            {
                var prev = flatWords[j];
                if (current.Line - prev.Line > 4) continue;
                if (current.Clean == prev.Clean) continue;

                var prevSuffixes = GetWordSuffixes(prev.Clean);
                if (prevSuffixes.Count == 0) continue;

                if (SuffixesOverlap(suffixes, prevSuffixes))
                {
                    if (wordToGroup.TryGetValue(prev.Clean, out var gid) && gid != null)
                    {
                        wordToGroup[current.Clean] = gid;
                        matchFound = true;
                        break;
                    }
                }
            }

            if (!matchFound)
            {
                for (int j = Math.Max(0, i - 20); j < i; j++)
                {
                    var prev = flatWords[j];
                    if (current.Line - prev.Line > 4) continue;
                    if (current.Clean == prev.Clean) continue;

                    var prevSuffixes = GetWordSuffixes(prev.Clean);
                    if (SuffixesOverlap(suffixes, prevSuffixes))
                    {
                        var gid = nextGroupId++;
                        wordToGroup[prev.Clean] = gid;
                        wordToGroup[current.Clean] = gid;
                        break;
                    }
                }
            }
        }

        return flatWords.Select(w => new RhymeGroupDto(w.Clean, wordToGroup.GetValueOrDefault(w.Clean, null))).ToList();
    }

    private bool SuffixesOverlap(List<string[]> first, List<string[]> second)
    {
        foreach (var f in first)
        {
            foreach (var s in second)
            {
                if (f.SequenceEqual(s)) return true;
            }
        }
        return false;
    }

    #endregion
}