import gradio as gr
import torch
import numpy as np
from PIL import Image, ImageDraw
try:
    from spaces import GPU
except ImportError:
    # Define a no-op decorator if running locally
    def GPU(func):
        return func
    
import os
import re
import json
import argparse
from datetime import datetime
from inference import GenerativeInferenceModel, get_inference_configs, get_imagenet_labels

# Parse command line arguments
parser = argparse.ArgumentParser(description='Run Generative Inference Demo')
parser.add_argument('--port', type=int, default=7860, help='Port to run the server on')
args = parser.parse_args()

# Create model directories if they don't exist
os.makedirs("models", exist_ok=True)
os.makedirs("stimuli", exist_ok=True)
SAVED_RUNS_DIR = "saved_runs"
os.makedirs(SAVED_RUNS_DIR, exist_ok=True)

# Load ImageNet labels for prior-bias category dropdown
IMAGENET_LABELS = get_imagenet_labels()

# Check if running on Hugging Face Spaces
if "SPACE_ID" in os.environ:
    default_port = int(os.environ.get("PORT", 7860))
else:
    default_port = 8861  # Local default port

# Initialize model
model = GenerativeInferenceModel()

# Define example images and their parameters with updated values from the research
examples = [
    {
        "image": os.path.join("stimuli", "ArtGallery1.jpg"),
        "name": "ArtGallery1",
        "method": "Prior-Guided Drift Diffusion",
        "reverse_diff": {
            "model": "resnet50_robust",
            "layer": "layer3",
            "initial_noise": 0.5,
            "diffusion_noise": 0.002,
            "step_size": 0.1,
            "iterations": 501,
            "epsilon": 40.0
        },
        "inference_normalization": "off",
        "use_adaptive_eps": False,
        "use_adaptive_step": True,
        "mask_center_x": 0.0,
        "mask_center_y": -1.0,
        "mask_radius": 0.1,
        "mask_sigma": 0.2,
        "eps_max_mult": 30.0,
        "eps_min_mult": 1.0,
        "step_max_mult": 100.0,
        "step_min_mult": 1.0,
    },
    {
        "image": os.path.join("stimuli", "farm1.jpg"),
        "name": "farm1",
        "method": "Prior-Guided Drift Diffusion",
        "reverse_diff": {
            "model": "resnet50_robust",
            "layer": "all",
            "initial_noise": 0.0,
            "diffusion_noise": 0.02,
            "step_size": 1.0,
            "iterations": 501,
            "epsilon": 40.0
        },
        "inference_normalization": "off",
        "use_adaptive_eps": False,
        "use_adaptive_step": False,
        "mask_center_x": 0.0,
        "mask_center_y": 0.0,
        "mask_radius": 0.2,
        "mask_sigma": 0.3,
        "eps_max_mult": 300.0,
        "eps_min_mult": 1.0,
        "step_max_mult": 10.0,
        "step_min_mult": 1.0,
    },
    {
        "image": os.path.join("stimuli", "urbanoffice1.jpg"),
        "name": "UrbanOffice1",
        "method": "Prior-Guided Drift Diffusion",
        "reverse_diff": {
            "model": "resnet50_robust",
            "layer": "all",
            "initial_noise": 1.0,
            "diffusion_noise": 0.002,
            "step_size": 1.0,
            "iterations": 500,
            "epsilon": 40.0
        },
        "inference_normalization": "off",
        "use_adaptive_eps": False,
        "use_adaptive_step": True,
        "mask_center_x": 0.5,
        "mask_center_y": 0.0,
        "mask_radius": 0.2,
        "mask_sigma": 0.2,
        "eps_max_mult": 20.0,
        "eps_min_mult": 1.0,
        "step_max_mult": 50.0,
        "step_min_mult": 0.2,
    },
]

def _input_image_stem(image):
    """Return a safe filename stem from the input image: known name or 'user_img'."""
    if image is None:
        return "user_img"
    path = None
    if isinstance(image, str) and (os.path.isfile(image) or os.path.exists(image)):
        path = image
    if isinstance(image, dict) and image.get("path") and os.path.exists(image.get("path", "")):
        path = image["path"]
    if path:
        name = os.path.splitext(os.path.basename(path))[0]
        # Safe for filenames: alphanumeric, underscore, hyphen only; max length
        safe = re.sub(r"[^\w\-]", "_", name).strip("_") or "user_img"
        return safe[:80] if len(safe) > 80 else safe
    return "user_img"


def _get_image_path_for_stem(img):
    """Extract file path from Gradio image value (path string, dict with path, or PIL) for stem tracking."""
    if img is None:
        return ""
    if isinstance(img, str) and (os.path.isfile(img) or os.path.exists(img)):
        return img
    if isinstance(img, dict) and img.get("path"):
        p = img["path"]
        if isinstance(p, str) and os.path.exists(p):
            return p
    return ""


def _update_tracked_image_path(img):
    """Keep path only when it's a known stimulus (e.g. from stimuli/); else '' so stem is 'user_img'."""
    path = _get_image_path_for_stem(img)
    if path and "stimuli" in path:
        return path
    return ""


def _config_to_json_serializable(c):
    """Return a copy of config with only JSON-serializable values."""
    if isinstance(c, dict):
        return {k: _config_to_json_serializable(v) for k, v in c.items()}
    if isinstance(c, (list, tuple)):
        return [_config_to_json_serializable(x) for x in c]
    if isinstance(c, (bool, int, float, str, type(None))):
        return c
    if hasattr(c, "item"):  # e.g. numpy scalar
        return c.item()
    return str(c)


@GPU 
def run_inference(image, model_type, inference_type, eps_value, num_iterations, 
                 initial_noise=0.05, diffusion_noise=0.3, step_size=0.8, model_layer="layer3",
                 use_adaptive_eps=False, use_adaptive_step=False,
                 mask_center_x=0.0, mask_center_y=0.0, mask_radius=0.3, mask_sigma=0.2,
                 eps_max_mult=4.0, eps_min_mult=1.0, step_max_mult=4.0, step_min_mult=1.0,
                 use_biased_inference=False, biased_class_name="",
                 current_image_path=""):
    # Check if image is provided
    if image is None:
        return None, [], "Please upload an image before running inference.", None
    
    # Convert eps to float
    eps = float(eps_value)
    step_size_f = float(step_size)
    # Coerce checkbox values (Gradio can send bool, None, or other)
    use_adaptive_eps = bool(use_adaptive_eps) if use_adaptive_eps is not None else False
    use_adaptive_step = bool(use_adaptive_step) if use_adaptive_step is not None else False
    
    # Load inference configuration based on the selected type
    config = get_inference_configs(inference_type=inference_type, eps=eps, n_itr=int(num_iterations))
    
    # Handle Prior-Guided Drift Diffusion specific parameters
    if inference_type == "Prior-Guided Drift Diffusion":
        config['initial_inference_noise_ratio'] = float(initial_noise)
        config['diffusion_noise_ratio'] = float(diffusion_noise)
        config['step_size'] = step_size_f
        config['top_layer'] = model_layer
    
    # Inference normalization off (option removed from UI)
    config['inference_normalization'] = 'off'
    config['recognition_normalization'] = 'off'
    
    # Adaptive epsilon (Gaussian mask)
    if use_adaptive_eps:
        config['adaptive_epsilon'] = {
            'enabled': True,
            'base_epsilon': eps,
            'center_x': float(mask_center_x),
            'center_y': float(mask_center_y),
            'flat_radius': float(mask_radius),
            'sigma': float(mask_sigma),
            'max_multiplier': float(eps_max_mult),
            'min_multiplier': float(eps_min_mult),
        }
    else:
        config['adaptive_epsilon'] = None
    
    # Adaptive step size (same Gaussian mask location/size, different multipliers)
    if use_adaptive_step:
        config['adaptive_step_size'] = {
            'enabled': True,
            'base_step_size': step_size_f,
            'center_x': float(mask_center_x),
            'center_y': float(mask_center_y),
            'flat_radius': float(mask_radius),
            'sigma': float(mask_sigma),
            'max_multiplier': float(step_max_mult),
            'min_multiplier': float(step_min_mult),
        }
    else:
        config['adaptive_step_size'] = None
    
    # Prior bias: steer perception toward a target ImageNet class
    use_biased_inference = bool(use_biased_inference) if use_biased_inference is not None else False
    biased_class_name = (biased_class_name or "").strip() if biased_class_name else ""
    if use_biased_inference and biased_class_name:
        config['biased_inference'] = {'enable': True, 'class': biased_class_name}
    else:
        config['biased_inference'] = config.get('biased_inference') or {'enable': False, 'class': None}
    
    # Run generative inference
    result = model.inference(image, model_type, config)
    
    # Extract results based on return type
    if isinstance(result, tuple):
        # Old format returning (output_image, all_steps)
        output_image, all_steps = result
    else:
        # New format returning dictionary
        output_image = result['final_image']
        all_steps = result['steps']
    
    # Create animation frames
    frames = []
    for i, step_image in enumerate(all_steps):
        # Convert tensor to PIL image
        step_pil = Image.fromarray((step_image.permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8))
        frames.append(step_pil)
    
    # Convert the final output image to PIL
    final_image = Image.fromarray((output_image.permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8))
    
    # Always save GIF and config and offer as downloads (browser will ask where to save)
    save_status = ""
    files_for_download = None
    if frames:
        # Use tracked path when available (e.g. from Load Parameters); else derive from image (PIL loses path)
        stem = _input_image_stem(current_image_path if (current_image_path and current_image_path.strip()) else image)
        unique_id = f"{datetime.now().strftime('%Y%m%d_%H%M%S')}_{stem}"
        gif_path = os.path.join(SAVED_RUNS_DIR, f"{unique_id}.gif")
        config_path = os.path.join(SAVED_RUNS_DIR, f"{unique_id}_config.json")
        try:
            frames[0].save(
                gif_path,
                save_all=True,
                append_images=frames[1:],
                loop=0,
                duration=200,
            )
            save_config = {
                "model_type": model_type,
                "input_image_name": stem,
                **_config_to_json_serializable(config),
            }
            with open(config_path, "w") as f:
                json.dump(save_config, f, indent=2)
            files_for_download = [gif_path, config_path]
            save_status = "**Download results** — Use the links below to save the GIF and config to your device (your browser may ask where to save)."
        except Exception as e:
            save_status = f"Save failed: {e}"
    
    return final_image, frames, save_status, files_for_download

def _image_to_pil(img):
    """Convert Gradio image value (PIL, numpy, path, or dict) to PIL Image; return None if invalid."""
    if img is None:
        return None
    if isinstance(img, Image.Image):
        return img
    if isinstance(img, np.ndarray):
        return Image.fromarray(img.astype(np.uint8))
    if isinstance(img, dict) and "path" in img:
        return Image.open(img["path"]).convert("RGB")
    if isinstance(img, str) and os.path.exists(img):
        return Image.open(img).convert("RGB")
    return None


def _image_size(img):
    """Return (width, height) from Gradio image value, or (1, 1) if unknown."""
    pil = _image_to_pil(img)
    if pil is not None:
        return pil.size
    return (1, 1)


def image_click_to_center(evt: gr.SelectData):
    """Convert click (x, y) on image to normalized mask center (-1 to 1). Returns (center_x, center_y)."""
    if evt.index is None or (isinstance(evt.index, (list, tuple)) and len(evt.index) < 2):
        return 0.0, 0.0
    x, y = float(evt.index[0]), float(evt.index[1])
    w, h = _image_size(evt.value)
    if w <= 0 or h <= 0:
        return 0.0, 0.0
    center_x = (x / w) * 2.0 - 1.0
    center_y = (y / h) * 2.0 - 1.0
    center_x = max(-1.0, min(1.0, center_x))
    center_y = max(-1.0, min(1.0, center_y))
    return center_x, center_y


def draw_mask_overlay(image, center_x, center_y, radius):
    """Draw the Gaussian mask center and radius on a copy of the image. Returns PIL or None."""
    pil = _image_to_pil(image)
    if pil is None:
        return None
    img = pil.convert("RGB").copy()
    w, h = img.size
    cx_px = (float(center_x) + 1.0) / 2.0 * w
    cy_px = (float(center_y) + 1.0) / 2.0 * h
    radius_px = float(radius) * min(w, h) / 2.0
    draw = ImageDraw.Draw(img)
    # Circle for radius
    draw.ellipse(
        [cx_px - radius_px, cy_px - radius_px, cx_px + radius_px, cy_px + radius_px],
        outline="#E11D48",
        width=2 * max(2, min(w, h) // 150),
    )
    # Center dot
    r = max(2, min(w, h) // 80)
    draw.ellipse([cx_px - r, cy_px - r, cx_px + r, cy_px + r], fill="#E11D48", outline="#FFF")
    return img


# Helper function to apply example parameters (adaptive mask off by default unless example defines it)
def apply_example(example):
    rd = example["reverse_diff"]
    mcx = example.get("mask_center_x", 0.0)
    mcy = example.get("mask_center_y", 0.0)
    mrad = example.get("mask_radius", 0.3)
    example_image_pil = _image_to_pil(example["image"])
    mask_img = draw_mask_overlay(example_image_pil, mcx, mcy, mrad)
    return [
        example_image_pil,
        rd.get("model", "resnet50_robust"),
        example["method"],
        rd["epsilon"],
        rd["iterations"],
        1.0,    # fixed Drift Noise
        0.002,  # fixed Diffusion Noise
        rd["step_size"],
        rd["layer"],
        example.get("use_adaptive_eps", False),
        example.get("use_adaptive_step", False),
        mcx,
        mcy,
        example.get("mask_radius", 0.3),
        0.2,  # fixed field sigma (fall-off); not exposed in simple UI
        example.get("eps_max_mult", 4.0),
        example.get("eps_min_mult", 1.0),
        example.get("step_max_mult", 4.0),
        example.get("step_min_mult", 1.0),
        example.get("use_biased_inference", False),
        example.get("biased_class_name", ""),
        example["image"],  # keep path for save filename (e.g. UrbanOffice1 -> urbanoffice1)
        mask_img,
        gr.Group(visible=True),
    ]

# Define the interface
with gr.Blocks(title="Human Hallucination Prediction", css="""
.purple-button {
    background-color: #8B5CF6 !important;
    color: white !important;
    border: none !important;
}
.purple-button:hover {
    background-color: #7C3AED !important;
}
""") as demo:
    gr.Markdown("# Human Hallucination Prediction")
    gr.Markdown("**Predict what visual hallucinations humans may experience** using neural networks.")
    
    gr.Markdown("""
    **How to predict hallucinations:**
    1. **Select an example image** below and click "Load Parameters" to set the prediction settings
    2. **Click "Run Generative Inference"** to predict what hallucination humans may perceive
    3. **View the prediction**: Watch as the model reveals the perceptual structures it expects—matching what humans typically hallucinate
    4. **You can upload your own images**
    5. **You can download the results** as a .gif file together with the configs.json
    """)
    with gr.Row():
        with gr.Column(scale=1):
            # Inputs (track path so save filenames use stimulus name when from example)
            default_image_path = os.path.join("stimuli", "urbanoffice1.jpg")
            image_input = gr.Image(
                label="Input Image (click to set field center)",
                type="pil",
                value=_image_to_pil(default_image_path),
            )
            current_image_path_state = gr.State(value=default_image_path)
            mask_preview = gr.Image(
                label="Field center preview (click to set center — circle shows field)",
                type="pil",
                interactive=False,
            )
            # Run Inference button right below the image
            run_button = gr.Button("🔮 Predict Hallucination", variant="primary", elem_classes="purple-button")
            
            # Parameters toggle button
            params_button = gr.Button("⚙️ Play with the parameters", variant="secondary")
            
            # Parameters section (initially hidden)
            with gr.Group(visible=False) as params_section:
                # Keep core method defaults fixed (not user-editable in simple UI)
                model_choice = gr.State(value="resnet50_robust")
                inference_type = gr.State(value="Prior-Guided Drift Diffusion")
                
                with gr.Row():
                    eps_slider = gr.State(value=40.0)  # fixed epsilon; hidden in simple UI
                    iterations_slider = gr.Slider(minimum=1, maximum=600, value=500, step=1, label="Number of Iterations")
                
                # Keep these internal defaults fixed (not user-editable in simple UI)
                initial_noise_slider = gr.State(value=1.0)      # Drift Noise
                diffusion_noise_slider = gr.State(value=0.002)  # Diffusion Noise
                    
                with gr.Row():
                    step_size_slider = gr.Slider(minimum=0.01, maximum=2.0, value=1.0, step=0.01, 
                                               label="Prior strength")
                    layer_choice = gr.Dropdown(
                        choices=[("Higher", "all"), ("Intermediate", "layer3"), ("Lower", "layer1")],
                        value="all",
                        label="Hierarchy Level"
                    )
                
                with gr.Row():
                    use_adaptive_eps_check = gr.State(value=False)
                    use_adaptive_step_check = gr.Checkbox(value=True, label="Use adaptive step size (stronger/weaker updates by field)")
                with gr.Row():
                    mask_center_x_slider = gr.Slider(minimum=-1.0, maximum=1.0, value=0.5, step=0.05, label="Field center X")
                    mask_center_y_slider = gr.Slider(minimum=-1.0, maximum=1.0, value=0.0, step=0.05, label="Field center Y")
                with gr.Row():
                    mask_radius_slider = gr.Slider(minimum=0.01, maximum=1.0, value=0.2, step=0.01, label="Field radius")
                mask_sigma_slider = gr.State(value=0.2)
                eps_max_mult_slider = gr.State(value=1.0)
                eps_min_mult_slider = gr.State(value=1.0)
                with gr.Row():
                    step_max_mult_slider = gr.Slider(minimum=0.1, maximum=150.0, value=50.0, step=0.1, label="Prior strength at center")
                    step_min_mult_slider = gr.Slider(minimum=0.1, maximum=10.0, value=0.2, step=0.1, label="Prior strength at periphery")
                gr.Markdown("### 🎯 Prior bias")
                with gr.Row():
                    use_biased_inference_check = gr.Checkbox(value=False, label="Prior bias")
                    biased_class_dropdown = gr.Dropdown(
                        choices=[("— No bias —", "")] + [(label, label) for label in sorted(IMAGENET_LABELS)],
                        value="",
                        label="Target category",
                        allow_custom_value=False,
                        filterable=True,
                    )
        with gr.Column(scale=2):
            # Outputs
            output_image = gr.Image(label="Predicted Hallucination")
            output_frames = gr.Gallery(label="Hallucination Prediction Process", columns=5, rows=2)
            save_status_md = gr.Markdown(value="")
            download_files = gr.File(label="Download results (GIF + config)", file_count="multiple")
    
    # Examples section with integrated explanations
    gr.Markdown("## Examples")
    gr.Markdown("Select an example and click Load Parameters to apply its settings")
    
    # For each example, create a row with the image and explanation side by side
    for i, ex in enumerate(examples):
        with gr.Row():
            # Left column for the image
            with gr.Column(scale=1):
                # Display the example image
                example_img = gr.Image(value=ex["image"], type="filepath", label=f"{ex['name']}")
                load_btn = gr.Button(f"Load Parameters", variant="primary")
                
                # Set up the load button to apply this example's parameters
                load_btn.click(
                    fn=lambda ex=ex: apply_example(ex),
                    outputs=[
                        image_input, model_choice, inference_type,
                        eps_slider, iterations_slider,
                        initial_noise_slider, diffusion_noise_slider,
                        step_size_slider, layer_choice,
                        use_adaptive_eps_check, use_adaptive_step_check,
                        mask_center_x_slider, mask_center_y_slider,
                        mask_radius_slider, mask_sigma_slider,
                        eps_max_mult_slider, eps_min_mult_slider,
                        step_max_mult_slider, step_min_mult_slider,
                        use_biased_inference_check, biased_class_dropdown,
                        current_image_path_state,
                        mask_preview,
                        params_section,
                    ],
                )
            
            # Right column for the explanation
            with gr.Column(scale=2):
                gr.Markdown(f"### {ex['name']}")
                
        
        if i < len(examples) - 1:  # Don't add separator after the last example
            gr.Markdown("---")
    
    # Set up event handler for the main inference
    run_button.click(
        fn=run_inference,
        inputs=[
            image_input, model_choice, inference_type,
            eps_slider, iterations_slider,
            initial_noise_slider, diffusion_noise_slider,
            step_size_slider, layer_choice,
            use_adaptive_eps_check, use_adaptive_step_check,
            mask_center_x_slider, mask_center_y_slider,
            mask_radius_slider, mask_sigma_slider,
            eps_max_mult_slider, eps_min_mult_slider,
            step_max_mult_slider, step_min_mult_slider,
            use_biased_inference_check, biased_class_dropdown,
            current_image_path_state,
        ],
        outputs=[output_image, output_frames, save_status_md, download_files]
    )
    
    # Toggle parameters visibility
    def toggle_params():
        return gr.Group(visible=True)

    params_button.click(
        fn=toggle_params,
        outputs=[params_section],
    )

    # Click on input image or field preview to set Gaussian field center
    def _mask_preview_inputs():
        return [image_input, mask_center_x_slider, mask_center_y_slider, mask_radius_slider]

    image_input.select(
        fn=image_click_to_center,
        outputs=[mask_center_x_slider, mask_center_y_slider],
    )
    mask_preview.select(
        fn=image_click_to_center,
        outputs=[mask_center_x_slider, mask_center_y_slider],
    )
    # Update mask preview when image or center/radius change
    image_input.change(
        fn=draw_mask_overlay,
        inputs=_mask_preview_inputs(),
        outputs=[mask_preview],
    )
    # Keep tracked path for save filename: known stimulus name or clear so stem becomes 'user_img'
    image_input.change(
        fn=_update_tracked_image_path,
        inputs=[image_input],
        outputs=[current_image_path_state],
    )
    mask_center_x_slider.change(
        fn=draw_mask_overlay,
        inputs=_mask_preview_inputs(),
        outputs=[mask_preview],
    )
    mask_center_y_slider.change(
        fn=draw_mask_overlay,
        inputs=_mask_preview_inputs(),
        outputs=[mask_preview],
    )
    mask_radius_slider.change(
        fn=draw_mask_overlay,
        inputs=_mask_preview_inputs(),
        outputs=[mask_preview],
    )
    # Populate mask preview on load
    demo.load(
        fn=draw_mask_overlay,
        inputs=_mask_preview_inputs(),
        outputs=[mask_preview],
    )
    
    # About section
    gr.Markdown("""
    ## 🧠 About Hallucination Prediction
    
    This tool predicts human visual hallucinations using **generative inference** with adversarially robust neural networks. Robust models develop human-like perceptual biases, allowing them to forecast what perceptual structures humans will experience.
    
    ### Prediction Methods:
    
    **Prior-Guided Drift Diffusion** (Primary Method)  
    Starting from a noisy representation, the model converges toward what it expects to perceive—revealing predicted hallucinations
    
    **IncreaseConfidence**  
    Moving away from unlikely interpretations to reveal the most probable perceptual experience
    
    ### Parameters:
    - **Prior strength**: How strongly each step moves toward the model’s expected percept
    - **Number of Iterations**: How many prediction steps to perform
    - **Hierarchy Level**: Which perceptual level to predict from (early edges vs. high-level objects)
    - **Epsilon (Stimulus Fidelity)**: How closely the prediction must match the input stimulus
    
    ### Why Does This Work?
    
    Adversarially robust neural networks develop perceptual representations similar to human vision. When we use generative inference to reveal what these networks "expect" to see, it matches what humans hallucinate in ambiguous images—allowing us to predict human perception.
    
    **Developed by [Tahereh Toosi](https://toosi.github.io)**
    """)

# Launch the demo
if __name__ == "__main__":
    print(f"Starting server on port {args.port}")
    demo.launch(
        server_name="0.0.0.0",
        server_port=args.port,
        share=False,
        debug=True
    )