from sklearn.svm import LinearSVC
import numpy as np
import cv2 as cv
import os

class Model:

    def __init__(self):
        self.model = LinearSVC()
        self.trained = False

    def train_model(self, counters):

        images = []
        labels = []

        # Class 1
        for i in range(1, counters[0]):
            path = f'1/frame{i}.jpg'
            if os.path.exists(path):
                img = cv.imread(path, cv.IMREAD_GRAYSCALE)
                img = cv.resize(img, (150, 150))
                images.append(img.flatten())
                labels.append(1)

        # Class 2
        for i in range(1, counters[1]):
            path = f'2/frame{i}.jpg'
            if os.path.exists(path):
                img = cv.imread(path, cv.IMREAD_GRAYSCALE)
                img = cv.resize(img, (150, 150))
                images.append(img.flatten())
                labels.append(2)

        if len(images) == 0:
            print("No training data found!")
            return

        X = np.array(images)
        y = np.array(labels)

        self.model.fit(X, y)
        self.trained = True

        print("Model trained successfully!")

    def predict(self, frame):

        if not self.trained:
            return None

        gray = cv.cvtColor(frame, cv.COLOR_RGB2GRAY)
        gray = cv.resize(gray, (150, 150))
        gray = gray.flatten()

        prediction = self.model.predict([gray])
        return prediction[0]