import numpy as np
import cv2 as cv
import matplotlib.pyplot as plt

"""
Macenko stain normalization — Hugging Face Space compatible version
Removed dependency on SPAMS (no native libraries required).
"""

def standardize_brightness(I):
    p = np.percentile(I, 95)
    return np.clip(I * 255.0 / p, 0, 255).astype(np.uint8)

def remove_zeros(I):
    mask = (I == 0)
    I[mask] = 1
    return I

def RGB_to_OD(I):
    I = remove_zeros(I)
    return -np.log((I.astype(np.float32) + 1) / 256)

def OD_to_RGB(OD):
    return np.clip(255 * np.exp(-OD), 0, 255).astype(np.uint8)

def normalize_rows(A):
    return A / np.linalg.norm(A, axis=1)[:, None]

def get_stain_matrix(I, beta=0.15, alpha=1):
    OD = RGB_to_OD(I).reshape((-1, 3))
    OD = OD[(OD > beta).any(axis=1), :]
    _, V = np.linalg.eigh(np.cov(OD, rowvar=False))
    V = V[:, [2, 1]]
    if V[0, 0] < 0: V[:, 0] *= -1
    if V[0, 1] < 0: V[:, 1] *= -1
    That = np.dot(OD, V)
    phi = np.arctan2(That[:, 1], That[:, 0])
    minPhi = np.percentile(phi, alpha)
    maxPhi = np.percentile(phi, 100 - alpha)
    v1 = np.dot(V, np.array([np.cos(minPhi), np.sin(minPhi)]))
    v2 = np.dot(V, np.array([np.cos(maxPhi), np.sin(maxPhi)]))
    HE = np.array([v1, v2])
    return normalize_rows(HE)

def get_concentrations(I, stain_matrix):
    OD = RGB_to_OD(I).reshape((-1, 3))
    return np.linalg.lstsq(stain_matrix.T, OD.T, rcond=None)[0].T

class macenko_normalizer(object):
    def __init__(self):
        self.stain_matrix_target = np.array(
            [[0.5626, 0.2159], [0.7201, 0.8012], [0.4062, 0.5581]], dtype=np.float32).T
        self.target_concentrations = None

    def fit(self, target):
        target = standardize_brightness(target)
        self.stain_matrix_target = get_stain_matrix(target)
        self.target_concentrations = get_concentrations(target, self.stain_matrix_target)

    def transform(self, I):
        I = standardize_brightness(I)
        stain_matrix_source = get_stain_matrix(I)
        source_concentrations = get_concentrations(I, stain_matrix_source)
        maxC_source = np.percentile(source_concentrations, 99, axis=0).reshape((1, 2))
        maxC_target = np.array([1.9705, 1.0308], dtype=float).reshape((1, 2))
        source_concentrations *= maxC_target / maxC_source
        return OD_to_RGB(np.dot(source_concentrations, self.stain_matrix_target).reshape(I.shape))