import spacy
import pandas as pd
import os
from dotenv import load_dotenv

class ProductNLPExtractor:
    types_db = ["lager", "beer", "can", "ale", "stout", "Wine", "Beer", "Soju & Sake", "Soju", "Sake", "Vodka", "Anejo", "Blanco", "Blended Whisky",
            "Bourbon","Campari","Craft","Extra Anejo","Gin","IPA & Pale Ale","Irish","Japanese","Lager","Mescal","Mezcal","Pilsner","Reposado","Rum","Rye",
            "Scotch","Single Malt","Sloe","Stout & Porter","Wheat Beer","Tequila","Vermouth","Bitters","Absinthe","Liqueur","Fortified Wines",
            "Champagnes & Sparkling Wines","Red Wine","White Wine","Rosé Wine","Cognac","Brandy", "Cabernet Sauvignon", "Merlot", "Cabernet Franc", ]


    def __init__(self):
        load_dotenv()
        nlp_model_location = os.getenv("NLP_MODEL_LOCATION")
        brands_file = os.path.join(nlp_model_location, 'nlp-data/brands.csv')
        self.brands = pd.read_csv(brands_file, delimiter='~')
        
        model_location = os.path.join(nlp_model_location, 'nlp-model')

        self.nlp = spacy.load(model_location)
        brand_patterns = [{"label": "BRAND", "pattern": brand} for brand in self.brands['Brands'].to_list()]
        types_patterns = [{"label": "TYPE", "pattern": type} for type in ProductNLPExtractor.types_db]

        ruler = self.nlp.add_pipe("entity_ruler", before="ner")
        ruler.add_patterns(brand_patterns)
        ruler.add_patterns(types_patterns)

    def extactTagsAndData(self, text):
        return self.nlp(text)