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


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

COLS = st.columns([0.60, 0.40]) 
SCROLLABLE_TEXT = COLS[1].container(height=500)


def extract_file_name(image_object):
    file_name = image_object.filename
    return file_name


def extract_index(file_name):
    return "flag: todo"


def classify_full_dataset(shosen_dataset_name, chosen_model_name):
    image_count = 0
     
    #modle instance
    classifier_pipeline = pipeline('image-classification', model=chosen_model_name)    
        
    #dataset
    dataset = load_dataset(shosen_dataset_name,"testedata_readme") 
     
    for i in range(len(dataset)):
        SCROLLABLE_TEXT.write("i-1:" + str(i-1))
        image_object = dataset['pasta'][i-1]["image"]
        SCROLLABLE_TEXT.image(image_object, caption="Uploaded Image", width=300)
        
        
        file_name = extract_file_name(image_object)
        SCROLLABLE_TEXT.write(f"file_name: {file_name}")
        image_index = extract_index(file_name)
        SCROLLABLE_TEXT.write(f"image_index: {image_index}")
        
        
        
        image_object_texte_data = dataset['pasta'][i-1]
        SCROLLABLE_TEXT.write(f"image_object_texte_data: {image_object_texte_data}")
        
        
 
        SCROLLABLE_TEXT.write(f"teste1 (image_object._getexif()): {teste1}")            
        
        teste2 = image_object._getexif().items()          
        SCROLLABLE_TEXT.write(f"teste2 (image_object._getexif().items()): {teste2}")
        
        
        
        #extract_metadata = { ExifTags.TAGS[k]: v for k, v in image_object._getexif().items() if k in ExifTags.TAGS }
        #SCROLLABLE_TEXT.write(f"extract_metadata: {extract_metadata}")
        
        #classification
        classification_result = classifier_pipeline(image_object)
        SCROLLABLE_TEXT.write(classification_result)
        #TODO save classification result in dataset
        image_count += 1
        SCROLLABLE_TEXT.write(f"Image count" + str(image_count))
        #SCROLLABLE_TEXT.write(image_count)    
 


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
    if chosen_model_name is not None and shosen_dataset_name is not None:
        if COLS[0].button("Classify images"):
        
            classify_full_dataset(shosen_dataset_name, chosen_model_name)
            COLS[0].write("Classification result {classification_result}")
            COLS[0].write("--- END ---")
            #COLS[0].write(classification_result)
 
                  
if __name__ == "__main__":
    main()