import os
from collections import OrderedDict
os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"

import gradio as gr
from shitsu import ShitsuScorer
from huggingface_hub import hf_hub_download

class OptimizedShitsuScorer:
    def __init__(self, max_models=2):
        self.scorers = OrderedDict()
        self.max_models = max_models
        self.current_language = None

    def get_scorer(self, language):
        if language in self.scorers:
            # Move the accessed language to the end (most recently used)
            self.scorers.move_to_end(language)
        else:
            # If we're at capacity, remove the least recently used model
            if len(self.scorers) >= self.max_models:
                self.scorers.popitem(last=False)
            
            # Load the new model
            self.scorers[language] = ShitsuScorer(language)
        
        self.current_language = language
        return self.scorers[language]

    def score(self, text, language):
        scorer = self.get_scorer(language)
        return scorer.score([text])[0]

    def get_loaded_languages(self):
        return list(self.scorers.keys())

optimized_scorer = OptimizedShitsuScorer(max_models=2)
# Preload English model
optimized_scorer.get_scorer('en')

example_inputs = [
    "The Beatles were a popular band in the 1960s. They released many hit songs.",
    "Chocolate is a type of sweet food that people often eat for dessert.",
    "I'm thinking of going to the beach this weekend. The weather is supposed to be great!",
    "Quantum mechanics is a fundamental theory in physics that provides a description of the physical properties of nature at the scale of atoms and subatomic particles.",
    "Can you believe it's already September? This year is flying by!"
]

def get_score(user_text, language):
    score = optimized_scorer.score(user_text, language)
    formatted_score = f"{score:.4g}"
    loaded_languages = optimized_scorer.get_loaded_languages()
    return f'<div class="nice-box"> Score: {formatted_score}</div>', f"Currently loaded languages: {', '.join(loaded_languages)}"

language_options = ['am', 'ar', 'bg', 'bn', 'cs', 'da', 'de', 'el', 'en', 'es', 'fa', 'fi', 'fr', 'gu', 'ha', 'hi', 'hu', 'id', 'it', 'ja', 'jv', 'kn', 'ko', 'lt', 'mr', 'nl', 'no', 'yo', 'zh']

css = '''
#gen_btn{height: 100%}
#title{text-align: center}
#title h1{font-size: 3em; display:inline-flex; align-items:center}
#title img{width: 100px; margin-right: 0.5em}
#gallery .grid-wrap{height: 10vh}
.card_internal{display: flex;height: 100px;margin-top: .5em}
.card_internal img{margin-right: 1em}
.styler{--form-gap-width: 0px !important}
.nice-box {
    border: 2px solid #007bff;
    border-radius: 10px;
    padding: 15px;
    background-color: #f8f9fa;
    font-size: 18px;
    text-align: center;
    min-height: 60px;
    display: flex;
    align-items: center;
    justify-content: center;
}
'''

theme = gr.themes.Soft(
    primary_hue="blue",
    secondary_hue="sky",
)

with gr.Blocks(theme=theme, css=css) as demo:
    title = gr.HTML(
        """<h1><img src="https://huggingface.co/spaces/Dusduo/shitsu-text-scorer-demo/resolve/main/shitsu-logo.jpeg" alt="LightBlue"> Shitsu Text Scorer</h1>""",
        elem_id="title",
    )
    gr.Markdown(
    """This is a demo of [Shitsu text scorer](https://huggingface.co/lightblue/shitsu_text_scorer) for multiple languages, which scores text based on the amount of useful, textbook-like information in it.
    
    It outputs a score generally between 0 and 1 but can exceed both of these bounds as it is a regressor.
    """
    )
    with gr.Row():
        user_text = gr.Textbox(label='Input text', placeholder='Type something here...')
        language_choice = gr.Dropdown(
            choices=language_options,
            label="Choose a language",
            info="Type to search",
            value="en",
            allow_custom_value=True,
        )
        with gr.Column(scale=0):
            submit_btn = gr.Button("Submit")
            score = gr.HTML(
                value='<div class="nice-box"> Score...  </div>',
                label="Output"
            )
    
    loaded_languages = gr.Markdown("Currently loaded languages: en")
    
    gr.Examples(examples=example_inputs, inputs=user_text)
    
    gr.Markdown(
    """
    This model is based on fasttext embeddings, meaning that it can be used on large amounts of data with limited compute quickly.

    This scorer can be used to filter useful information from large text corpora in many languages.  
    
    This model can also be found on [Github](https://github.com/lightblue-tech/shitsu) and has its own pip installable package. 
    """
    )
    
    submit_btn.click(get_score, inputs=[user_text, language_choice], outputs=[score, loaded_languages])
    
demo.launch()