"""Provide BlockedSongClassifier, which detects certain hard-coded songs in lyrics."""

import re

import numpy as np
import scipy  # type: ignore [import-untyped]
from sklearn.feature_extraction.text import TfidfVectorizer  # type: ignore [import-untyped]

from suno_utils.utils.lyrics import remove_speakers
from suno_utils.utils.text import normalize_whitespace


def _clean_text(text: str) -> str:
    """Preprocess lyrics for TfIdfVectorizer."""
    text = "\n" + text
    text = text.lower()
    text = text.replace("`", "'")
    text = text.replace("'", "")
    text = text.replace('"', "")
    text = re.sub(r"\[.+?\]", " ", text)  # tags
    text = re.sub(r"\n.+?\:", " ", text)  # new line ends with :
    text = re.sub(r"\n.+?\:", " ", text)  # new line ends with :
    text = re.sub(r"\n\(.+?\)", " ", text)  # new line with ()
    text = re.sub(r"[\d]", " ", text)  # digits
    text = re.sub(r"▁", "", text)  # special stuff
    text = remove_speakers(text)
    text = re.sub(r"\(.*?\)", "", text)  # remove parenthetical material
    text = re.sub(r"[^\w\'\s]", " ", text)  # keep only the words
    text = re.sub(r"ing\b", "in", text)  # normalize ..ing to ..in  to deal with lyrical informality
    text = normalize_whitespace(text)
    return text


BLOCKED_SONGS = [
    """
Sleigh bells ring, are you listening? (Doo)
In the lane, snow is glistening
A beautiful sight, we're happy tonight
Walking in a winter wonderland
Gone away is the bluebird
Here to stay is a new bird (Ooo)
He sings a love song as we go along
Walking in a winter wonderland

In the meadow, we can build a snowman
And pretend that he is Parson Brown
He'll say, "Are you married?" We'll say, "No man"
But you can do the job when you're in town

Later on, we'll conspire
As we dream by the fire
To face unafraid, the plans that we've made
Walking in a winter wonderland

Sleigh bells ring, are you listening?
In the lane, snow is glistening
A beautiful sight, we're happy tonight
Walking in a winter wonderland
Gone away is the bluebird
Here to stay is a new bird
He sings a love song as we go along (Ooo)
Walking in a winter wonderland

In the meadow, we can build a snowman
And pretend that he's a circus clown
We'll have lots of fun with Mr. Snowman
Yes, until the other kiddies knock him down

Later on, we'll conspire (Ooo)
As we dream by the fire
To face unafraid all the plans that we've made
Walking in a winter wonderland
Walking in a winter wonderland""",
    """
Chestnuts roasting on an open fire, Jack Frost nipping at your nose
Yuletide carols being sung by a choir and folks dressed up like Eskimos
Everybody knows a turkey and some mistletoe help to make the season bright
Tiny tots with their eyes all aglow will find it hard to sleep tonight


They know that Santa's on his way
He's loaded lots of toys and goodies on his sleigh
And every mother's child is gonna spy
To see if reindeer really know how to fly


And so I'm offering this simple phrase to kids from one to ninety-two
Although it's been said many times, many ways, Merry Christmas to you

And so I'm offering this simple phrase to kids from one to ninety-two
Although it's been said many times, many ways, Merry Christmas to you""",
]

# params eyeballed: perfect classification on random sample of
# 100_000 genius lyrics.
_DEFAULT_THRESHOLD = 0.7

_NGRAM_RANGE = (4, 4)  # using 4-grams results in a very precise classifier


def singleton(cls):
    instances = {}

    def wrapper(*args, **kwargs):
        if cls not in instances:
            instances[cls] = cls(*args, **kwargs)
        return instances[cls]

    return wrapper


class SingletonBase:
    _instances = {}

    def __new__(cls, *args, **kwargs):
        if cls not in cls._instances:
            instance = super().__new__(cls)
            cls._instances[cls] = instance
        return cls._instances[cls]


class BlockedSongClassifier(SingletonBase):
    """Classify songs as blocked with TfIdf + linear classification."""

    def __init__(
        self,
        vectorizer: TfidfVectorizer,
        train_tfidf: scipy.sparse.spmatrix,
        threshold: float = _DEFAULT_THRESHOLD,
    ) -> None:
        self.vectorizer = vectorizer
        self.train_tfidf = train_tfidf
        self.threshold = threshold

    @classmethod
    def from_songs(
        cls,
        songs: list[str],
        threshold: float = _DEFAULT_THRESHOLD,
    ) -> "BlockedSongClassifier":
        """Instantiate a classifier from a list of songs to be blocked."""
        vectorizer = TfidfVectorizer(ngram_range=_NGRAM_RANGE)
        train_tfidf = vectorizer.fit_transform([_clean_text(song) for song in songs])
        return cls(vectorizer, train_tfidf, threshold)

    def _get_score(self, lyrics: str) -> float:
        """Calculate dot product of lyrics with best match in training data."""
        clean_lyrics = _clean_text(lyrics)
        y = self.vectorizer.transform([clean_lyrics])
        best_match_score = np.max(self.train_tfidf @ y.T)
        return best_match_score

    def is_blocked(self, lyrics: str) -> bool:
        """Classify lyrics as blocked or not."""
        return self._get_score(lyrics) > self.threshold


blocked_song_classifier = BlockedSongClassifier.from_songs(BLOCKED_SONGS)