import torch
from guided_diffusion.script_util import create_model_and_diffusion, model_and_diffusion_defaults
import lpips
import clip


from encoders.modules import BERTEmbedder
from models.clipseg import CLIPDensePredT

from huggingface_hub import hf_hub_download

STEPS = 100
USE_DDPM = False
USE_DDIM = False
USE_CPU = False
CLIP_SEG_PATH = './weights/rd64-uni.pth'
CLIP_GUIDANCE = False

def make_models():
    segmodel = CLIPDensePredT(version='ViT-B/16', reduce_dim=64)
    segmodel.eval()

    # non-strict, because we only stored decoder weights (not CLIP weights)
    segmodel.load_state_dict(torch.load(CLIP_SEG_PATH, map_location=torch.device('cpu')), strict=False)
    # segmodel.save_pretrained("./weights/hf_clipseg")

    device = torch.device('cuda:0' if (torch.cuda.is_available() and not USE_CPU) else 'cpu')
    print('Using device:', device)

    hf_inpaint_path = hf_hub_download("alvanlii/rdm_inpaint", "inpaint.pt")
    model_state_dict = torch.load(hf_inpaint_path, map_location='cpu')

    # print(
    #     'hey',
    #     'clip_proj.weight' in model_state_dict, # True
    #     model_state_dict['input_blocks.0.0.weight'].shape[1] == 8, # True
    #     'external_block.0.0.weight' in model_state_dict # False
    # )

    model_params = {
        'attention_resolutions': '32,16,8',
        'class_cond': False,
        'diffusion_steps': 1000,
        'rescale_timesteps': True,
        'timestep_respacing': STEPS,  # Modify this value to decrease the number of
                                    # timesteps.
        'image_size': 32,
        'learn_sigma': False,
        'noise_schedule': 'linear',
        'num_channels': 320,
        'num_heads': 8,
        'num_res_blocks': 2,
        'resblock_updown': False,
        'use_fp16': False,
        'use_scale_shift_norm': False,
        'clip_embed_dim': 768,
        'image_condition': True,
        'super_res_condition': False,
    }

    if USE_DDPM:
        model_params['timestep_respacing'] = '1000'
    if USE_DDIM:
        if STEPS:
            model_params['timestep_respacing'] = 'ddim'+str(STEPS)
        else:
            model_params['timestep_respacing'] = 'ddim50'
    elif STEPS:
        model_params['timestep_respacing'] = str(STEPS)

    model_config = model_and_diffusion_defaults()
    model_config.update(model_params)

    if USE_CPU:
        model_config['use_fp16'] = False


    model, diffusion = create_model_and_diffusion(**model_config)
    model.load_state_dict(model_state_dict, strict=False)

    model.requires_grad_(CLIP_GUIDANCE).eval().to(device)

    if model_config['use_fp16']:
        model.convert_to_fp16()
    else:
        model.convert_to_fp32()

    def set_requires_grad(model, value):
        for param in model.parameters():
            param.requires_grad = value


    lpips_model = lpips.LPIPS(net="vgg").to(device)
    hf_kl_path = hf_hub_download("alvanlii/rdm_inpaint", "kl-f8.pt")

    ldm = torch.load(hf_kl_path, map_location="cpu")
    
    # torch.save(ldm, "./weights/hf_ldm")
    ldm.to(device)
    ldm.eval()
    ldm.requires_grad_(CLIP_GUIDANCE)
    set_requires_grad(ldm, CLIP_GUIDANCE)

    bert = BERTEmbedder(1280, 32)
    hf_bert_path = hf_hub_download("alvanlii/rdm_inpaint", 'bert.pt')
    # bert = BERTEmbedder.from_pretrained("alvanlii/rdm_bert")
    sd = torch.load(hf_bert_path, map_location="cpu")
    bert.load_state_dict(sd)
    # bert.save_pretrained("./weights/hf_bert")

    bert.to(device)
    bert.half().eval()
    set_requires_grad(bert, False)


    clip_model, clip_preprocess = clip.load('ViT-L/14', device=device, jit=False)
    clip_model.eval().requires_grad_(False)

    return segmodel, model, diffusion, ldm, bert, clip_model, model_params


if __name__ == "__main__":
    make_models()