from typing import Any, Self
import numpy as np
import os
import pickle

from classes.concept import Concept
from classes.feature import Feature


class FeatureVector:

	def __init__(self, concept: Concept, features: dict = {}, features_list: list = None, loaded: bool = False) -> None:
		if loaded == False:
			self.loaded: bool = False
			self.features: dict = features
			self.concept: Concept = concept
		else:
			self.loaded: bool = True
			self.concept: Concept = concept
			self.features_list: list = features_list


	def _get_values_of_feature(self, feature: Feature) -> list[int]:
		try:
			feature_data = self.features[feature]
		except KeyError:
			print("missing key: ", feature)
			return [-1]

		ret = []
		if feature == Feature.OVERALL_COLOR_PERCENTAGE:
			feature_data: dict
			for val in feature_data.values():
				ret.append(val)
			return ret
		
		elif feature == Feature.RASTER_COLOR_PERCENTAGE:
			feature_data: list[list[dict]]
			
			for feature_data1d in feature_data:
				for data in feature_data1d:
					for val in data.values():
						ret.append(val)
			return ret
		
		elif feature == Feature.CORNERS:
			ret.append(feature_data)
			return ret
		elif feature == Feature.EDGES:
			ret.append(feature_data)
			return ret
		elif feature == Feature.COUNTOURS:
			ret.append(feature_data)
			return ret


	def add_feature(self, key: Feature, value: Any) -> None:
		self.features.update({key: value})


	def get_vector(self) -> list:
		if self.loaded:
			return self.features_list
		
		ret = []
		for feature in Feature:
			ret = ret + self._get_values_of_feature(feature)
		return ret


	def get_concept(self) -> Concept:
		return self.concept


	def get_num_features(self) -> int:
		return len(self.features)


	def get_feature_value(self, key: str) -> Any:
		return self.features[key]
	

	def get_features(self) -> Any:
		return self.features
	

	def save(self, path: os.path) -> None:
		with open(path, "wb") as write:
			pickle.dump(self, write, pickle.HIGHEST_PROTOCOL)


	def load(path: os.path) -> Self:
		with open(path, "rb") as read:
			self: FeatureVector = pickle.load(read)
		return self