import gradio as gr
import os
import gc

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

from collections import Counter

from sklearn.preprocessing import LabelEncoder

from keras.models import Model
from keras.regularizers import l2
from keras.constraints import max_norm
from keras.utils import to_categorical
from keras.preprocessing.text import Tokenizer
from keras.utils import pad_sequences
from keras.callbacks import EarlyStopping
from keras.layers import Input, Dense, Dropout, Flatten, Activation
from keras.layers import Conv1D, Add, MaxPooling1D, BatchNormalization
from keras.layers import Embedding, Bidirectional, LSTM, CuDNNLSTM, GlobalMaxPooling1D

import tensorflow as tf
from huggingface_hub import hf_hub_url, cached_download


def residual_block(data, filters, d_rate):
    """
    _data: input
    _filters: convolution filters
    _d_rate: dilation rate
    """
    
    shortcut = data
    
    bn1 = BatchNormalization()(data)
    act1 = Activation('relu')(bn1)
    conv1 = Conv1D(filters, 1, dilation_rate=d_rate, padding='same', kernel_regularizer=l2(0.001))(act1)
    
    #bottleneck convolution
    bn2 = BatchNormalization()(conv1)
    act2 = Activation('relu')(bn2)
    conv2 = Conv1D(filters, 3, padding='same', kernel_regularizer=l2(0.001))(act2)
    
    #skip connection
    x = Add()([conv2, shortcut])
    
    return x

def get_model():
    # model
    x_input = Input(shape=(100, 21))
    
    #initial conv
    conv = Conv1D(128, 1, padding='same')(x_input) 
    
    # per-residue representation
    res1 = residual_block(conv, 128, 2)
    res2 = residual_block(res1, 128, 3)
    
    x = MaxPooling1D(3)(res2)
    x = Dropout(0.5)(x)
    
    # softmax classifier
    x = Flatten()(x)
    x_output = Dense(1000, activation='softmax', kernel_regularizer=l2(0.0001))(x)
    
    model2 = Model(inputs=x_input, outputs=x_output)
    model2.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])
    weights = cached_download(hf_hub_url("jonathang", "Protein_Family_CNN"))
    model2.load_weights(weights)

    return model2


def greet(name):
    get_model()
    return "Hello " + name + "!!"

iface = gr.Interface(fn=greet, inputs="text", outputs="text")
iface.launch()