import neural_style

import streamlit as st
import os 
import random
import numpy as np
from  PIL import Image, ImageEnhance
from io import BytesIO
import matplotlib.pyplot as plt
import streamlit_ext as ste #for download button not to rerun
from huggingface_hub import upload_file

HF_TOKEN = os.environ.get("HF_TOKEN")

st.set_page_config(layout="wide")

st.markdown('<p class="font">Afrodreams.AI</p>', unsafe_allow_html=True)
st.subheader("This app takes in your image and styles it with a unique african art.")

#Create two columns with different width
col1, col2 = st.columns( [0.8, 0.2])
import time



with col1:               # To display the header text using css style
    st.markdown(""" <style> .font {
    font-size:35px ; font-family: 'Cooper Black'; color: #FF9633;} 
    </style> """, unsafe_allow_html=True)
    st.markdown('<p class="font">Upload your photo here...</p>', unsafe_allow_html=True)
    
    



#Add file uploader to allow users to upload photos
uploaded_file = st.file_uploader("", type=['jpg','png','jpeg'])

# add slider to side bar
style_weight = st.slider("Select Style Weight", min_value=10, max_value=100, value=12)
img_size_slider= st.select_slider(label= 'Seleet Output Quality Level',
                                         options = ['Very Low', 'Low', 'Normal', 'High', 'Very High'],
                                         value='Normal')
img_size_mapping  = {'Very Low':128, 'Low':300, 'Normal':400, 'High':500, 'Very High':600}


def get_random_subset(list_, num_imgs):
    return random.sample(list_, num_imgs)


def display_random_images(five_rand_imgs, display_type, size= (15, 6)):
    fig = plt.figure(figsize=size)
    fig.subplots_adjust(wspace=0.2) 
    for i in range(1, len(five_rand_imgs)+1):
        ith_image = Image.open(five_rand_imgs[i-1])
    
        ax = fig.add_subplot(1, 5, i)
        ax.imshow(ith_image)
        ax.set_title(f'{display_type} {i}')
        plt.axis('off')
    
    st.pyplot(fig)



path = 'stylesv2'
    
    
#expander for style selection
with st.expander("Expand to select style type"):
    img_names = [os.path.join(path, img) for img in os.listdir(path)]
    five_rand_imgs0 = get_random_subset(img_names, 5)
    if 'selected_image' not in st.session_state:
        st.session_state.selected_image = five_rand_imgs0
    five_rand_imgs = st.session_state.selected_image   
    display_random_images(five_rand_imgs, 'Style')
    chosen_style = st.selectbox(
    'Select the style you want to use',
    options = five_rand_imgs, format_func = lambda x: "Style " + str(five_rand_imgs.index(x) + 1),
    key= 'expander1'
    )
    
 

#put notificaation
#with st.empty():
    #for seconds in range(5):
        #st.info('Please note that by using this app, you agree that your image  be will be showcased on this app.')
        #time.sleep(1)
    #st.empty()

#Add 'before' and 'after' columns
if uploaded_file is not None:
    image = Image.open(uploaded_file)
    
    col1, col2 = st.columns( [0.5, 0.5])
    with col1:
        st.markdown('<p style="text-align: center;">Before</p>',unsafe_allow_html=True)
        st.image(image,width=300)  

    with col2:
        st.markdown('<p style="text-align: center;">After</p>',unsafe_allow_html=True)
        
    # add a button
    run = st.button('Generate Art')
    my_bar = st.progress(0)
    params = neural_style.TransferParams()
    params.gpu = "c" #0
    params.backend = "mkl"

    
    params.image_size = img_size_mapping[img_size_slider]
    
    params.content_image = uploaded_file
    params.style_weight = style_weight



    keep_style = False
    if run==True:
        # run image selection if keep style is false
        if keep_style==False:
            
            styles = os.listdir(path)
            #params.style_image = path + '/' + random.choice(styles)
            params.style_image = chosen_style
            
        st.session_state.submitted = True   
        with st.spinner('Wait for it...'): 
            neural_style.transfer(params)

        #display image when done.
        with col2:
            if 'submitted' in st.session_state:
                result = Image.open('out.png')
                st.image(result, width=300)
                buf = BytesIO()
                result.save(buf, format="png")

                img_file_name = f"generated_samples/{str(len(os.listdir('generated_samples')))}.png"
            
                _ = upload_file(path_or_fileobj = 'out.png',
                        path_in_repo = img_file_name,
                        repo_id='AfrodreamsAI/afrodreams',
                        repo_type='space',
                        token=HF_TOKEN
                    ) 
                    
                byte_im = buf.getvalue()
                run = ste.download_button("Download Image", data=byte_im, file_name="afrodreams.png")
                

                #if run==True: 
# selectiuing random iamges to be displayed
img_names = [os.path.join('generated_samples', img) for img in os.listdir('generated_samples')]
five_rand_imgs1 = get_random_subset(img_names, 5)
st.subheader('\n\n\n\n\n\n\n\n\n Examples of some Generate Images')
display_random_images(five_rand_imgs1, 'Generate image', size=(20, 15))