import streamlit as st
import matplotlib.pyplot as plt

from keras.datasets import mnist,imdb
from sklearn.preprocessing import MinMaxScaler
from keras.models import Model,Sequential
from keras.layers import Conv2D,MaxPooling2D,InputLayer,Dense,Flatten

m = MinMaxScaler()
with st.sidebar:
    option = st.selectbox("Datasets",["Select dataset","Hand Writen Digit Dataset","imdb"])
    if option == "Hand Writen Digit Dataset":
        (x_train,y_train),(x_test,y_test) = mnist.load_data()
        st.success(f"Successfully Load the mnist Dataset")
    elif option == "imdb":
        (x_train,y_train),(x_test,y_test) = imdb.load_data()
        st.success(f"Successfully Load the imdb Dataset")
    
    col1, col2 = st.columns(2)
    with col1:
        epochs = st.number_input("Epochs (100 - 2000)", 5, 2000, step=5)
    with col2:
        neuron_input = st.text_input("Neurons per Layer (comma-separated, e.g. 5,4,3)")
        neurons_per_layer = [int(n.strip()) for n in neuron_input.split(",") if n.strip().isdigit()]

    submit = st.button("Train")


if submit:
    model = Sequential()
    model.add(InputLayer(shape=(28,28,1)))
    model.add(Conv2D(filters=6,kernel_size=(5,5),activation="tanh",padding="valid",strides=(1,1)))
    model.add(MaxPooling2D(pool_size=(2,2),strides=(2,2)))
    model.add(Conv2D(filters=16,kernel_size=(5,5),activation="tanh",padding="same",strides=(1,1)))
    model.add(MaxPooling2D(pool_size=(2,2),strides=(2,2)))
    model.add(Conv2D(filters=30,kernel_size=(5,5),activation="tanh",padding="valid",strides=(1,1)))
    model.add(Flatten())
    for i in neurons_per_layer:
        model.add(Dense(units=i,activation="tanh"))
    model.add(Dense(10,activation="softmax"))
    model.compile(optimizer="sgd",loss="sparse_categorical_crossentropy",metrics=["accuracy"])
    model.fit(x_train, y_train, epochs=3, batch_size=128, validation_split=0.4,verbose=False)
    history = model.fit(
        x_train, y_train,
        batch_size=128,
        epochs=epochs,
        validation_split=0.4,
        verbose=0
    )

    fig, axs = plt.subplots(6, 1, figsize=(8, 6))
    col1,col2,col3,col4,col5,col6 = st.columns(6)
    with col1:
        st.write("Filters")
        for i in range(6):
            axs[i].imshow(
                m.fit_transform(model.layers[0].weights[0][:, :, :, i][:, :, 0]),
                cmap="gray"
            )
            axs[i].axis("off")
        st.pyplot(fig)
    with col2:
        st.write("Conv2D Layer-1")
        sub = Model(inputs=model.inputs[0],outputs=model.layers[0].output)
        fig,ax = plt.subplots(6,1,figsize=(8,6))
        for i in range(6):
            ax[i].imshow(sub.predict(x_train[0:1,:].reshape(1,28,28,1))[0,:,:,i],cmap="gray")
            ax[i].axis("off")
        st.pyplot(fig)
    with col3:
        st.write("Average Pooling Layer-1")
        max1 = Model(inputs=model.inputs[0],outputs=model.layers[1].output)
        fig,ax1 = plt.subplots(6,1,figsize=(8,6))
        for i in range(6):
            ax1[i].imshow(max1.predict(x_train[0:1,:].reshape(1,28,28,1))[0,:,:,i],cmap="gray")
            ax1[i].axis("off")
        st.pyplot(fig)
    
    with col4:
        st.write("Conv2D Layer-2")
        sub1 = Model(inputs=model.inputs[0],outputs=model.layers[2].output)
        fig,ax2 = plt.subplots(16,1,figsize=(8,6))
        for i in range(16):
            ax2[i].imshow(sub1.predict(x_train[0:1,:].reshape(1,28,28,1))[0,:,:,i],cmap="gray")
            ax2[i].axis("off")
        st.pyplot(fig)
    with col5:
        st.write("Average Pooling Layer-2")
        max2 = Model(inputs=model.inputs[0],outputs=model.layers[3].output)
        fig,ax3 = plt.subplots(16,1,figsize=(8,6))
        for i in range(16):
            ax3[i].imshow(max2.predict(x_train[0:1,:].reshape(1,28,28,1))[0,:,:,i],cmap="gray")
            ax3[i].axis("off")
        st.pyplot(fig)
    with col6:
        st.write("Conv2D Layer-3")
        sub3 = Model(inputs=model.inputs[0],outputs=model.layers[4].output)
        fig,ax4 = plt.subplots(120,1,figsize=(8,6))
        for i in range(120):
            ax4[i].imshow(sub3.predict(x_train[0:1,:].reshape(1,28,28,1))[0,:,:,i],cmap="gray")
            ax4[i].axis("off")
        st.pyplot(fig)