import os
import torch
import torch.nn as nn
import numpy as np
from fastapi import FastAPI, Request
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from typing import List, Dict, Any

app = FastAPI()

# A "Realistic Mock" of the AnyTop Architecture
# Based on the SIGGRAPH 2025 paper "AnyTop: Character Animation Diffusion with Any Topology"
# Features: Enrichment Block, Skeletal Attention, Temporal Attention

class AnyTopMockModel(nn.Module):
    def __init__(self, feat_dim=6, latent_dim=128):
        super().__init__()
        # feat_dim: 3 for pos, 3 for rot (simplified)
        self.enrichment = nn.Linear(feat_dim + 3, latent_dim) # +3 for rest-pose embedding
        self.skeletal_transformer = nn.TransformerEncoderLayer(d_model=latent_dim, nhead=4, batch_first=True)
        self.temporal_transformer = nn.TransformerEncoderLayer(d_model=latent_dim, nhead=4, batch_first=True)
        self.output_head = nn.Linear(latent_dim, feat_dim)

    def forward(self, motion_noise, rest_pose, topology_dist):
        # motion_noise: (batch, frames, joints, feat_dim)
        # rest_pose: (batch, joints, 3)
        b, f, j, d = motion_noise.shape
        
        # Enrichment: Concat rest pose to features
        # Broadcast rest_pose to all frames
        rp_expanded = rest_pose.unsqueeze(1).expand(-1, f, -1, -1)
        x = torch.cat([motion_noise, rp_expanded], dim=-1)
        x = self.enrichment(x) # (b, f, j, latent_dim)

        # Skeletal Attention (per frame)
        x = x.view(b * f, j, -1)
        # In a real model, topology_dist would bias the attention matrix here
        x = self.skeletal_transformer(x)
        x = x.view(b, f, j, -1)

        # Temporal Attention (per joint)
        x = x.transpose(1, 2).reshape(b * j, f, -1)
        x = self.temporal_transformer(x)
        x = x.view(b, j, f, -1).transpose(1, 2)

        return self.output_head(x)

# Global Instance
model = AnyTopMockModel()

class InferenceRequest(BaseModel):
    skeleton: Dict[str, Any]
    steps: int = 10

@app.post("/infer")
async def infer(data: InferenceRequest):
    # Simulate Diffusion Iteration
    # Input: joints and hierarchy
    joints = data.skeleton.get("joints", [])
    num_joints = len(joints)
    num_frames = 60
    
    # Random Noise Start
    motion = torch.randn(1, num_frames, num_joints, 6)
    rest_pose = torch.tensor([j["pos"] for j in joints]).unsqueeze(0).float()
    
    # Mocking the denoising loop
    with torch.no_grad():
        for _ in range(min(data.steps, 5)): # Max 5 for speed in demo
            motion = model(motion, rest_pose, None)
    
    # Convert to list for JSON response
    # Shape: (frames, joints, XYZ)
    result = motion.squeeze(0).numpy()[..., :3].tolist()
    return {"motion": result}

app.mount("/static", StaticFiles(directory="static"), name="static")

@app.get("/", response_class=HTMLResponse)
async def index():
    with open("templates/index.html") as f:
        return f.read()

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)