import os
import torch
from PIL import Image

from lvdm.models.modules.lora import net_load_lora
from lvdm.utils.common_utils import instantiate_from_config


# ------------------------------------------------------------------------------------------
def load_model(config, ckpt_path, gpu_id=None, inject_lora=False, lora_scale=1.0, lora_path=''):
    print(f"Loading model from {ckpt_path}")
    
    # load sd
    pl_sd = torch.load(ckpt_path, map_location="cpu")
    try:
        global_step = pl_sd["global_step"]
        epoch = pl_sd["epoch"]
    except:
        global_step = -1
        epoch = -1
    
    # load sd to model
    try:
        sd = pl_sd["state_dict"]
    except:
        sd = pl_sd
    model = instantiate_from_config(config.model)
    model.load_state_dict(sd, strict=True)

    if inject_lora:
        net_load_lora(model, lora_path, alpha=lora_scale)
    
    # move to device & eval
    if gpu_id is not None:
        model.to(f"cuda:{gpu_id}")
    else:
        model.cuda()
    model.eval()

    return model, global_step, epoch


# ------------------------------------------------------------------------------------------
@torch.no_grad()
def get_conditions(prompts, model, batch_size, cond_fps=None,):
    
    if isinstance(prompts, str) or isinstance(prompts, int):
        prompts = [prompts]
    if isinstance(prompts, list):
        if len(prompts) == 1:
            prompts = prompts * batch_size
        elif len(prompts) == batch_size:
            pass
        else:
            raise ValueError(f"invalid prompts length: {len(prompts)}")
    else:
        raise ValueError(f"invalid prompts: {prompts}")
    assert(len(prompts) == batch_size)
    
    # content condition: text / class label
    c = model.get_learned_conditioning(prompts)
    key = 'c_concat' if model.conditioning_key == 'concat' else 'c_crossattn'
    c = {key: [c]}

    # temporal condition: fps
    if getattr(model, 'cond_stage2_config', None) is not None:
        if model.cond_stage2_key == "temporal_context":
            assert(cond_fps is not None)
            batch = {'fps': torch.tensor([cond_fps] * batch_size).long().to(model.device)}
            fps_embd = model.cond_stage2_model(batch)
            c[model.cond_stage2_key] = fps_embd
    
    return c


# ------------------------------------------------------------------------------------------
def make_model_input_shape(model, batch_size, T=None):
    image_size = [model.image_size, model.image_size] if isinstance(model.image_size, int) else model.image_size
    C = model.model.diffusion_model.in_channels
    if T is None:
        T = model.model.diffusion_model.temporal_length
    shape = [batch_size, C, T, *image_size]
    return shape


# ------------------------------------------------------------------------------------------
def custom_to_pil(x):
    x = x.detach().cpu()
    x = torch.clamp(x, -1., 1.)
    x = (x + 1.) / 2.
    x = x.permute(1, 2, 0).numpy()
    x = (255 * x).astype(np.uint8)
    x = Image.fromarray(x)
    if not x.mode == "RGB":
        x = x.convert("RGB")
    return x

def torch_to_np(x):
    # saves the batch in adm style as in https://github.com/openai/guided-diffusion/blob/main/scripts/image_sample.py
    sample = x.detach().cpu()
    sample = ((sample + 1) * 127.5).clamp(0, 255).to(torch.uint8)
    if sample.dim() == 5:
        sample = sample.permute(0, 2, 3, 4, 1)
    else:
        sample = sample.permute(0, 2, 3, 1)
    sample = sample.contiguous()
    return sample

def make_sample_dir(opt, global_step=None, epoch=None):
    if not getattr(opt, 'not_automatic_logdir', False):
        gs_str = f"globalstep{global_step:09}" if global_step is not None else "None"
        e_str = f"epoch{epoch:06}" if epoch is not None else "None"
        ckpt_dir = os.path.join(opt.logdir, f"{gs_str}_{e_str}")
        
        # subdir name
        if opt.prompt_file is not None:
            subdir = f"prompts_{os.path.splitext(os.path.basename(opt.prompt_file))[0]}"
        else:
            subdir = f"prompt_{opt.prompt[:10]}"
        subdir += "_DDPM" if opt.vanilla_sample else f"_DDIM{opt.custom_steps}steps"
        subdir += f"_CfgScale{opt.scale}"
        if opt.cond_fps is not None:
            subdir += f"_fps{opt.cond_fps}"
        if opt.seed is not None:
            subdir += f"_seed{opt.seed}"

        return os.path.join(ckpt_dir, subdir)
    else:
        return opt.logdir