import torch
import os
import sys
import safetensors.torch
from safetensors.torch import load_file, save_file

BLOCKID26=["BASE","IN00","IN01","IN02","IN03","IN04","IN05","IN06","IN07","IN08","IN09","IN10","IN11","M00","OUT00","OUT01","OUT02","OUT03","OUT04","OUT05","OUT06","OUT07","OUT08","OUT09","OUT10","OUT11"]
BLOCKID17=["BASE","IN01","IN02","IN04","IN05","IN07","IN08","M00","OUT03","OUT04","OUT05","OUT06","OUT07","OUT08","OUT09","OUT10","OUT11"]
BLOCKID12=["BASE","IN04","IN05","IN07","IN08","M00","OUT00","OUT01","OUT02","OUT03","OUT04","OUT05"]
BLOCKID20=["BASE","IN00","IN01","IN02","IN03","IN04","IN05","IN06","IN07","IN08","M00","OUT00","OUT01","OUT02","OUT03","OUT04","OUT05","OUT06","OUT07","OUT08"]
BLOCKNUMS = [12,17,20,26]
BLOCKIDS=[BLOCKID12,BLOCKID17,BLOCKID20,BLOCKID26]

BLOCKIDX=["encoder",
"lora_unet_input_blocks_0_",
"lora_unet_input_blocks_1_",
"lora_unet_input_blocks_2_",
"lora_unet_input_blocks_3_",
"lora_unet_input_blocks_4_",
"lora_unet_input_blocks_5_",
"lora_unet_input_blocks_6_",
"lora_unet_input_blocks_7_",
"lora_unet_input_blocks_8_",
"lora_unet_input_blocks_9_",
"lora_unet_input_blocks_10_",
"lora_unet_input_blocks_11_",
"lora_unet_middle_block_",
"lora_unet_output_blocks_0_",
"lora_unet_output_blocks_1_",
"lora_unet_output_blocks_2_",
"lora_unet_output_blocks_3_",
"lora_unet_output_blocks_4_",
"lora_unet_output_blocks_5_",
"lora_unet_output_blocks_6_",
"lora_unet_output_blocks_7_",
"lora_unet_output_blocks_8_",
"lora_unet_output_blocks_9_",
"lora_unet_output_blocks_10_",
"lora_unet_output_blocks_11_",
]

def to26(ratios):
    if len(ratios) == 26: return ratios
    ids = BLOCKIDS[BLOCKNUMS.index(len(ratios))]
    output = [0]*26
    for i, id in enumerate(ids):
        output[BLOCKID26.index(id)] = ratios[i]
    return output

device = "cuda"
filename = "d:/ai/models/lora/jrpg_characters_dora-000012.safetensors"
outfilename = "C:/ai/automatic/models/Lora/EnvyStarlighJRPGCharacters.safetensors"
weights = [0,0,0,0,0,0,1,1,1,0,0,0]

def load_lora(file_name):
    sd = load_file(file_name,device=device)

    for key in list(sd.keys()):
        #print(f"{key}: {type(sd[key])}")
        if type(sd[key]) == torch.Tensor:
            sd[key] = sd[key].to(device = device)
            if "resblocks" in key:
                isv2 = True

    return sd


def load_metadata_from_safetensors(safetensors_file: str) -> dict:
    """
    This method locks the file. see https://github.com/huggingface/safetensors/issues/164
    If the file isn't .safetensors or doesn't have metadata, return empty dict.
    """
    if os.path.splitext(safetensors_file)[1] != ".safetensors":
        return {}

    with safetensors.safe_open(safetensors_file, framework="pt", device="cpu") as f:
        metadata = f.metadata()
    if metadata is None:
        metadata = {}
    return metadata

def get_idx(key):
    for i, k in enumerate(BLOCKIDX):
        if k in key:
            return i
    return None

lora = load_lora(filename)
metadata = load_metadata_from_safetensors(filename)
weights = to26(weights)

for key in list(lora.keys()):
    idx = get_idx(key)
    if idx != None:
        print(f"{key}: {idx}")
    else:
        print(f"{key}: no index")
        continue

    if 'lora_up' in key or 'lora_down' in key:
        lora[key] *= weights[idx]



save_file(lora, outfilename)