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'

# COL1, COL2 = st.columns([3, 1])
# CONTAINER_TOP = st.container()
# CONTAINER_BODY = st.container()
# CONTAINER_FULL = st.container()
# CONTAINER_LOOP = st.container()
COL1=""        
COL2="" 
COLS = st.columns([3, 1]) 
CONTAINER_TOP = st.container() 
CONTAINER_BODY = st.container()
CONTAINER_FULL = st.container() 
CONTAINER_LOOP = st.container() 



#(image_object, classifier_pipeline)
#def classify_one_image(classifier_model, dataset_to_classify):
#classify_one_image(image_object, classifier_pipeline)
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 CONTAINER_LOOP:
        #Image teste load 
        image_object = dataset['pasta'][0]["image"]
        st.image(image_object, caption="Uploaded Image", width=300)
        st.write("### FLAG 3")
        
    #modle instance
    classifier_pipeline = pipeline('image-classification', model=chosen_model_name)
    CONTAINER_LOOP.write("### FLAG 4")
    
    #classification
    classification_result = classifier_pipeline(image_object)
    CONTAINER_LOOP.write(classification_result)
    CONTAINER_LOOP.write("### FLAG 5")
    #classification_array.append(classification_result)
    
    #save classification
    
    image_count += 1
    CONTAINER_LOOP.write(f"Image count: {image_count}")
    
    return image_count


def make_template():
    
    
    
    with CONTAINER_FULL:
        CONTAINER_TOP
        CONTAINER_BODY
        with CONTAINER_BODY:  
            #COL1, COL2 = st.columns([3, 1])
            with COLS[1]:
                CONTAINER_LOOP


def main():
    
    make_template()
    
    with CONTAINER_TOP:
        st.write("# Bulk Image Classification DEMO")

    
    # TODO Restart or reset your app
    # if st.button("Restart"):
    # # Code to restart or reset your app goes here
        #     import subprocess
        #     subprocess.call(["shutdown", "-r", "-t", "0"])
        
    with CONTAINER_BODY:

        with COL1:
            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 = st.selectbox("Select the model to use",  MODELS, index=0)
        if chosen_model_name is not None:
            COL1.st.write("You selected", chosen_model_name) 

        #Dataset
        shosen_dataset_name = st.selectbox("Select the dataset to use",  DATASETS, index=0)
        if shosen_dataset_name is not None:
            COL1.st.write("You selected", 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 COL1.button("Classify images"):
            
                #classification_array =[]
                classification_result = classify_full_dataset(shosen_dataset_name, chosen_model_name)
                CONTAINER_LOOP.write(f"Classification result: {classification_result}")
                #classification_array.append(classification_result)
                #st.write("# FLAG 6")
                #st.write(classification_array)   

if __name__ == "__main__":
    main()