import cv2
import numpy as np
import tensorflow as tf
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, Flatten

class RecycleModel:
    def __init__(self):
        self.vgg16_model = Sequential()

        self.pretrained_model = tf.keras.applications.VGG16(
            include_top=False,
            weights="imagenet",
            input_shape=(180, 180, 3),
            pooling='max',
            classes=7
        )
        for layer in self.pretrained_model.layers:
            layer.trainable = False

        self.vgg16_model.add(self.pretrained_model)
        self.vgg16_model.add(Flatten())
        self.vgg16_model.add(Dense(512, activation='relu'))
        self.vgg16_model.add(Dense(7, activation='softmax'))
        
        self.vgg16_model.compile(optimizer=Adam(learning_rate=0.001), 
                                 loss='sparse_categorical_crossentropy', 
                                 metrics=['accuracy'])

        self.vgg16_model.load_weights('updated_vgg16_weights.h5')
        
        self.class_names = ['1_polyethylene_PET', '2_high_density_polyethylene_PE-HD', '3_polyvinylchloride_PVC',
                            '4_low_density_polyethylene_PE-LD', '5_polypropylene_PP', '6_polystyrene_PS', '8_no_plastic']
        
    def recyclePrediction(self, img):
        image_resized = cv2.resize(img, (180, 180))
        image = np.expand_dims(image_resized, axis=0)
        pred = self.vgg16_model.predict(image)
        output_class = self.class_names[np.argmax(pred)]
        
        return output_class