import streamlit as st
from transformers import pipeline
from PIL import Image
from datasets import load_dataset, Image, list_datasets
from PIL import Image

MODELS = [
            "google/vit-base-patch16-224", #Classifição geral
            "nateraw/vit-age-classifier" #Classifição de idade
]
DATASETS = [
            "Nunt/testedata",
            "Nunt/backup_leonardo_2024-02-01" 
]
MAX_N_LABELS = 5
SPLIT_TO_CLASSIFY = 'pasta'

COLS = st.columns([0.75, 0.25]) 
SCROLLABLE_TEXT = COLS[1].text_area("Conteúdo da segunda coluna", height=500)



def classify_one_image(classifier_model, dataset_to_classify):
       
       
          
    #image_object = dataset[SPLIT_TO_CLASSIFY][i]["image"]
    #st.image(image_object, caption="Uploaded Image", width=300)

       
    #for i in range(len(dataset_to_classify)): 
    #for image in dataset_to_classify:
        #image_object = dataset[SPLIT_TO_CLASSIFY][i]["image"]
        #st.image(image_object, caption="Uploaded Image", width=300)
        
        #st.write(f"Image classification: ", image['file'])

        # image_path = image['file']
        # img = Image.open(image_path)
        # st.image(img, caption="Original image", use_column_width=True)
        # results = classifier(image_path, top_k=MAX_N_LABELS)
        # st.write(results)
        # st.write("----")
    
    return "done" 



def classify_full_dataset(shosen_dataset_name, chosen_model_name):
    image_count = 0
     
    #dataset
    dataset = load_dataset(shosen_dataset_name,"testedata_readme")
    #with SCROLLABLE_TEXT:
    #Image teste load 
    image_object = SCROLLABLE_TEXT.dataset['pasta'][0]["image"]
    SCROLLABLE_TEXT.image(image_object, caption="Uploaded Image", width=300)
    #st.write("### FLAG 3")
        
    #modle instance
    classifier_pipeline = pipeline('image-classification', model=chosen_model_name)
    #COLS[1].write("### FLAG 4")
    
    #classification
    classification_result = SCROLLABLE_TEXT.classifier_pipeline(image_object)
    SCROLLABLE_TEXT.write(classification_result)
    #COLS[1].write("### FLAG 5")
    #classification_array.append(classification_result)
    
    #save classification
    
    image_count += 1
    SCROLLABLE_TEXT.write("Image count")
    SCROLLABLE_TEXT.write(image_count)  
    return image_count


def make_template():
    
    tile = CONTAINER_TOP.title(":balloon:")
    tile.title(":balloon:")
    
    with CONTAINER_FULL:
        CONTAINER_TOP.title("titulo de teste dentro do container CONTAINER_TOP")
        with CONTAINER_BODY:  
            #COL1, COL2 = st.columns([3, 1])
            with COLS[1]:
                CONTAINER_LOOP.write("### OUTPUT")


def main():
    
     
    COLS[0].write("# Bulk Image Classification App")
    
      
    #with CONTAINER_BODY:
    with COLS[0]:
        st.markdown("This app uses several 🤗 models to classify images stored in 🤗 datasets.")
        st.write("Soon we will have a dataset template")
        
    #Model
    chosen_model_name = COLS[0].selectbox("Select the model to use",  MODELS, index=0)
    if chosen_model_name is not None:
        COLS[0].write("You selected")
        COLS[0].write(chosen_model_name)
        
    #Dataset
    shosen_dataset_name = COLS[0].selectbox("Select the dataset to use",  DATASETS, index=0)
    if shosen_dataset_name is not None:
        COLS[0].write("You selected")
        COLS[0].write(shosen_dataset_name)
        
    #click to classify
    #image_object = dataset['pasta'][0]  
    if chosen_model_name is not None and shosen_dataset_name is not None:
        if COLS[0].button("Classify images"):
        
            #classification_array =[]
            classification_result = classify_full_dataset(shosen_dataset_name, chosen_model_name)
            COLS[0].write("Classification result {classification_result}")
            COLS[0].write(classification_result)
            #classification_array.append(classification_result)
            #st.write("# FLAG 6")
            #st.write(classification_array) 
                  
if __name__ == "__main__":
    main()