import gradio as gr
import torch
import numpy as np
from PIL import Image
import os
import plyfile
import open_clip

# Load OpenCLIP for prompt conditioning (runs efficiently on CPU or GPU)
model, _, preprocess = open_clip.create_model_and_transforms('ViT-B-32', pretrained='laion2b_s34b_b79K')
model.eval()
tokenizer = open_clip.get_tokenizer('ViT-B-32')

class PersistentCortex:
    def __init__(self, num_gaussians=8000):
        self.num = num_gaussians
        # Initialize a compact spherical distribution (similar to your working PLY)
        angles = torch.rand(num_gaussians, 2) * 2 * np.pi
        radius = torch.rand(num_gaussians).pow(1/3) * 0.6
        self.positions = torch.stack([
            radius * torch.sin(angles[:, 0]) * torch.cos(angles[:, 1]),
            radius * torch.sin(angles[:, 0]) * torch.sin(angles[:, 1]),
            radius * torch.cos(angles[:, 0])
        ], dim=1)
        
        self.scales = torch.exp(torch.randn(num_gaussians, 3) * -2.5 - 2.0)
        self.colors = torch.rand(num_gaussians, 3) * 0.7 + 0.3  # neutral gray start
        self.opacities = torch.sigmoid(torch.randn(num_gaussians) * 1.5 + 2.0)
        self.rotations = torch.nn.functional.normalize(torch.randn(num_gaussians, 4), dim=-1)

    def evolve_from_image(self, image: Image.Image, steps=800):
        img_array = np.array(image.resize((256, 256))) / 255.0
        target = torch.tensor(img_array, dtype=torch.float32)

        for _ in range(steps):
            # Simple orthographic projection to 2D for guidance
            proj = self.positions[:, :2].clone()
            proj = proj / proj.abs().max()  # normalize to [-1, 1]
            grid = proj.unsqueeze(0).unsqueeze(0)  # (1,1,N,2)

            sampled = torch.nn.functional.grid_sample(
                target.permute(2, 0, 1).unsqueeze(0),
                grid,
                mode='bilinear',
                padding_mode='border',
                align_corners=True
            ).squeeze(0).squeeze(0)  # (N, 3)

            # Attract positions toward brighter areas
            brightness = sampled.mean(dim=1)
            attraction = (brightness - brightness.mean()) * 0.015
            self.positions[:, :2] += attraction.unsqueeze(1) * proj

            # Gradually adopt average color from target
            self.colors = torch.lerp(self.colors, target.mean(dim=[0,1]).repeat(self.num, 1), 0.02)

        return self

    def condition_on_prompt(self, prompt: str):
        if not prompt.strip():
            return self
        text_tokens = tokenizer([prompt])
        with torch.no_grad():
            text_emb = model.encode_text(text_tokens).float()
            text_emb = text_emb / text_emb.norm(dim=-1, keepdim=True)

        # Simple linear projection to RGB shift (deterministic for reproducibility)
        proj = torch.nn.Linear(512, 3, bias=False)
        torch.nn.init.normal_(proj.weight, std=0.2)
        color_shift = proj(text_emb)

        self.colors = torch.clamp(self.colors + color_shift.repeat(self.num, 1) * 0.3, 0, 1)
        return self

    def export_ply(self, path="output.ply"):
        # Prepare vertex data matching gsplat.js expected properties
        zeros = np.zeros((self.num, 3), dtype=np.float32)
        log_scales = np.log(self.scales.cpu().numpy()).astype(np.float32)

        vertex_data = np.core.records.fromarrays([
            self.positions.cpu().numpy().astype(np.float32),
            zeros,  # normals (unused)
            self.colors.cpu().numpy().astype(np.float32),
            self.opacities.cpu().numpy().astype(np.float32),
            log_scales,
            self.rotations.cpu().numpy().astype(np.float32)
        ], names='x,y,z,nx,ny,nz,f_dc_0,f_dc_1,f_dc_2,opacity,scale_0,scale_1,scale_2,rot_0,rot_1,rot_2,rot_3')

        el = plyfile.PlyElement.describe(vertex_data, 'vertex')
        plyfile.PlyData([el], text=True).write(path)
        return os.path.abspath(path)

def process(image: Image.Image, prompt: str = ""):
    cortex = PersistentCortex(num_gaussians=8000)
    cortex.evolve_from_image(image, steps=800)
    if prompt:
        cortex.condition_on_prompt(prompt)

    ply_path = cortex.export_ply("/tmp/output.ply")

    # In Hugging Face Spaces, /tmp files are automatically served under /files/
    viewer_html = f"""
    <div id="viewer" style="width:100%; height:600px; background:#000;"></div>
    <script type="module">
    import * as SPLAT from "https://cdn.jsdelivr.net/npm/gsplat@latest";

    const container = document.getElementById('viewer');
    const canvas = document.createElement('canvas');
    canvas.style.width = '100%';
    canvas.style.height = '100%';
    container.appendChild(canvas);

    const scene = new SPLAT.Scene();
    const camera = new SPLAT.Camera();
    const renderer = new SPLAT.WebGLRenderer({{ canvas }});
    const controls = new SPLAT.OrbitControls(camera, canvas);

    controls.autoRotate = false;
    controls.enableDamping = true;
    controls.dampingFactor = 0.05;
    controls.rotateSpeed = 1.0;
    controls.zoomSpeed = 1.2;
    controls.panSpeed = 0.8;

    await SPLAT.Loader.LoadAsync("/files/output.ply", scene);

    function animate() {{
        controls.update();
        renderer.render(scene, camera);
        requestAnimationFrame(animate);
    }}
    animate();
    </script>
    """

    status = "3D Gaussian splat generated and evolved from your image"
    if prompt:
        status += f" with prompt conditioning: '{prompt}'."
    else:
        status += "."

    return viewer_html, status

with gr.Blocks(title="Persistent 3D Cortex Demo") as demo:
    gr.Markdown("# Persistent 3D Cortex – Interactive Demo")
    gr.Markdown("""
    Upload an image and optionally add a text prompt.  
    The system evolves a persistent 3D Gaussian splat representation influenced by the image content and prompt.
    """)
    
    with gr.Row():
        img_input = gr.Image(type="pil", label="Input Image")
        prompt_input = gr.Textbox(label="Prompt (e.g., 'shiny red apple', 'futuristic city')", placeholder="Optional text prompt")
    
    generate_btn = gr.Button("Generate & Evolve 3D", variant="primary")
    
    viewer_output = gr.HTML(label="Interactive 3D Viewer")
    status_output = gr.Textbox(label="Status")

    generate_btn.click(
        fn=process,
        inputs=[img_input, prompt_input],
        outputs=[viewer_output, status_output]
    )

demo.launch()