from collections import namedtuple
from typing import List

ModelInfo = namedtuple("ModelInfo", ["simple_name", "link", "description"])
model_info = {}

def register_model_info(
    full_names: List[str], simple_name: str, link: str, description: str
):
    info = ModelInfo(simple_name, link, description)

    for full_name in full_names:
        model_info[full_name] = info

def get_model_info(name: str) -> ModelInfo:
    if name in model_info:
        return model_info[name]
    else:
        # To fix this, please use `register_model_info` to register your model
        return ModelInfo(
            name, "", "Register the description at fastchat/model/model_registry.py"
        )

def get_model_description_md(model_list):
    model_description_md = """
| | | |
| ---- | ---- | ---- |
"""
    ct = 0
    visited = set()
    for i, name in enumerate(model_list):
        minfo = get_model_info(name)
        if minfo.simple_name in visited:
            continue
        visited.add(minfo.simple_name)
        one_model_md = f"[{minfo.simple_name}]({minfo.link}): {minfo.description}"

        if ct % 3 == 0:
            model_description_md += "|"
        model_description_md += f" {one_model_md} |"
        if ct % 3 == 2:
            model_description_md += "\n"
        ct += 1
    return model_description_md

# regist text-to-shape generation models

register_model_info(
    ["dreamfusion"],
    "DreamFusion",
    "https://dreamfusion3d.github.io/",
    "Text-to-3D using 2D Diffusion and SDS Loss",
)

register_model_info(
    ["fantasia3d"],
    "Fantasia3D",
    "https://fantasia3d.github.io/",
    "Disentangling Geometry and Appearance for High-quality Text-to-3D Content Creation",
)

register_model_info(
    ["instant3d"],
    "Instant3D",
    "https://jiahao.ai/instant3d/",
    "Fast Text-to-3D with Sparse-View Generation and Large Reconstruction Model",
)

register_model_info(
    ["latent-nerf"],
    "Latent-NeRF",
    "https://github.com/eladrich/latent-nerf",
    "Latent-NeRF for Shape-Guided Generation of 3D Shapes and Textures",
)

register_model_info(
    ["magic3d"],
    "Magic3D",
    "https://research.nvidia.com/labs/dir/magic3d/",
    "High-Resolution Text-to-3D Content Creation",
)

register_model_info(
    ["geodream"],
    "GeoDream",
    "https://mabaorui.github.io/GeoDream_page/",
    "Disentangling 2D and Geometric Priors for High-Fidelity and Consistent 3D Generation",
)


register_model_info(
    ["lucid-dreamer"],
    "LucidDreamer",
    "https://github.com/EnVision-Research/LucidDreamer",
    "Towards High-Fidelity Text-to-3D Generation via Interval Score Matching",
)

register_model_info(
    ["mvdream"],
    "MVDream",
    "https://github.com/bytedance/MVDream",
    "Multi-view Diffusion for 3D Generation",
)

register_model_info(
    ["point-e"],
    "Point·E",
    "https://github.com/openai/point-e",
    "A System for Generating 3D Point Clouds from Complex Prompts",
)

register_model_info(
    ["shap-e"],
    "Shap-E",
    "https://github.com/openai/shap-e",
    "Generating Conditional 3D Implicit Functions",
)

register_model_info(
    ["prolificdreamer"],
    "ProlificDreamer",
    "https://ml.cs.tsinghua.edu.cn/prolificdreamer/",
    "High-Fidelity and Diverse Text-to-3D Generation with Variational Score Distillation",
)

register_model_info(
    ["syncdreamer"],
    "SyncDreamer",
    "https://github.com/liuyuan-pal/SyncDreamer",
    "Generating Multiview-consistent Images from a Single-view Image",
)

register_model_info(
    ["sjc"],
    "Score Jacobian Chaining",
    "https://pals.ttic.edu/p/score-jacobian-chaining",
    "Lifting Pretrained 2D Diffusion Models for 3D Generation",
)

# register_model_info(
#     [],
#     "",
#     "",
#     "",
# )


## regist image-to-shape generation models
register_model_info(
    ["dreamgaussian"],
    "DreamGaussian",
    "https://github.com/dreamgaussian/dreamgaussian",
    "Generative Gaussian Splatting for Efficient 3D Content Creation",
)


register_model_info(
    ["wonder3d"],
    "Wonder3D",
    "https://github.com/xxlong0/Wonder3D",
    "Single Image to 3D using Cross-Domain Diffusion",
)

register_model_info(
    ["dreamcraft3d"],
    "Dreamcraft3d",
    "https://github.com/deepseek-ai/DreamCraft3D",
    "Hierarchical 3d generation with bootstrapped diffusion prior",
)

register_model_info(
    ["zero123"],
    "Zero-1-to-3",
    "https://github.com/cvlab-columbia/zero123",
    "Zero-shot One Image to 3D Object",
)

register_model_info(
    ["stable-zero123", "zero123-xl"],
    "Stable Zero123",
    "https://stability.ai/news/stable-zero123-3d-generation",
    "Quality 3D Object Generation from Single Images",
)

register_model_info(
    ["magic123"],
    "Magic123",
    "https://guochengqian.github.io/project/magic123/",
    "One Image to High-Quality 3D Object Generation Using Both 2D and 3D Diffusion Priors",
)
    
register_model_info(
    ["imagedream"],
    "ImageDream",
    "https://github.com/bytedance/ImageDream",
    "Image-Prompt Multi-view Diffusion for 3D Generation",
)
register_model_info(
    ["make-it-3d"],
    "Make-It-3D",
    "https://github.com/junshutang/Make-It-3D",
    "High-Fidelity 3D Creation from A Single Image with Diffusion Prior",
)

register_model_info(
    ["triplane-gaussian"],
    "TriplaneGaussian",
    "https://github.com/VAST-AI-Research/TriplaneGaussian",
    "Triplane Meets Gaussian Splatting: Fast and Generalizable Single-View 3D Reconstruction with Transformers",
)

register_model_info(
    ["free3d"],
    "Free3D",
    "https://github.com/lyndonzheng/Free3D",
    "Consistent Novel View Synthesis without 3D Representation",
)

register_model_info(
    ["escher-net"],
    "EscherNet",
    "https://github.com/kxhit/EscherNet",
    "A Generative Model for Scalable View Synthesis",
)

register_model_info(
    ["v3d"],
    "V3D",
    "https://github.com/heheyas/V3D",
    "Video Diffusion Models are Effective 3D Generators",
)

register_model_info(
    ["lgm"],
    "LGM",
    "https://github.com/3DTopia/LGM",
    "Large Multi-View Gaussian Model for High-Resolution 3D Content Creation",
)

register_model_info(
    ["gsgen"],
    "GSGEN",
    "https://github.com/gsgen3d/gsgen",
    "Text-to-3D using Gaussian Splatting",
)

register_model_info(
    ["openlrm"],
    "OpenLRM",
    "https://github.com/3DTopia/OpenLRM",
    "Open-Source Large Reconstruction Models",
)

register_model_info(
    ["hifa"],
    "HiFA",
    "https://github.com/JunzheJosephZhu/HiFA",
    "High-fidelity Text-to-3D Generation with Advanced Diffusion Guidance",
)

# register_model_info(
#     [],
#     "",
#     "",
#     "",
# )