# -*- coding: utf-8 -*-
#
# @File:   app.py
# @Author: Haozhe Xie
# @Date:   2024-03-02 16:30:00
# @Last Modified by: Haozhe Xie
# @Last Modified at: 2024-03-03 10:39:25
# @Email:  root@haozhexie.com

import logging
import os
import torch
import gradio as gr
import subprocess
import urllib.request
import ssl
import sys

# Fix: ssl.SSLCertVerificationError: [SSL: CERTIFICATE_VERIFY_FAILED] certificate verify failed
ssl._create_default_https_context = ssl._create_unverified_context

sys.path.append(os.path.join(os.path.dirname(__file__), "citydreamer"))
# Import CityDreamer modules
# import citydreamer.model
# import citydreamer.inference


def setup_runtime_env():
    subprocess.call(["pip", "freeze"])
    ext_dir = os.path.join(os.path.dirname(__file__), "citydreamer", "extensions")
    for e in os.listdir(ext_dir):
        if not os.path.isdir(e):
            continue
        subprocess.call(["pip", "install", "."], workdir=os.path.join(ext_dir, e))


def get_models():
    if not os.path.exists("CityDreamer-Fgnd.pth"):
        urllib.request.urlretrieve(
            "https://huggingface.co/hzxie/city-dreamer/resolve/main/CityDreamer-Fgnd.pth",
            "CityDreamer-Fgnd.pth",
        )
    if not os.path.exists("CityDreamer-Bgnd.pth"):
        urllib.request.urlretrieve(
            "https://huggingface.co/hzxie/city-dreamer/resolve/main/CityDreamer-Bgnd.pth",
            "CityDreamer-Bgnd.pth",
        )

    bgm_ckpt = torch.load("CityDreamer-Bgnd.pth")
    fgm_ckpt = torch.load("CityDreamer-Fgnd.pth")
    bgm = citydreamer.model.GanCraftGenerator(bgm_ckpt["cfg"])
    fgm = citydreamer.model.GanCraftGenerator(fgm_ckpt["cfg"])
    if torch.cuda.is_available():
        fgm = torch.nn.DataParallel(fgm).cuda().eval()
        bgm = torch.nn.DataParallel(bgm).cuda().eval()

    return bgm, fgm


def get_generated_city(radius, altitude, azimuth):
    print(radius, altitude, azimuth)


def main(debug):
    title = "CityDreamer Demo 🏙️"
    with open("README.md", "r") as f:
        markdown = f.read()
        desc = markdown[markdown.rfind("---") + 3:]
    with open("ARTICLE.md", "r") as f:
        arti = f.read()

    app = gr.Interface(
        get_generated_city,
        [
            gr.Slider(
                128, 512, value=320, step=5, label="Camera Radius (m)"
            ),
            gr.Slider(
                256, 512, value=384, step=5, label="Camera Altitude (m)"
            ),
            gr.Slider(0, 360, value=180, step=5, label="Camera Azimuth (°)"),
        ],
        [gr.Image(type="numpy", label="Generated City")],
        title=title,
        description=desc,
        article=arti,
        allow_flagging="never",
    )
    app.queue(api_open=False)
    app.launch(debug=debug)


if __name__ == "__main__":
    logging.basicConfig(
        format="[%(levelname)s] %(asctime)s %(message)s", level=logging.INFO
    )
    logging.info("Compile CUDA extensions...")
    # setup_runtime_env()
    logging.info("Downloading pretrained models...")
    # fgm, bgm = get_models()
    logging.info("Starting the main application...")
    main(os.getenv("DEBUG") == "1")