from PIL import Image
import numpy as np

def remove_black_background(img_path, out_path, t_low=32, t_high=80):
    # Carregar imagem e garantir modo RGBA
    img = Image.open(img_path).convert("RGBA")
    data = np.array(img)
    
    # Extrair canais
    r, g, b, a = data[:,:,0], data[:,:,1], data[:,:,2], data[:,:,3]
    
    # O valor de brilho de cada pixel é dado pelo máximo dos canais R, G, B
    max_rgb = np.fmax(np.fmax(r, g), b)
    
    # Criar nova imagem com o mesmo tamanho
    new_data = np.zeros_like(data)
    
    # Calcular o novo alpha com base na rampa de threshold
    # Se max_rgb <= t_low -> alpha = 0
    # Se max_rgb >= t_high -> alpha = max_rgb (ou 255 se quisermos opacidade total)
    # Se t_low < max_rgb < t_high -> interpolação linear
    
    new_alpha = np.zeros_like(max_rgb, dtype=float)
    
    # Máscara para pixels totalmente opacos (ou usando o brilho original)
    mask_full = max_rgb >= t_high
    # Opcionalmente, definimos opacidade 255 para o miolo brilhante do logotipo para ficar bem nítido
    new_alpha[mask_full] = 255.0
    
    # Máscara para pixels de transição (antialiasing)
    mask_trans = (max_rgb > t_low) & (max_rgb < t_high)
    new_alpha[mask_trans] = 255.0 * (max_rgb[mask_trans] - t_low) / (t_high - t_low)
    
    # Converter alpha final para uint8
    final_alpha = np.clip(new_alpha, 0, 255).astype(np.uint8)
    
    # Máscara de pixels visíveis
    mask_visible = final_alpha > 0
    
    # Inicializar os canais de cores com zeros
    new_data[:,:,0] = 0
    new_data[:,:,1] = 0
    new_data[:,:,2] = 0
    new_data[:,:,3] = final_alpha
    
    # Aplicar unmultiplying alpha nas cores para remover o halo escuro e clarear pixels de transição
    alpha_float = final_alpha.astype(float)
    for i in range(3):
        channel = data[:,:,i].astype(float)
        adjusted = np.zeros_like(channel)
        # Onde for visível, recuperamos a cor original. 
        # Para evitar divisão por zero, usamos a máscara de pixels visíveis
        adjusted[mask_visible] = np.minimum(255, (channel[mask_visible] * 255.0 / np.maximum(1.0, max_rgb[mask_visible])))
        new_data[:,:,i] = adjusted.astype(np.uint8)
        
    # Salvar a nova imagem
    out_img = Image.fromarray(new_data, "RGBA")
    out_img.save(out_path, "PNG")
    print(f"Processada com sucesso: {img_path} -> {out_path} (T_low: {t_low}, T_high: {t_high})")

if __name__ == "__main__":
    # O logotipo tem fundo com brilho aproximado de 28
    remove_black_background("logotipo_steelXR.png", "public/logotipo_steelXR_transparente.png", t_low=32, t_high=75)
    # O ícone tem fundo com brilho de 9
    remove_black_background("iconeXR.png", "public/iconeXR_transparente.png", t_low=12, t_high=50)