import streamlit as st
import tensorflow as tf
from tensorflow import keras
from keras import models
from PIL import Image
import numpy as np
import cv2
import io

# Some constants to be used in the program
IMG_SIZE = (32,32)

# Character mapping for the character prediction
char_map = {
    0:'𑑐(0)',   1:'𑑑(1)',    2:'𑑒(2)',   3:'𑑓(3)',      4: '𑑔(4)',     5: '𑑕(5)',    6: '𑑖(6)',    7: '𑑗(7)',
    8:'𑑘(8)',   9:'𑑙(9)',    10:'𑑉(OM)', 11:'𑐀(A)',    12: '𑐁(AA)',   13: '𑐀𑑅(AH)',  14: '𑐂(I)',    
    15:'𑐃(II)',16:'𑐄(U)',   17:'𑐅(UU)',  18:'𑐆(R)',    19: '𑐆𑐺(RR)',  20: '𑐊(E)',   21: '𑐋(AI)',    22: '𑐌(O)',    
    23:'𑐍(AU)', 24:'𑐈(L)',  25:'𑐉(LL)',   26:'𑐎(KA)',   27: '𑐎𑑂𑐳(KSA)', 28: '𑐏(KHA)',29: '𑐐(GA)',    30: '𑐑(GHA)',    
    31:'𑐒(NGA)',32:'𑐔(CA)',  33:'𑐕(CHA)', 34:'𑐖(JA)',   35: '𑐖𑑂𑐘(JñA)',  36: '𑐗(JHA)',37: '𑐗(JHA-alt)',38: '𑐘(NYA)',    
    39:'𑐚(TA)', 40:'𑐛(TTHA)', 41:'𑐜(DDA)', 42:'𑐝(DHA)',  43: '𑐞(NNA)', 44: '𑐟(TA)',  45: '𑐟𑑂𑐬(TRA)',    46: '𑐠(THA)',
    47:'𑐡(DA)', 49:'𑐣(NA)',   50:'𑐥(PA)',  51:'𑐦(PHA)',  52: '𑐧(BA)',  53: '𑐨(BHA)',  54: '𑐩(MA)',    55: '𑐫(YA)', 
    56:'𑐬(RA)', 57: '𑐮(LA)', 58:'𑐰(WA)', 59:'𑐱(SHA)',    60: '𑐱(SHA-alt)', 61: '𑐲(SSA)',    62: '𑐳(SA)', 63: '𑐴(HA)'
}

# Importing the model
model = models.load_model('tf_model.h5')

# Function for reading image
def file_to_array(file) -> np.ndarray:
    image = np.array(Image.open(io.BytesIO(file)))

    return image 

# Main Streamlit app
def main():
    st.title("Character Recognition")
    st.write("Upload an image and the model will predict the character")

    uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])

    if uploaded_file is not None:
        image = Image.open(uploaded_file)
        st.image(image, caption='Uploaded Image.', use_column_width=True)

        if st.button('Predict'):
            image = cv2.resize(np.array(image), IMG_SIZE)
            image = image.astype('float32')
            image = np.expand_dims(image, axis=0)

            output = model.predict(image)
            result = char_map[np.argmax(output)]
            
            st.success(f'Prediction: {result}')

if __name__ == "__main__":
    main()