import spaces
import gradio as gr
import numpy as np
import pandas as pd
import time

css="""
#col-left {
    margin: 0 auto;
    max-width: 640px;
}
#col-right {
    margin: 0 auto;
    max-width: 640px;
}
.grid-container {
  display: flex;
  align-items: center;
  justify-content: center;
  gap:10px
}

.image {
  width: 128px; 
  height: 128px; 
  object-fit: cover; 
}

.text {
  font-size: 16px;
}
"""
emotion_columns = ['admiration', 'amusement', 'anger', 'annoyance', 'approval',
                   'caring', 'confusion', 'curiosity', 'desire', 'disappointment',
                   'disapproval', 'disgust', 'embarrassment', 'excitement', 'fear',
                   'gratitude', 'grief', 'joy', 'love', 'nervousness', 'optimism',
                   'pride', 'realization', 'relief', 'remorse', 'sadness', 'surprise','neutral']

def load_text(text_path: str) -> str:
    with open(text_path, 'r', encoding='utf-8') as f:
        text = f.read()

    return text

def select_checkbox(name):
    if name =="All":
        return gr.CheckboxGroup(value=emotion_columns)
    elif name =="None":
        return []
    elif name =="Positive":
        return ["admiration","amusement","approval","caring","curiosity","desire","excitement","gratitude","joy","love","optimism","pride","relief"]
    elif name =="Negative":
        return ["anger","annoyance","disappointment","disapproval","disgust","fear","grief","embarrassment","remorse","sadness"]
    else:
        return ["confusion","nervousness","neutral","realization","surprise"]
def process_datas(checked_emotions,mode="filter",max_data=100,skip_data=0):
    checked_emotions = sorted(checked_emotions)
    
    df = pd.read_parquet("hf://datasets/google-research-datasets/go_emotions/raw/train-00000-of-00001.parquet")
    
     
    def filter_emotions(emotions):
        unchecked = emotion_columns.copy()
        condition_checked = np.all(df[emotions] == 1, axis=1)
        for emotion in checked_emotions:
            unchecked.remove(emotion)
        condition_unchecked = np.all(df[unchecked] == 0, axis=1)
        filtered_df = df[condition_checked & condition_unchecked]
        return filtered_df
    
    def df_to_text(df):
        df = df.iloc[skip_data:]
        if len(filtered_df) == 0:
            return ""
        texts=(df.head(max_data)[['text']].to_string(index=False,max_colwidth=None))
        trimmed_texts = [line.strip() for line in texts.split('\n')[1:] if line.strip()]
        return "\n".join(trimmed_texts)
    
    if mode == "filter":
        filtered_df = filter_emotions(checked_emotions)
        count = (len(filtered_df))
        trimmed_texts = df_to_text(filtered_df)
        
        last_count = min(count,(skip_data+max_data))
        label = f"{skip_data+1} - {last_count} of {count}"

        label_texts = [f"[{emotion}]" for emotion in checked_emotions]

        output_text = "+".join(label_texts)+"\n"+trimmed_texts
        output_label = label
    else:
        max_data = max(1,int(max_data/len(checked_emotions)))
        text_arrays = []
        for emotion in checked_emotions:
            text_arrays.append(f"[{emotion}]")
            filtered_df = filter_emotions([emotion])
            trimmed_texts = df_to_text(filtered_df)
            text_arrays.append(trimmed_texts)
            text_arrays.append("\n")
        print(text_arrays)
        output_text = "\n".join(text_arrays)
        output_label = f"{len(checked_emotions)} x {max_data}"
             

    return output_text,output_label,",".join(checked_emotions)


with gr.Blocks(css=css, elem_id="demo-container") as demo:
    with gr.Column():
        gr.HTML(load_text("demo_header.html"))
        gr.HTML(load_text("demo_tools.html"))
    with gr.Row():
        with gr.Column():
            with gr.Row(equal_height=True):
                mode_group = gr.Radio(choices=["filter","list"],label="Mode",value="filter")
                selection_name = gr.Dropdown(label="Select",choices=["All","None","Positive","Negative","Neutral"],value="All")
                selection_btn= gr.Button("Update Selection")
            checkbox_group = gr.CheckboxGroup(choices=emotion_columns,label="Emotions",value=["love"])
            btn= gr.Button("View Data",variant="primary")
            with gr.Row():
                max_data = gr.Slider(
                                label="Max Data",
                                minimum=0,
                                maximum=540,
                                step=10,
                                value=50,info="returning large data is heavy action,if you need more copy the space")
                skip_data = gr.Slider(
                                label="Skip Data",
                                minimum=0,
                                maximum=100000,
                                step=1,
                                value=0)
        with gr.Column():
            with gr.Row():
                data_size = gr.Textbox(label="Data Count",scale=1)
                checked_size = gr.Textbox(label="Checked",scale=2)

            text_out = gr.TextArea(label="Output", elem_id="output-text")
                    
    btn.click(fn=process_datas, inputs=[checkbox_group,mode_group,max_data,skip_data], outputs =[text_out,data_size,checked_size], api_name='infer')
    selection_btn.click(fn=select_checkbox,inputs=[selection_name],outputs=[checkbox_group])
    gr.Examples(
               examples=[
                    
                    ],
                inputs=[],
                
    )
    gr.HTML(
       gr.HTML(load_text("demo_footer.html"))
    )

    if __name__ == "__main__":
        demo.launch()