import copy

import torch

_orig_torch_load = torch.load

def _torch_load(*args, **kwargs):
    kwargs.setdefault('weights_only', False)
    return _orig_torch_load(*args, **kwargs)

torch.load = _torch_load

from safetensors.torch import load_file
from diffusers import StableDiffusionPipeline, DDIMScheduler, UNet2DConditionModel


def get_module_kohya_state_dict(module, prefix: str, dtype: torch.dtype, adapter_name: str = "default"):
    kohya_ss_state_dict = {}
    for peft_key, weight in module.items():
        kohya_key = peft_key.replace("unet.base_model.model", prefix)
        kohya_key = kohya_key.replace("lora_A", "lora_down")
        kohya_key = kohya_key.replace("lora_B", "lora_up")
        kohya_key = kohya_key.replace(".", "_", kohya_key.count(".") - 2)
        kohya_ss_state_dict[kohya_key] = weight.to(dtype)
        # Set alpha parameter
        if "lora_down" in kohya_key:
            alpha_key = f'{kohya_key.split(".")[0]}.alpha'
            kohya_ss_state_dict[alpha_key] = torch.tensor(8).to(dtype)

    return kohya_ss_state_dict


def _make_pipe(vae, text_encoder, tokenizer, unet, scheduler):
    """Build a StableDiffusionPipeline around one specific UNet.

    The VAE / text encoder / tokenizer are shared (read-only at inference
    time), but every pipeline gets its OWN UNet module so the reverse and
    forward LoRA deltas are fused into separate weight sets — the same end
    state as the original copy.deepcopy(pipe) implementation.

    Everything stays on CPU while building: ZeroGPU rejects copy.deepcopy of
    pipelines whose tensors sit on a mix of devices, so devices are assigned
    only once, at the very end of loading (see load_consistency).
    """
    return StableDiffusionPipeline(
        vae=vae,
        text_encoder=text_encoder,
        tokenizer=tokenizer,
        unet=unet,
        scheduler=scheduler,
        safety_checker=None,
        feature_extractor=None,
        requires_safety_checker=False,
    )


# Load models (DM, CM)
def load_models(model_id,
                device,
                reverse_checkpoint,
                forward_checkpoint,
                r=64,
                w_embed_dim=0,
                teacher_checkpoint=None,
                dtype='fp32',
                ):
    # Diffusion
    # ------------------------------------------------------------
    dtype = torch.float32 if dtype == 'fp32' else torch.float16
    scheduler = DDIMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", clip_sample=False,
                              set_alpha_to_one=False)
    ldm_stable = StableDiffusionPipeline.from_pretrained(
        model_id, scheduler=scheduler, safety_checker=None,
        requires_safety_checker=False).to('cpu', dtype=dtype)
    ldm_stable.enable_attention_slicing()
    # ------------------------------------------------------------

    # Reverse consistency
    # ------------------------------------------------------------
    if w_embed_dim > 0:
        print(f'Forward CD is initialized with guidance embedding, dim {w_embed_dim}')
        unet = UNet2DConditionModel.from_pretrained(
            model_id, subfolder="unet",
            time_cond_proj_dim=w_embed_dim, low_cpu_mem_usage=False, device_map=None,
            attn_implementation="sdpa"
        ).to('cpu')
        if teacher_checkpoint is not None:
            print(f'Embedded model is loading from {teacher_checkpoint}')
            unet.load_state_dict(torch.load(teacher_checkpoint, map_location='cpu'))
            ldm_stable.unet = unet
            ldm_stable.to(dtype=dtype)
        else:
            print('PROVIDE TEACHER')
    else:
        unet = UNet2DConditionModel.from_pretrained(
            model_id, subfolder="unet", attn_implementation="sdpa"
        ).to('cpu')

    if reverse_checkpoint is not None:
        print(f'Reverse CD is loading from {reverse_checkpoint}')
        reverse_cons_model = _make_pipe(
            ldm_stable.vae, ldm_stable.text_encoder, ldm_stable.tokenizer,
            copy.deepcopy(unet), ldm_stable.scheduler)
        lora_weight = load_file(reverse_checkpoint)
        lora_state_dict = get_module_kohya_state_dict(lora_weight, "lora_unet", torch.float16)
        reverse_cons_model.load_lora_weights(lora_state_dict)
        reverse_cons_model.fuse_lora()
        reverse_cons_model.to(dtype=dtype)
    else:
        reverse_cons_model = None
    # ------------------------------------------------------------

    # Inverse consistency
    # ------------------------------------------------------------
    if forward_checkpoint is not None:
        print(f'Forward CD is loading from {forward_checkpoint}')
        forward_cons_model = _make_pipe(
            ldm_stable.vae, ldm_stable.text_encoder, ldm_stable.tokenizer,
            copy.deepcopy(unet), ldm_stable.scheduler)
        lora_weight = load_file(forward_checkpoint)
        lora_state_dict = get_module_kohya_state_dict(lora_weight, "lora_unet", torch.float16)
        forward_cons_model.load_lora_weights(lora_state_dict)
        forward_cons_model.fuse_lora()
        forward_cons_model.to(dtype=dtype)
    else:
        forward_cons_model = None
    # ------------------------------------------------------------

    return ldm_stable, reverse_cons_model, forward_cons_model