import os
import csv
import cv2 as cv
import re

from util.featue_extraction import get_color_percentage, get_raster_color_percentage, get_edges, get_corners, get_contours
from classes.feature_vector import FeatureVector
from classes.concept import Concept
from classes.feature import Feature

img_path = os.path.abspath(os.path.join(__file__, "..", "..", "data", "processed"))
vector_path = os.path.abspath(os.path.join(__file__, "..", "..", "data", "vectors"))
csv_path = os.path.abspath(os.path.join(__file__, "..", "all_feature_vectors.csv"))

def get_concept(path) -> Concept:
    if "fahrtrichtung_links" in path:
        return Concept.LINKS_ABBIEGEN
    elif "fahrtrichtung_rechts" in path:
        return Concept.RECHTS_ABBIEGEN
    elif "rechts_vor_links" in path:
        return Concept.RECHTS_VOR_LINKS
    elif "stop" in path:
        return Concept.STOP
    elif "vorfahrt_gewaehren" in path:
        return Concept.VORFAHRT_GEWAEHREN
    elif "vorfahrtsstrasse" in path:
        return Concept.VORFAHRT_STRASSE
    else:
        return Concept.UNKNOWN

if __name__ == "__main__":
    # create csv file and insert headers
    with open(csv_path, "w") as csv_file:
                csv_writer = csv.writer(csv_file, delimiter=";")
                csv_writer.writerow(["brightness","filename", "concept", "vector"])

    # go to every file in img_path
    for dirpath, dnames, fnames in os.walk(img_path):
        save_path = os.path.join(dirpath.replace(img_path, vector_path))
        
        # create directories if not exist
        if not os.path.exists(save_path):
            os.makedirs(save_path)

        for fname in fnames:
            image_path = os.path.join(dirpath, fname)

            # only create vectors of images in directory 0 (normal brightness)
            # if not "0\\3500" in image_path:
            #     continue
            
            try:
                # prepare image
                image = cv.imread(image_path)
                gray_image = cv.imread(image_path, cv.IMREAD_GRAYSCALE)
                #extract features
                color_percentage_overall = get_color_percentage(image)
                color_percentage_rasters = get_raster_color_percentage(image, 4, 4)
                edges_amount, edged_image = get_edges(gray_image)
                edges_amount = round(edges_amount / 1000)
                corners_amount = get_corners(gray_image)
                contours_amount = get_contours(edged_image)

            except Exception:
                print("Failed for ", image_path)
                continue
                
            concept: Concept = get_concept(image_path)
            if concept == Concept.UNKNOWN:
                print("no concept found for ", image_path)
                continue

            vec: FeatureVector = FeatureVector(concept)
            vec.add_feature(Feature.CORNERS, corners_amount)
            vec.add_feature(Feature.EDGES, edges_amount)
            vec.add_feature(Feature.COUNTOURS, contours_amount)
            vec.add_feature(Feature.OVERALL_COLOR_PERCENTAGE, color_percentage_overall)
            vec.add_feature(Feature.RASTER_COLOR_PERCENTAGE, color_percentage_rasters)

            brightness = re.search(r'(-2|-1|0|\+1|\+2)', image_path)[1]

            with open(csv_path, "a") as csv_file:
                csv_writer = csv.writer(csv_file, delimiter=";", lineterminator="\n")
                csv_writer.writerow([brightness, fname, concept, vec.get_vector()])
            
            vec_path = os.path.join(save_path, fname.replace(".jpg", ".pkl"))
            vec.save(vec_path)