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 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")
#Create two columns with different width
col1, col2 = st.columns( [0.8, 0.2])
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)
    st.subheader("This app takes in your image and styles it with a unique african art.")
    
#Add a header and expander in side bar
st.sidebar.markdown('<p class="font">Afrodreams.AI</p>', unsafe_allow_html=True)
with st.sidebar.expander("About the App"):
     st.write("""
        This app takes in your image and styles it with a unique african art.""")


#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)

#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 = 0#"c"
    params.backend = "mkl"
    params.image_size = 400
    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:
            path = 'stylesv2'
            styles = os.listdir(path)
            params.style_image = path + '/' + random.choice(styles)
            
        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")
                if len(os.listdir('generated_samples')) <= 10:
                    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(button_text="Download Image", data=byte_im, download_filename='afrodreams.jpg', mime="image/png")
        #keeping the current style by update the weight 
        keep_style = st.sidebar.checkbox("Keep current style")