import os
import csv
import cv2
import tensorflow as tf
import numpy as np

class DataSetLoader:
    def __init__(self, p: str, x: int, y: int) -> None:
        if not os.path.exists(p) or p is None:
            raise Exception("Path is not correct")
        else:
            self.path_to_dataset = p
        
        self.x_size = x if x is not None else 25
        self.y_size = y if y is not None else 25


    def get_concept(self, path) -> int:
        if "fahrtrichtung_links" in path:
            return 1
        elif "fahrtrichtung_rechts" in path:
            return 2
        elif "rechts_vor_links" in path:
            return 3
        elif "stop" in path:
            return 4
        elif "vorfahrt_gewaehren" in path:
            return 5
        elif "vorfahrtsstrasse" in path:
            return 6
        else:
            return 7


    def get_classified_csv(self):
        images = []
        for dirpath, dnames, fnames in os.walk(self.path_to_dataset):
            for fname in fnames:
                image_path = os.path.join(dirpath, fname)
                concept = self.get_concept(image_path)
                if any(ext in fname for ext in ['.bmp', '.jpeg', '.jpg', '.png']):
                    images.append((image_path, concept))

        with open(os.path.abspath(os.path.join(__file__, "..", "..", "all_data_cropped.csv")), "w") as csv_output:
            csv_writer = csv.writer(csv_output, delimiter=";")
            csv_writer.writerow(["name", "concept"])
            for image in images:
                csv_writer.writerow(image)


    def create_dataset_from_csv(self, path: str):
        images = []
        concepts = []

        with open(path, "r") as csv_data:
            csv_reader = csv.reader(csv_data, delimiter=";")
            next(csv_reader)
            for row in csv_reader:
                images.append(row[0])
                concepts.append(int(row[1]))

        all_images = []
        all_labels = []

        for idx, (img_path, label) in enumerate(zip(images, concepts)):
            try:
                image = self.load_image_with_opencv(img_path)
                all_images.append(image)
                all_labels.append(label)
                tf.print(idx)
            except Exception as e:
                print(f"Error loading image {img_path}: {e}")
                continue  # Skip this image if error occurs

        all_images = np.array(all_images)
        all_labels = np.array(all_labels)

        dataset = tf.data.Dataset.from_tensor_slices((all_images, all_labels))

        return dataset


    def load_image_with_opencv(self, file_path):
        image = cv2.imread(file_path)

        if image is None:
            raise Exception(f"Could not load image: {file_path}")
        
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

        image = cv2.resize(image, (self.x_size, self.y_size))

        image = image.astype(np.float32) / 255.0

        return image