import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.impute import SimpleImputer
from sklearn.preprocessing import StandardScaler
from sklearn.manifold import TSNE
import gradio as gr
import os

training_tsne_path = "training_tsne.xlsx"  

if os.path.exists(training_tsne_path):
    training_tsne_df = pd.read_excel(training_tsne_path)
    print("Training t-SNE data loaded successfully!")
else:
    training_tsne_df = None
    print("Error: training_tsne.xlsx not found. Please upload it.")

columns_to_transform = ['Mn', 'Co', 'Ni', 'Cu', 'Zn', 'As', 'Se', 'Mo', 'Ag', 'Cd', 'Sb', 'Te', 'Tl', 'Pb', 'Bi']

imputer = SimpleImputer(strategy='mean')
scaler = StandardScaler()

def process_test_data(test_file):
    test_df = pd.read_excel(test_file.name)
    missing_columns = [col for col in columns_to_transform if col not in test_df.columns]
    if missing_columns:
        return f"Error: Missing columns in the uploaded data: {', '.join(missing_columns)}"
    test_df[columns_to_transform] = np.log1p(test_df[columns_to_transform])
    X_test_imputed = imputer.fit_transform(test_df[columns_to_transform])
    X_test_scaled = scaler.fit_transform(X_test_imputed)
    return X_test_scaled

def visualize_tsne_interface(test_file):
    if training_tsne_df is None:
        return "Error: Training t-SNE data not found. Please upload `training_tsne.xlsx`."
    X_test_scaled = process_test_data(test_file)
    if isinstance(X_test_scaled, str):
        return X_test_scaled 
    tsne_test = TSNE(n_components=2, random_state=42)
    X_tsne_test = tsne_test.fit_transform(X_test_scaled)
    plt.figure(figsize=(10, 7))
    sns.scatterplot(data=training_tsne_df,
                    x='x',
                    y='y',
                    hue='label',
                    palette=sns.color_palette("hsv", len(training_tsne_df['label'].unique())),
                    legend='full',
                    alpha=0.6)
    plt.scatter(X_tsne_test[:, 0], X_tsne_test[:, 1],
                c='black', label='Test Data', marker='x')
    plt.title('t-SNE Visualization: Training Data with Test Data Points')
    plt.xlabel('t-SNE Component 1')
    plt.ylabel('t-SNE Component 2')
    plt.legend(title='Tectonic Environments')
    output_file = "tsne.png"
    plt.savefig(output_file)
    plt.close()  
    return output_file


css = """
#footer {
    position: fixed !important;
    left: 10px !important;
    bottom: 10px !important;
    font-size: 12px;
    color: #555 !important;
    z-index: 9999;
}
"""

footer = """
<div id="footer">
    <strong>Dr Muhammad Amar Gul</strong><br>
    School of Earth and Space Sciences<br>
    University of Science and Technology of China
</div>
"""

with gr.Blocks(css=css) as demo:
    gr.HTML("""
    <div style="text-align: center;">
        <p style="font-size: 20px;">
            Upload an Excel file containing the required features ('Mn', 'Co', 'Ni', 'Cu', 'Zn', 'As', 'Se', 'Mo', 'Ag', 'Cd', 'Sb', 'Te', 'Tl', 'Pb', 'Bi'). 
            This application will provide t-SNE visualization. The t-SNE plot will show how your test data points relate to the training data.
        </p>
    </div>
    """)
    
    with gr.Row():
        test_file = gr.File(label="Upload Test Data Excel File")
        image_output = gr.Image(type="filepath", label="t-SNE Plot")
    test_button = gr.Button("Generate Plot")
    test_button.click(visualize_tsne_interface, inputs=test_file, outputs=image_output)
    gr.HTML(footer)

demo.launch(share=True)