import os
os.system("python -m pip install --upgrade pip")
import gradio as gr
import numpy as np
import os
import torch
os.system('pip install --upgrade --no-cache-dir gdown')
import openai

from PIL import Image

from utils.logger import get_logger
from config.defaults import get_config
from inference import preprocess, run_one_inference
from models.build import build_model
from argparse import Namespace
import gdown

# Gradio로 입력 폼을 생성합니다.
def summarize_and_classify(text):
    # 상황을 GPT-3.5 모델을 사용하여 분류하고 대응 방안을 생성합니다.
    response = openai.ChatCompletion.create(
        model="gpt-3.5-turbo",
        messages=[
            {"role": "system", "content": "고객 취향을 듣고, 어떤 인테리어가 어울릴지 추천해줘."},
            {"role": "user", "content": text},
        ],
        max_tokens=300,
        temperature=0.4
    )

    return response.choices[0].message["content"]

def down_ckpt(model_cfg, ckpt_dir):
    model_ids = [
        ['src/config/mp3d.yaml', '1CCEEwY12EkMrokF6TEeGjMao5SEl26FQ'],
        ['src/config/zind.yaml', '1ie4KQsB3uhd4ghZ7OTYt-XYbVRJ8w4ce'],
        ['src/config/pano.yaml', '1Ria7Bl4THDqQ3z8pL6bIGUmkTgtEzsuK'],
        ['src/config/s2d3d.yaml', '1P_EK8SreFCmA_LLScWCcvvpn-0l3sU1m'],
        ['src/config/ablation_study/full.yaml', '16wQFUC76-tfVlGes_avh7QwE9rcBrDXX']
    ]

    for model_id in model_ids:
        if model_id[0] != model_cfg:
            continue
        path = os.path.join(ckpt_dir, 'best.pkl')
        if not os.path.exists(path):
            logger.info(f"Downloading {model_id}")
            os.makedirs(ckpt_dir, exist_ok=True)
            gdown.download(f"https://drive.google.com/uc?id={model_id[1]}", path, False)
            gdown.download(f"https://drive.google.com/uc?id={model_id[1]}", path, False)


def greet(img_path, pre_processing, weight_name, post_processing, visualization, mesh_format, mesh_resolution):
    args.pre_processing = pre_processing
    args.post_processing = post_processing
    if weight_name == 'mp3d':
        model = mp3d_model
    elif weight_name == 'zind':
        model = zind_model
    else:
        logger.error("unknown pre-trained weight name")
        raise NotImplementedError

    img_name = os.path.basename(img_path).split('.')[0]
    img = np.array(Image.open(img_path).resize((1024, 512), Image.Resampling.BICUBIC))[..., :3]

    vp_cache_path = 'src/demo/default_vp.txt'
    if args.pre_processing:
        vp_cache_path = os.path.join('src/output', f'{img_name}_vp.txt')
        logger.info("pre-processing ...")
        img, vp = preprocess(img, vp_cache_path=vp_cache_path)

    img = (img / 255.0).astype(np.float32)
    run_one_inference(img, model, args, img_name,
                      logger=logger, show=False,
                      show_depth='depth-normal-gradient' in visualization,
                      show_floorplan='2d-floorplan' in visualization,
                      mesh_format=mesh_format, mesh_resolution=int(mesh_resolution))

    return [os.path.join(args.output_dir, f"{img_name}_pred.png"),
            os.path.join(args.output_dir, f"{img_name}_3d{mesh_format}"),
            os.path.join(args.output_dir, f"{img_name}_3d{mesh_format}"),
            vp_cache_path,
            os.path.join(args.output_dir, f"{img_name}_pred.json")]


def get_model(args):
    config = get_config(args)
    down_ckpt(args.cfg, config.CKPT.DIR)
    if ('cuda' in args.device or 'cuda' in config.TRAIN.DEVICE) and not torch.cuda.is_available():
        logger.info(f'The {args.device} is not available, will use cpu ...')
        config.defrost()
        args.device = "cpu"
        config.TRAIN.DEVICE = "cpu"
        config.freeze()
    model, _, _, _ = build_model(config, logger)
    return model


if __name__ == '__main__':
    logger = get_logger()
    args = Namespace(device='cuda', output_dir='src/output', visualize_3d=False, output_3d=True)
    os.makedirs(args.output_dir, exist_ok=True)

    args.cfg = 'src/config/mp3d.yaml'
    mp3d_model = get_model(args)

    args.cfg = 'src/config/zind.yaml'
    zind_model = get_model(args)

    description = "(주)플랜포스_인트로우 3D 스캐닝 고도화"

    demo = gr.Interface(fn=greet,
                        inputs=[gr.Image(type='filepath', label='input rgb panorama', value='src/demo/demo.jpg'),
                                gr.Checkbox(label='pre-processing', value=True),
                                gr.Radio(['mp3d', 'zind'],

                                         value='mp3d',
         visible=False),
                                gr.Radio(['manhattan', 'atalanta', 'original'],

                                         value='manhattan',
         visible=False),
                                gr.CheckboxGroup(['depth-normal-gradient', '2d-floorplan'],

                                                 value=['2d-floorplan'],
         visible=False),
                                gr.Radio(['.gltf', '.obj', '.glb'],
                                        
                                         value='.obj',
         visible=False),
                                gr.Radio(['128', '256', '512', '1024'],

                                         value='128',
         visible=False),
                                ],
                        outputs=[gr.Image(label='predicted result 2d-visualization', type='filepath'),
                                 gr.Model3D(label='3d mesh reconstruction', clear_color=[1.0, 1.0, 1.0, 1.0]),
                                 gr.File(label='3d mesh')
                                 
                                 # gr.File(label='vanishing point information'),
                                 # gr.File(label='layout json')
                                ],
                        
                        title='room3d', allow_flagging="never", cache_examples=False, description=description)

    demo.launch(debug=True, enable_queue=False)