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
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 compact spherical distribution
        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
        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).permute(2, 0, 1).unsqueeze(0)  # (1, 3, 256, 256)

        # Random projection points in [-1,1] for sampling
        proj = torch.rand(self.num, 2) * 2 - 1  # (N, 2)

        for _ in range(steps):
            # Normalize grid to [-1,1]
            grid = proj.unsqueeze(0).unsqueeze(0)  # (1, 1, N, 2)

            sampled = torch.nn.functional.grid_sample(
                target,
                grid,
                mode='bilinear',
                padding_mode='border',
                align_corners=True
            ).squeeze(1).squeeze(0)  # Squeeze H (1) and batch if needed, to (3, N)

            # Transpose to (N, 3)
            sampled = sampled.t()

            brightness = sampled.mean(dim=1)
            attraction = (brightness - brightness.mean()) * 0.015
            proj += attraction.unsqueeze(1) * (proj / (proj.norm(dim=1, keepdim=True) + 1e-6))

            # Color lerp to image average
            avg_color = target.mean(dim=[2,3]).squeeze()
            self.colors = torch.lerp(self.colors, avg_color.repeat(self.num, 1), 0.02)

        # Apply evolved projection to positions (scale to sphere radius)
        radii = self.positions.norm(dim=1, keepdim=True)
        self.positions[:, :2] = proj * radii

        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)

        # Projection to color shift
        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"):
        pos = self.positions.cpu().numpy()
        col = self.colors.cpu().numpy()
        opa = self.opacities.cpu().numpy()
        sca = np.log(self.scales.cpu().numpy())
        rot = self.rotations.cpu().numpy()
        nor = np.zeros_like(pos)  # Normals placeholder

        vertex_data = np.core.records.fromarrays([
            pos[:,0], pos[:,1], pos[:,2],
            nor[:,0], nor[:,1], nor[:,2],
            col[:,0], col[:,1], col[:,2],
            opa,
            sca[:,0], sca[:,1], sca[:,2],
            rot[:,0], rot[:,1], rot[:,2], rot[:,3]
        ], 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 = ""):
    if image is None:
        raise gr.Error("Please upload an image to proceed.")
    
    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")

    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 = "Persistent 3D representation evolved from image"
    if prompt:
        status += f" and conditioned on prompt: '{prompt}'"
    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 to generate an evolving 3D Gaussian splat representation.")
    
    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()