import gradio as gr
import torch
import jax
import jax.numpy as jnp
import numpy as np
from PIL import Image
import pickle
import warnings
import logging
from datetime import datetime
from huggingface_hub import hf_hub_download
from diffusers import StableDiffusionXLImg2ImgPipeline
from transformers import DPTImageProcessor, DPTForDepthEstimation
from model import build_thera

# Configuração de logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler("processing.log"),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

# Configurações e supressão de avisos
warnings.filterwarnings("ignore", category=FutureWarning)
warnings.filterwarnings("ignore", category=UserWarning)

# Configurar dispositivos
JAX_DEVICE = jax.devices("cpu")[0]
TORCH_DEVICE = "cpu"


# 1. Carregar modelos do Thera ----------------------------------------------------------------
def load_thera_model(repo_id, filename):
    try:
        logger.info(f"Iniciando carregamento do modelo Thera de {repo_id}")
        model_path = hf_hub_download(repo_id=repo_id, filename=filename)
        with open(model_path, 'rb') as fh:
            check = pickle.load(fh)
            variables = check['model']
            backbone, size = check['backbone'], check['size']
        model = build_thera(3, backbone, size)
        logger.info("Modelo Thera carregado com sucesso")
        return model, variables
    except Exception as e:
        logger.error(f"Falha ao carregar modelo Thera: {str(e)}")
        raise


logger.info("Carregando Thera EDSR...")
model_edsr, variables_edsr = load_thera_model("prs-eth/thera-edsr-pro", "model.pkl")

# 2. Carregar SDXL + LoRA ---------------------------------------------------------------------
try:
    logger.info("Iniciando carregamento do SDXL + LoRA...")
    pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained(
        "stabilityai/stable-diffusion-xl-base-1.0",
        torch_dtype=torch.float32
    ).to(TORCH_DEVICE)
    pipe.load_lora_weights("KappaNeuro/bas-relief", weight_name="BAS-RELIEF.safetensors")
    logger.info("SDXL + LoRA carregado com sucesso")
except Exception as e:
    logger.error(f"Falha ao carregar SDXL: {str(e)}")
    raise

# 3. Carregar modelo de profundidade ----------------------------------------------------------
try:
    logger.info("Iniciando carregamento do DPT Depth...")
    feature_extractor = DPTImageProcessor.from_pretrained("Intel/dpt-large")
    depth_model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large").to(TORCH_DEVICE)
    logger.info("Modelo DPT carregado com sucesso")
except Exception as e:
    logger.error(f"Falha ao carregar DPT: {str(e)}")
    raise


# Pipeline principal --------------------------------------------------------------------------
def full_pipeline(image, prompt, scale_factor=2.0, progress=gr.Progress()):
    try:
        progress(0, desc="Iniciando processamento...")

        # 1. Super Resolução com Thera
        progress(0.1, desc="Convertendo imagem para RGB...")
        image = image.convert("RGB")

        progress(0.2, desc="Preparando entrada para super-resolução...")
        source = np.array(image) / 255.0
        original_size = image.size
        target_shape = (int(image.height * scale_factor), int(image.width * scale_factor))
        logger.info(f"Super-resolução: {original_size} → {target_shape} (scale: {scale_factor}x)")

        progress(0.3, desc="Processando com Thera...")
        source_jax = jax.device_put(source, JAX_DEVICE)
        t = jnp.array([1.0 / (scale_factor ** 2)], dtype=jnp.float32)

        start_time = datetime.now()
        upscaled = model_edsr.apply(
            variables_edsr,
            source_jax,
            t,
            target_shape
        )
        logger.info(f"Super-resolução concluída em {datetime.now() - start_time}")

        progress(0.5, desc="Convertendo resultado...")
        upscaled_pil = Image.fromarray((np.array(upscaled) * 255).astype(np.uint8))
        logger.info(f"Tamanho após super-resolução: {upscaled_pil.size}")

        # 2. Gerar Bas-Relief
        progress(0.6, desc="Gerando Bas-Relief...")
        full_prompt = f"BAS-RELIEF {prompt}, insanely detailed and complex engraving relief, ultra-high definition, rich in detail, 16K resolution"
        logger.info(f"Prompt final: {full_prompt}")

        start_time = datetime.now()
        bas_relief = pipe(
            prompt=full_prompt,
            image=upscaled_pil,
            strength=0.7,
            num_inference_steps=25,
            guidance_scale=7.5
        ).images[0]
        logger.info(f"Bas-Relief gerado em {datetime.now() - start_time}")

        # 3. Calcular Depth Map
        progress(0.8, desc="Calculando mapa de profundidade...")
        start_time = datetime.now()
        inputs = feature_extractor(bas_relief, return_tensors="pt").to(TORCH_DEVICE)

        with torch.no_grad():
            outputs = depth_model(**inputs)
            depth = outputs.predicted_depth

        depth_map = torch.nn.functional.interpolate(
            depth.unsqueeze(1),
            size=bas_relief.size[::-1],
            mode="bicubic"
        ).squeeze().cpu().numpy()

        progress(0.9, desc="Processando mapa de profundidade...")
        depth_min = depth_map.min()
        depth_max = depth_map.max()
        depth_normalized = (depth_map - depth_min) / (depth_max - depth_min + 1e-8)
        depth_pil = Image.fromarray((depth_normalized * 255).astype(np.uint8))
        logger.info(f"Profundidade calculada em {datetime.now() - start_time} | Range: {depth_min:.2f}-{depth_max:.2f}")

        progress(1.0, desc="Finalizado!")
        return upscaled_pil, bas_relief, depth_pil

    except Exception as e:
        logger.error(f"Erro no processamento: {str(e)}", exc_info=True)
        raise gr.Error(f"Erro no processamento: {str(e)}")


# Interface Gradio ----------------------------------------------------------------------------
with gr.Blocks(title="Super Res + Bas-Relief") as app:
    gr.Markdown("## 🔍 Super Resolução + 🗿 Bas-Relief + 🗺️ Profundidade")

    with gr.Row():
        with gr.Column():
            img_input = gr.Image(type="pil", label="Imagem de Entrada")
            prompt = gr.Textbox(
                label="Descrição do Relevo",
                value="insanely detailed and complex engraving relief, ultra-high definition, rich in detail, and 16K resolution."
            )
            scale = gr.Slider(1.0, 4.0, value=2.0, label="Fator de Escala")
            btn = gr.Button("Processar")

        with gr.Column():
            img_upscaled = gr.Image(label="Imagem Super Resolvida")
            img_basrelief = gr.Image(label="Resultado Bas-Relief")
            img_depth = gr.Image(label="Mapa de Profundidade")

    btn.click(
        full_pipeline,
        inputs=[img_input, prompt, scale],
        outputs=[img_upscaled, img_basrelief, img_depth]
    )

if __name__ == "__main__":
    logger.info("Iniciando aplicação Gradio")
    app.launch(share=False)