import os
from multiprocessing import Pool as ProcessPool

import cv2 as cv
import numpy as np
import pandas as pd
import scipy.ndimage
from PIL import Image
from tqdm import tqdm

Image.MAX_IMAGE_PIXELS = None


def get_x_and_y(name):
    x, y = os.path.splitext(name)[0].split('_')[-2:]
    return int(x), int(y)


def flatten_list(l):
    return [item for sublist in l for item in sublist]


class NucleiFeatures:

    def position(self, img, orig, **kwargs):
        x, y = scipy.ndimage.measurements.center_of_mass(img)
        x = x + self.x_min + kwargs['x_tile'] * img.shape[0]
        y = y + self.y_min + kwargs['y_tile'] * img.shape[1]
        return [x, y]

    def ellips(self, img, orig, **kwargs):
        try:
            cont = cv.findContours(img.astype(np.uint8).T.copy(), cv.RETR_EXTERNAL, cv.CHAIN_APPROX_NONE)[1][0][:, 0, :]
            ellipse_center, axles, angle = cv.fitEllipse(cont)
            x, y = ellipse_center
            x = x + self.x_min + kwargs['x_tile'] * img.shape[0]
            y = y + self.y_min + kwargs['y_tile'] * img.shape[1]
            if axles[1] > 100:
                axles = (30, 30)
            return [*axles, x, y, angle]
        except:
            return [0] * 5

    def size(self, img, orig, **kwargs):
        return [img.sum()]

    def color(self, img, orig, **kwargs):
        cell_pixels = orig[img]
        return [*cell_pixels.mean(axis=0), *cell_pixels.std(axis=0)]

    def __init__(self, tif_folder, png_folder, features, x_min=0, y_min=0):
        self.tif_folder = tif_folder
        self.png_folder = png_folder
        self.computed_features = None
        self.feature_dict = {'position': (self.position, ['x', 'y']),
                             'size': (self.size, ['size']),
                             'ellipse': (
                                 self.ellips, ['first_axis', 'second_axis', 'ellipse_x', 'ellipse_y', 'ellipse_angle']),
                             'color': (
                                 self.color,
                                 ['Blue_mean', 'Red_mean', 'Green_mean', 'Blue_std', 'Red_std', 'Green_std']),
                             'color_gray': (
                                 self.color,
                                 ['Color_mean', 'Color_std'])
                             }
        if features == 'all':
            self.features = self.feature_dict.keys()
        else:
            self.features = features

        self.x_min = x_min
        self.y_min = y_min

    @property
    def feature_names(self):
        names = []
        for f in self.features:
            names += self.feature_dict[f][1]
        return names

    def compute(self):
        self.computed_features = []
        base_names = [os.path.splitext(i)[0] for i in os.listdir(self.tif_folder)]
        for filename in tqdm(base_names):
            img = cv.imread(f'{self.tif_folder}/{filename}.tif', -1)
            orig = cv.imread(f'{self.png_folder}/{filename}/images/{filename}.png', 1)

            img = np.rot90(img, k=3)
            img = np.flip(img, axis=1)
            orig = np.rot90(orig, k=3)
            orig = np.flip(orig, axis=1)
            for i in range(1, img.max()):
                tmp_img = (img == i)
                tmp = []
                x_tile, y_tile = get_x_and_y(filename)
                for f in self.features:
                    tmp += self.feature_dict[f][0](tmp_img, orig, x_tile=x_tile, y_tile=y_tile)
                self.computed_features.append(tmp)
        return self

    def compute_multipricess(self, n_workers=10):

        orig_list = []
        img_list = []
        filenames = []
        base_names = [os.path.splitext(i)[0] for i in os.listdir(self.tif_folder)]
        for filename in base_names:
            img = cv.imread(f'{self.tif_folder}/{filename}.tif', -1)
            orig = cv.imread(f'{self.png_folder}/{filename}/images/{filename}.png', 1)

            img = np.rot90(img, k=3)
            img = np.flip(img, axis=1)
            orig = np.rot90(orig, k=3)
            orig = np.flip(orig, axis=1)

            orig_list.append(orig)
            img_list.append(img)
            filenames.append(filename)

        global compute_one

        def compute_one(data):
            computed_features = []
            img = data[0]
            orig = data[1]
            filename = data[2]

            for i in range(1, img.max()):
                tmp_img = (img == i)
                tmp = []
                x_tile, y_tile = get_x_and_y(filename)
                for f in self.features:
                    tmp += self.feature_dict[f][0](tmp_img, orig, x_tile=x_tile, y_tile=y_tile)
                computed_features.append(tmp)
            return computed_features

        with ProcessPool(n_workers) as p:
            result = list(tqdm(p.imap(compute_one, zip(img_list, orig_list, filenames)), total=len(orig_list)))
            self.computed_features = flatten_list(result)

        return self

    def df(self):
        if self.computed_features is None:
            self.compute()
        df = pd.DataFrame(self.computed_features, columns=self.feature_names)
        return df