import sys
import spaces
sys.path.append("flash3d")  # Add the flash3d directory to the system path for importing local modules

from omegaconf import OmegaConf
import gradio as gr
import torch
import torchvision.transforms as TT
import torchvision.transforms.functional as TTF
from huggingface_hub import hf_hub_download
import numpy as np
from einops import rearrange

from networks.gaussian_predictor import GaussianPredictor
from util.vis3d import save_ply

def main():
    print("[INFO] Starting main function...")
    # Determine if CUDA (GPU) is available and set the device accordingly
    if torch.cuda.is_available():
        device = "cuda:0"
        print("[INFO] CUDA is available. Using GPU device.")
    else:
        device = "cpu"
        print("[INFO] CUDA is not available. Using CPU device.")

    # Download model configuration and weights from Hugging Face Hub
    print("[INFO] Downloading model configuration...")
    model_cfg_path = hf_hub_download(repo_id="einsafutdinov/flash3d", 
                                     filename="config_re10k_v1.yaml")
    print("[INFO] Downloading model weights...")
    model_path = hf_hub_download(repo_id="einsafutdinov/flash3d", 
                                 filename="model_re10k_v1.pth")

    # Load model configuration using OmegaConf
    print("[INFO] Loading model configuration...")
    cfg = OmegaConf.load(model_cfg_path)
    
    # Initialize the GaussianPredictor model with the loaded configuration
    print("[INFO] Initializing GaussianPredictor model...")
    model = GaussianPredictor(cfg)
    try:
        device = torch.device(device)
        model.to(device)  # Move the model to the specified device (CPU or GPU)
    except Exception as e:
        print(f"[ERROR] Failed to set device: {e}")
        raise
    
    # Load the pre-trained model weights
    print("[INFO] Loading model weights...")
    model.load_model(model_path)

    # Define transformation functions for image preprocessing
    pad_border_fn = TT.Pad((cfg.dataset.pad_border_aug, cfg.dataset.pad_border_aug))  # Padding to augment the image borders
    to_tensor = TT.ToTensor()  # Convert image to tensor

    # Function to check if an image is uploaded by the user
    def check_input_image(input_images):
        print("[DEBUG] Checking input images...")
        if not input_images or len(input_images) == 0:
            print("[ERROR] No images uploaded!")
            raise gr.Error("No images uploaded!")
        print("[INFO] Input images are valid.")

    # Function to preprocess the input images before passing them to the model
    def preprocess(images, padding_value):
        processed_images = []
        for image in images:
            # Resize and pad each image
            print("[DEBUG] Preprocessing image...")
            image = TTF.resize(image, (cfg.dataset.height, cfg.dataset.width), interpolation=TT.InterpolationMode.BICUBIC)
            pad_border_fn = TT.Pad((padding_value, padding_value))
            image = pad_border_fn(image)
            print("[INFO] Image preprocessing complete.")
            processed_images.append(image)
        return processed_images

    # Function to reconstruct the 3D model from the input images and export it as a PLY file
    @spaces.GPU(duration=120)  # Decorator to allocate a GPU for this function during execution
    def reconstruct_and_export(images, num_gauss):
        """
        Passes images through model, outputs reconstruction in form of a dict of tensors.
        """
        outputs_list = []
        for image in images:
            print("[DEBUG] Starting reconstruction and export...")
            # Convert the preprocessed image to a tensor and move it to the specified device
            image = to_tensor(image).to(device).unsqueeze(0)  # Add a batch dimension to the image tensor
            inputs = {
                ("color_aug", 0, 0): image,  # The input dictionary expected by the model
            }

            # Pass the image through the model to get the output
            print("[INFO] Passing image through the model...")
            outputs = model(inputs)  # Perform inference to get model outputs
            outputs_list.append(outputs)

        # Combine or process outputs from multiple images here if necessary
        # For now, we'll just save the first one for illustration
        gauss_means = outputs_list[0][('gauss_means', 0, 0)]
        if gauss_means.size(0) < num_gauss or gauss_means.size(0) % num_gauss != 0:
            adjusted_num_gauss = max(1, gauss_means.size(0) // (gauss_means.size(0) // num_gauss))
            print(f"[WARNING] Adjusting num_gauss from {num_gauss} to {adjusted_num_gauss} to avoid shape mismatch.")
            num_gauss = adjusted_num_gauss  # Adjust num_gauss to prevent errors during tensor reshaping

        # Debugging tensor shape
        print(f"[DEBUG] gauss_means tensor shape: {gauss_means.shape}")

        # Export the reconstruction to a PLY file
        print(f"[INFO] Saving output to {ply_out_path}...")
        save_ply(outputs_list[0], ply_out_path, num_gauss=num_gauss)  # Save the output 3D model to a PLY file
        print("[INFO] Reconstruction and export complete.")

        return ply_out_path  # Return the path to the saved PLY file
    
    # Path to save the output PLY file
    ply_out_path = f'./mesh.ply'

    # CSS styling for the Gradio interface
    css = """
        h1 {
            text-align: center;
            display:block;
        }
        """

    # Create the Gradio user interface
    with gr.Blocks(css=css) as demo:
        gr.Markdown(
            """
            # Flash3D
            """
        )
        with gr.Row(variant="panel"):
            with gr.Column(scale=1):
                with gr.Row():
                    # Input images component for the user to upload multiple images
                    input_images = gr.Images(
                        label="Input Images",
                        image_mode="RGBA",  # Accept RGBA images
                        sources="upload",  # Allow users to upload images
                        type="pil",  # The images are returned as PIL images
                        elem_id="content_images",
                        tool="editor",  # Optional, for editing images
                        multiple=True  # Allow multiple image uploads
                    )
                with gr.Row():
                    # Sliders for configurable parameters
                    num_gauss = gr.Slider(minimum=1, maximum=20, step=1, label="Number of Gaussians per Pixel", value=1)  # Slider to set the number of Gaussians per pixel
                    padding_value = gr.Slider(minimum=0, maximum=128, step=8, label="Padding Amount for Output Processing", value=32)  # Slider to set padding value
                with gr.Row():
                    # Button to trigger the generation process
                    submit = gr.Button("Generate", elem_id="generate", variant="primary")

                with gr.Row(variant="panel"):
                    # Examples panel to provide sample images for users
                    gr.Examples(
                        examples=[
                            './demo_examples/bedroom_01.png',
                            './demo_examples/kitti_02.png',
                            './demo_examples/kitti_03.png',
                            './demo_examples/re10k_04.jpg',
                            './demo_examples/re10k_05.jpg',
                            './demo_examples/re10k_06.jpg',
                        ],
                        inputs=[input_images],  # Load the example images into the input component
                        cache_examples=False,
                        label="Examples",  # Label for the examples section
                        examples_per_page=20,
                    )

                with gr.Row():
                    # Display the preprocessed images (after resizing and padding)
                    processed_images = gr.Gallery(label="Processed Images", interactive=False)  # Output component to show the processed images

            with gr.Column(scale=2):
                with gr.Row():
                    with gr.Tab("Reconstruction"):
                        # 3D model viewer to display the reconstructed model
                        output_model = gr.Model3D(
                            height=512,  # Height of the 3D model viewer
                            label="Output Model",
                            interactive=False  # The viewer is not interactive
                        )

        # Define the workflow for the Generate button
        submit.click(fn=check_input_image, inputs=[input_images]).success(
            fn=preprocess,
            inputs=[input_images, padding_value],  # Pass the input images and padding value to the preprocess function
            outputs=[processed_images],  # Output the processed images
        ).success(
            fn=reconstruct_and_export,
            inputs=[processed_images, num_gauss],  # Pass the processed images and number of Gaussians to the reconstruction function
            outputs=[output_model],  # Output the reconstructed 3D model
        )

    # Queue the requests to handle them sequentially (to avoid GPU resource conflicts)
    demo.queue(max_size=1)
    print("[INFO] Launching Gradio demo...")
    demo.launch(share=True)  # Launch the Gradio interface and allow public sharing

if __name__ == "__main__":
    print("[INFO] Running application...")
    main()