import numpy as np
from PIL import Image
from skimage.feature import hog
from skimage.color import rgb2gray


class FeatureExtractor:
    """
    Extract features from images.
    Current methods accept either file paths or numpy arrays representing images.
    Provides HOG and color histogram features to extract a concatenated feature vector.
    """

    def __init__(self, resize=None, hog_params=None, hist_bins=32):
        # resize: (width, height) or None
        self.resize = resize
        self.hist_bins = hist_bins
        self.hog_params = {} if hog_params is None else hog_params

    def _load_image(self, img):
        if isinstance(img, str):
            im = Image.open(img).convert('RGB')
            if self.resize:
                im = im.resize(self.resize)
            return np.asarray(im)
        elif hasattr(img, 'shape'):
            # numpy array must be either HxWxC or HxW
            im = img
            if self.resize:
                im = np.asarray(Image.fromarray(im).resize(self.resize))
            if im.ndim == 2:
                # grayscale to RGB
                im = np.stack([im, im, im], axis=-1)
            return im
        else:
            raise ValueError('Unsupported image input type')

    def extract_hog(self, img):
        im = self._load_image(img)
        gray = rgb2gray(im)
        features = hog(gray, **self.hog_params)
        return features

    def color_histogram(self, img, bins=None, normalize=True):
        im = self._load_image(img)
        bins = self.hist_bins if bins is None else bins
        chans = []
        for c in range(3):
            h, _ = np.histogram(im[:, :, c], bins=bins, range=(0, 255))
            chans.append(h.astype(np.float32))
        feat = np.concatenate(chans)
        if normalize:
            s = feat.sum()
            if s > 0:
                feat = feat / s
        return feat

    def extract(self, img, which=('hog', 'hist')):
        pieces = []
        if 'hog' in which:
            pieces.append(self.extract_hog(img))
        if 'hist' in which:
            pieces.append(self.color_histogram(img))
        if not pieces:
            raise ValueError('No feature types requested')
        return np.concatenate(pieces)

    def extract_from_list(self, images, which=('hog', 'hist'), verbose=False):
        """images: list of file paths or numpy arrays. Returns numpy array of features."""
        feats = []
        for i, im in enumerate(images):
            feats.append(self.extract(im, which=which))
            if verbose and (i + 1) % 100 == 0:
                print(f"Extracted features from {i+1} images")
        return np.vstack(feats)