import gradio as gr
import cv2
import numpy as np
import torch
import os
from pathlib import Path

# Ensure directories exist
INPUT_DIR = Path("input")
OUTPUT_DIR = Path("output")
MODELS_DIR = Path("models")
INPUT_DIR.mkdir(exist_ok=True)
OUTPUT_DIR.mkdir(exist_ok=True)

# Function to load pre-trained models
def load_model(model_path, use_cpu=False):
    if not model_path.exists():
        raise FileNotFoundError(f"Model file not found: {model_path}")

    device = "cpu" if use_cpu or not torch.cuda.is_available() else "cuda"

    # Load state_dict if the model was saved that way
    model_state = torch.load(model_path, map_location=device)

    # If a full model object was saved, load it directly
    if isinstance(model_state, torch.nn.Module):
        model = model_state
    else:
        # If saved as state_dict, we need a model architecture (Assuming CNN or custom model)
        model = torch.nn.Sequential(
            torch.nn.Conv2d(3, 64, kernel_size=3, padding=1),
            torch.nn.ReLU(),
            torch.nn.Conv2d(64, 3, kernel_size=3, padding=1)
        )
        model.load_state_dict(model_state)

    model.to(device)
    model.eval()
    return model

# Process image and save to output folder
def process_image(input_path, tile_size=512, seamless=False, use_cpu=False):
    # Read input image
    img = cv2.imread(str(input_path))
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

    # Load models
    normal_model = load_model(MODELS_DIR / "1x_NormalMapGenerator-CX-Lite_200000_G.pth", use_cpu)
    franken_model = load_model(MODELS_DIR / "1x_frankenMapGenerator-CX-Lite_215000_G.pth", use_cpu)

    # Convert to tensor
    img_tensor = torch.from_numpy(img.transpose(2, 0, 1)).float() / 255.0
    img_tensor = img_tensor.unsqueeze(0)  # Add batch dimension

    device = "cpu" if use_cpu or not torch.cuda.is_available() else "cuda"
    img_tensor = img_tensor.to(device)

    # Generate maps
    with torch.no_grad():
        normal_map = normal_model(img_tensor).cpu().numpy().squeeze()
        franken_map = franken_model(img_tensor).cpu().numpy().squeeze()

    # Post-process maps
    normal_map = (normal_map.transpose(1, 2, 0) * 255).clip(0, 255).astype(np.uint8)
    disp_map = (franken_map[0] * 255).clip(0, 255).astype(np.uint8)
    rough_map = (franken_map[1] * 255).clip(0, 255).astype(np.uint8)

    # Convert grayscale to RGB for Gradio display
    disp_map = np.stack([disp_map] * 3, axis=-1)
    rough_map = np.stack([rough_map] * 3, axis=-1)

    # Define output paths
    base_name = input_path.stem
    normal_path = OUTPUT_DIR / f"{base_name}_normal.png"
    disp_path = OUTPUT_DIR / f"{base_name}_displacement.png"
    rough_path = OUTPUT_DIR / f"{base_name}_roughness.png"

    # Save outputs
    cv2.imwrite(str(normal_path), cv2.cvtColor(normal_map, cv2.COLOR_RGB2BGR))
    cv2.imwrite(str(disp_path), cv2.cvtColor(disp_map, cv2.COLOR_RGB2BGR))
    cv2.imwrite(str(rough_path), cv2.cvtColor(rough_map, cv2.COLOR_RGB2BGR))

    return normal_path, disp_path, rough_path

# Gradio function
def generate_maps(input_image, tile_size, seamless, use_cpu):
    # Save uploaded image to input folder
    input_path = INPUT_DIR / "uploaded_texture.png"
    input_img = np.array(input_image)
    cv2.imwrite(str(input_path), cv2.cvtColor(input_img, cv2.COLOR_RGB2BGR))

    # Process image
    normal_path, disp_path, rough_path = process_image(input_path, tile_size, seamless, use_cpu)

    # Read outputs for display
    normal_map = cv2.cvtColor(cv2.imread(str(normal_path)), cv2.COLOR_BGR2RGB)
    disp_map = cv2.cvtColor(cv2.imread(str(disp_path)), cv2.COLOR_BGR2RGB)
    rough_map = cv2.cvtColor(cv2.imread(str(rough_path)), cv2.COLOR_BGR2RGB)

    return input_image, normal_map, disp_map, rough_map

# Gradio interface
interface = gr.Interface(
    fn=generate_maps,
    inputs=[
        gr.Image(type="pil", label="Diffuse Texture"),
        gr.Slider(minimum=256, maximum=1024, step=64, value=512, label="Tile Size"),
        gr.Checkbox(label="Seamless", value=False),
        gr.Checkbox(label="Use CPU", value=False),
    ],
    outputs=[
        gr.Image(type="numpy", label="Input Diffuse Texture"),
        gr.Image(type="numpy", label="Normal Map"),
        gr.Image(type="numpy", label="Displacement Map"),
        gr.Image(type="numpy", label="Roughness Map"),
    ],
    title="Material Map Generator",
    description="Upload a diffuse texture to generate Normal, Displacement, and Roughness maps."
)

if __name__ == "__main__":
    interface.launch()