import numpy as np
import pandas as pd
from keras.models import load_model
from sklearn.preprocessing import StandardScaler
import matplotlib.pyplot as plt
import pickle
import gradio as gr
import io
import requests

# Load the model and scaler
transformer = load_model('transformer_pile.h5')
sc_f = pickle.load(open('scaler_f.pkl', 'rb'))

# URL to the default Excel file
DEFAULT_EXCEL_URL = "https://huggingface.co/spaces/neng123/Pile_deform/resolve/main/data_pile.xlsx"

def download_default_excel():
    response = requests.get(DEFAULT_EXCEL_URL)
    return io.BytesIO(response.content)

def process_excel(file):
    if file == "default":
        file = download_default_excel()
    df = pd.read_excel(file, sheet_name='soil')
    df_y = pd.read_excel(file, sheet_name='pile')
    df_p = pd.read_excel(file, sheet_name='pile_length')
    data = np.array(df)
    data_y = np.array(df_y)
    data_pile = np.array(df_p)[:, 1:61]
    x_feature = data_y[:, 0:4]
    bh = data[:, 1:61]
    bh2 = data[:, 61:122] / 2
    x_train = bh / 50
    x_feature = sc_f.transform(x_feature)
    soil_data = np.stack([x_train, bh2, data_pile], axis=2)
    return soil_data

def validate_inputs(pile_length, section_width, section_length, pile_type):
    if pile_length < 39 or pile_length > 60:
        return False, "Pile length must be between 39 and 60 meters."
    
    if pile_type == "Circular":
        if section_width < 0.8 or section_width > 1.8 or section_length < 0.8 or section_length > 1.8:
            return False, "For circular piles, section width and length must be between 0.8 and 1.8 meters."
    elif pile_type == "Barrette":
        if section_width < 1 or section_width > 4 or section_length < 1 or section_length > 4:
            return False, "For barrette piles, section width and length must be between 1 and 4 meters."
    
    return True, ""

def predict_pile(file_choice, uploaded_file, pile_length, section_width, section_length, pile_type):
    # Validate inputs
    is_valid, error_message = validate_inputs(pile_length, section_width, section_length, pile_type)
    if not is_valid:
        return None, error_message

    if file_choice == "default":
        file = "default"
    else:
        if uploaded_file is None:
            return None, "Please upload an Excel file or choose the default option."
        file = uploaded_file
    
    X = process_excel(file)
    
    # Convert pile type to numerical value
    pile_type_num = 1 if pile_type == "Circular" else 2
    
    # Create feature array
    feature = np.array([pile_length, section_width, section_length, pile_type_num])
    feature = np.reshape(feature, (1, -1))
    fd = sc_f.transform(feature)
    x_feature = fd
    
    # Use the first sample for demonstration
    Xd = X[0:1]
    X_train_CNN = np.zeros((Xd.shape[0], Xd.shape[1], x_feature.shape[1] + 3))
    X_train_CNN[:, :, 0:3] = Xd
    for i in range(Xd.shape[0]):
        X_train_CNN[i, :, 3] = x_feature[i, 1]
        X_train_CNN[i, :, 4] = x_feature[i, 2]
        X_train_CNN[i, :, 5] = x_feature[i, 3]
        X_train_CNN[i, :, 6] = x_feature[i, 3]
    
    XT = X_train_CNN
    print(XT.shape)
    
    y_ini = np.zeros((1, 40))
    y_ini[0, 0] = 0
    for step in range(39):
        y = transformer.predict([XT, y_ini, fd])
        y_ini[0, step+1] = y
    
    y_pred = y_ini * 40000
    
    plt.figure(figsize=(10, 6))
    ydist = range(1, 41)
    plt.plot(ydist, y_pred[0], color='blue', label='predict')
    plt.legend()
    plt.xlabel("Deformation (mm)")
    plt.ylabel("Load (kN)")
    plt.title(f"Pile Prediction (Length: {pile_length}m, {pile_type})")
    return plt, ""

def update_file_input(choice):
    return gr.update(visible=choice == "upload")

with gr.Blocks() as iface:
    gr.Markdown("# Pile Prediction Model in Bangkok Subsoil")
    gr.Markdown("Choose the default Excel file or upload your own, then enter pile characteristics to predict pile behavior.")
    
    with gr.Row():
        file_choice = gr.Radio(["default", "upload"], label="Choose Excel File Source", value="default")
        uploaded_file = gr.File(label="Upload Excel File", type="binary", visible=False)
    
    with gr.Row():
        pile_length = gr.Number(label="Pile Length (m) please change data in excel file together", value=42)
        section_width = gr.Number(label="Section Width (m)", value=1)
        section_length = gr.Number(label="Section Length (m)", value=1)
        pile_type = gr.Radio(["Circular", "Barrette"], label="Pile Type", value="Circular")
    
    output = gr.Plot()
    error_output = gr.Textbox(label="Error", visible=True)
    
    submit_btn = gr.Button("Predict")
    
    # Add a download link for the default Excel file
    gr.Markdown(f"[Download Default Excel File]({DEFAULT_EXCEL_URL})")
    
    file_choice.change(
        fn=update_file_input,
        inputs=[file_choice],
        outputs=[uploaded_file]
    )
    
    submit_btn.click(
        fn=predict_pile,
        inputs=[file_choice, uploaded_file, pile_length, section_width, section_length, pile_type],
        outputs=[output, error_output]
    )

iface.launch()