import os
import sys
import torch
import torch.nn as nn
import numpy as np
import inspect

# --- PROACTIVE PYTHON 3.12 / NUMPY 2.0 PATCH ---
if not hasattr(np, 'bool'):
    np.bool = np.bool_
if not hasattr(np, 'float'):
    np.float = np.float64
if not hasattr(np, 'int'):
    np.int = np.int_
if not hasattr(np, 'complex'):
    np.complex = complex
if not hasattr(np, 'object'):
    np.object = object
if not hasattr(np, 'unicode'):
    np.unicode = str
if not hasattr(np, 'str'):
    np.str = str

if not hasattr(inspect, 'getargspec'):
    inspect.getargspec = inspect.getfullargspec
# -----------------------------------------------

# Path setup based on your layout
# Download from https://github.com/yfeng95/DECA
DECA_DIR = os.path.abspath("DECA")
sys.path.insert(0, DECA_DIR)

from decalib.deca import DECA
from decalib.utils.config import cfg as deca_cfg

# --- THE ULTIMATE PYTORCH3D BYPASS ---
DECA._setup_renderer = lambda self, *args, **kwargs: None
# -------------------------------------

# ---------------------------------------------------------
# The Custom Rust Wrapper
# ---------------------------------------------------------
class DecaVisionONNXWrapper(nn.Module):
    def __init__(self, deca_model):
        super().__init__()
        # Extract just the ResNet50 vision encoder from DECA
        self.encoder = deca_model.E_flame

    def forward(self, images):
        parameters = self.encoder(images)
        flame_vector = parameters[:, 0:150]
        return flame_vector

# ---------------------------------------------------------
def main():
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    print(f"Loading DECA on {device.upper()}...")

    # Initialize standard DECA
    deca_cfg.model.use_tex = False
    deca_cfg.rasterizer_type = 'pytorch3d'

    deca = DECA(config=deca_cfg, device=device)
    deca.eval()

    print("Wrapping model for seamless 150-dim output...")
    wrapper = DecaVisionONNXWrapper(deca).to(device)
    wrapper.eval()

    print("Tracing and exporting to ONNX...")
    dummy_input = torch.randn(1, 3, 224, 224).to(device)

    torch.onnx.export(
        wrapper,
        dummy_input,
        "deca_vision.onnx",
        export_params=True,
        opset_version=17,
        do_constant_folding=True,
        input_names=['image_input'],
        output_names=['flame_vector'],
        dynamic_axes={
            'image_input': {0: 'batch_size'},
            'flame_vector': {0: 'batch_size'}
        }
    )

    print("\nSuccess! 'deca_vision.onnx' is fully baked.")

if __name__ == "__main__":
    main()