# -*- coding: utf-8 -*-
"""
Created on Tue Jan 19 21:52:54 2021

@author: femiogundare
"""

import numpy as np



class StainNormalization:
    """
    Adopted and modified from "Classification of breast cancer histology images using Convolutional Neural Networks",
    Teresa Araújo , Guilherme Aresta, Eduardo Castro, José Rouco, Paulo Aguiar, Catarina Eloy, António Polónia,
    Aurélio Campilho. https://doi.org/10.1371/journal.pone.0177544
    
    Performs staining normalization.
    """
    
    def __init__(self, Io=240, beta=0.15, alpha=1):
        # Store the image
        self.Io = Io
        self.beta = beta
        self.alpha = alpha
        
    def preprocess(self, img):
        """
        # Arguments
            img: Numpy image array.
        # Returns
            Normalized Numpy image array.
        """
        
        HERef = np.array([[0.5626, 0.2159],
                      [0.7201, 0.8012],
                      [0.4062, 0.5581]])
        maxCRef = np.array([1.9705, 1.0308])
        h, w, c = img.shape
        image = img.reshape(h * w, c)
        
        OD = -np.log((image.astype("uint16") + 1) / self.Io)
        ODhat = OD[(OD >= self.beta).all(axis=1)]
        
        if (OD >= self.beta).all(axis=1).sum() <= 1:
            return img.astype('uint8')
        
        W, V = np.linalg.eig(np.cov(ODhat, rowvar=False))

        Vec = -V.T[:2][::-1].T
        That = np.dot(ODhat, Vec)
        phi = np.arctan2(That[:, 1], That[:, 0])
        minPhi = np.percentile(phi, self.alpha)
        maxPhi = np.percentile(phi, 100 - self.alpha)
        vMin = np.dot(Vec, np.array([np.cos(minPhi), np.sin(minPhi)]))
        vMax = np.dot(Vec, np.array([np.cos(maxPhi), np.sin(maxPhi)]))
        if vMin[0] > vMax[0]:
            HE = np.array([vMin, vMax])
        else:
            HE = np.array([vMax, vMin])

        HE = HE.T
        Y = OD.reshape(h * w, c).T

        C = np.linalg.lstsq(HE, Y)
        maxC = np.percentile(C[0], 99, axis=1)

        C = C[0] / maxC[:, None]
        C = C * maxCRef[:, None]
        Inorm = self.Io * np.exp(-np.dot(HERef, C))
        Inorm = Inorm.T.reshape(h, w, c).clip(0, 255).astype("uint8")
        
        return Inorm
    
    
    
    
    
"""
def normalize_staining(img):
    
    #Adopted from "Classification of breast cancer histology images using Convolutional Neural Networks",
    #Teresa Araújo , Guilherme Aresta, Eduardo Castro, José Rouco, Paulo Aguiar, Catarina Eloy, António Polónia,
    #Aurélio Campilho. https://doi.org/10.1371/journal.pone.0177544
    #Performs staining normalization.
    # Arguments
    #    img: Numpy image array.
    # Returns
    #   Normalized Numpy image array.
    
    Io = 240
    beta = 0.15
    alpha = 1
    HERef = np.array([[0.5626, 0.2159],
                      [0.7201, 0.8012],
                      [0.4062, 0.5581]])
    maxCRef = np.array([1.9705, 1.0308])

    h, w, c = img.shape
    img = img.reshape(h * w, c)
    OD = -np.log((img.astype("uint16") + 1) / Io)
    ODhat = OD[(OD >= beta).all(axis=1)]
    W, V = np.linalg.eig(np.cov(ODhat, rowvar=False))

    Vec = -V.T[:2][::-1].T  # desnecessario o sinal negativo
    That = np.dot(ODhat, Vec)
    phi = np.arctan2(That[:, 1], That[:, 0])
    minPhi = np.percentile(phi, alpha)
    maxPhi = np.percentile(phi, 100 - alpha)
    vMin = np.dot(Vec, np.array([np.cos(minPhi), np.sin(minPhi)]))
    vMax = np.dot(Vec, np.array([np.cos(maxPhi), np.sin(maxPhi)]))
    if vMin[0] > vMax[0]:
        HE = np.array([vMin, vMax])
    else:
        HE = np.array([vMax, vMin])

    HE = HE.T
    Y = OD.reshape(h * w, c).T

    C = np.linalg.lstsq(HE, Y)
    maxC = np.percentile(C[0], 99, axis=1)

    C = C[0] / maxC[:, None]
    C = C * maxCRef[:, None]
    Inorm = Io * np.exp(-np.dot(HERef, C))
    Inorm = Inorm.T.reshape(h, w, c).clip(0, 255).astype("uint8")

    return Inorm
"""