import streamlit as st
import io, os, argparse, torch, random
import pytorch_lightning as pl
import numpy as np
from predict import main
from tools.utils import plot_ls

seed = 37
torch.manual_seed(seed)
pl.seed_everything(seed)
torch.manual_seed(seed)
np.random.seed(seed)
random.seed(seed)

st.title('DeepStruc')

st.write('Welcome to DeepStruc that is a Deep Generative Model which has been trained to solve a mono-metallic structure (<200 atoms) based on a PDF!')
st.write('Upload a PDF to use DeepStruc to predict the structure.')


# Define the file upload widget
pdf_file = st.file_uploader("Upload PDF file in .gr format", type=["gr"])

# Define the form to get the other parameters
num_structures = st.number_input("Number of structures to generate", min_value=1, max_value=100, value=10)
structure_index = st.number_input("Index of structure to visualize", min_value=0, value=3)
sigma = st.number_input("Standard deviation for sampling", min_value=0.1, value=3.0)

# Define parser
parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
args = parser.parse_args()
args.num_samples = num_structures
args.index_plot = structure_index
args.sigma = sigma
# Fixed for DeepStruc app
args.model = 'DeepStruc'
args.save_path = './'

# Define the predict button and its behavior
if st.button("Generate structures"):
    if pdf_file is None:
        st.warning("Please upload a PDF file.")
    else:
        # Get the contents of the file as bytes
        file_bytes = pdf_file.read()
    
        # Save the contents of the file to disk
        with open("uploaded_file.gr", "wb") as f:
            f.write(file_bytes)
        
        df, index_highlight, these_cords = main(args)
        
        # Plot the latent space
        fig = plot_ls(df, index_highlight)
        st.pyplot(fig)
        st.write('**The two-dimensional latent space with location of the input.** The size of the points relates to the size of the embedded structure. Each point is coloured after its structure type, FCC (light blue), octahedral (dark grey), decahedral (orange), BCC (green), icosahedral (dark blue), HCP (pink) and SC (red). Each point in the latent space corresponds to a structure based on its simulated PDF. Test data point are plotted on top of the training and validation data, which is made semi-transparent. The latent space locations of the reconstructed structures from the input are shown with black markers and the specific reconstructed structure that is shown in the next box is shown with a black and white marker.')

        # Define the save directory and file name
        file_name = "DeepStruc_prediction.xyz"
        
        # Define a download button to download the file
        def download_button(file_name, button_text):
            with open(file_name, "rb") as f:
                bytes = f.read()
                st.download_button(
                label=button_text,
                data=bytes,
                file_name=file_name,
                mime="text/xyz",)
        
        # Save the coordinates to a file and display a download button
        np.savetxt(file_name, these_cords, fmt="%s")
        download_button(file_name, "Download XYZ file")



st.subheader('Cite')

st.write('If you use DeepStruc, our code or results, please consider citing our papers. Thanks in advance!')

st.write('DeepStruc: Towards structure solution from pair distribution function data using deep generative models **2023** (https://pubs.rsc.org/en/content/articlehtml/2022/dd/d2dd00086e)')
st.write('Characterising the atomic structure of mono-metallic nanoparticles from x-ray scattering data using conditional generative models **2020** (https://par.nsf.gov/biblio/10300745)')

st.subheader('LICENSE')

st.write('This project is licensed under the Apache License Version 2.0, January 2004 - see the LICENSE file at https://github.com/EmilSkaaning/DeepStruc/blob/main/LICENSE.md for details.')
st.write("")

st.subheader('Github')
st.write('https://github.com/EmilSkaaning/DeepStruc')

st.subheader('Questions')
st.write('andy@chem.ku.dk or etsk@chem.ku.dk')