# heavily based on https://github.com/KohakuBlueleaf/LyCORIS/blob/main/lycoris/utils.py

from typing import *

import numpy as np

import torch
import torch.nn as nn
import torch.nn.functional as F

import torch.linalg as linalg

from tqdm import tqdm
from collections import OrderedDict


def make_sparse(t: torch.Tensor, sparsity=0.95):
    abs_t = torch.abs(t)
    np_array = abs_t.detach().cpu().numpy()
    quan = float(np.quantile(np_array, sparsity))
    sparse_t = t.masked_fill(abs_t < quan, 0)
    return sparse_t


def extract_conv(
        weight: Union[torch.Tensor, nn.Parameter],
        mode='fixed',
        mode_param=0,
        device='cpu',
        is_cp=False,
) -> Tuple[nn.Parameter, nn.Parameter]:
    weight = weight.to(device)
    out_ch, in_ch, kernel_size, _ = weight.shape

    U, S, Vh = linalg.svd(weight.reshape(out_ch, -1))

    if mode == 'fixed':
        lora_rank = mode_param
    elif mode == 'threshold':
        assert mode_param >= 0
        lora_rank = torch.sum(S > mode_param)
    elif mode == 'ratio':
        assert 1 >= mode_param >= 0
        min_s = torch.max(S) * mode_param
        lora_rank = torch.sum(S > min_s)
    elif mode == 'quantile' or mode == 'percentile':
        assert 1 >= mode_param >= 0
        s_cum = torch.cumsum(S, dim=0)
        min_cum_sum = mode_param * torch.sum(S)
        lora_rank = torch.sum(s_cum < min_cum_sum)
    else:
        raise NotImplementedError('Extract mode should be "fixed", "threshold", "ratio" or "quantile"')
    lora_rank = max(1, lora_rank)
    lora_rank = min(out_ch, in_ch, lora_rank)
    if lora_rank >= out_ch / 2 and not is_cp:
        return weight, 'full'

    U = U[:, :lora_rank]
    S = S[:lora_rank]
    U = U @ torch.diag(S)
    Vh = Vh[:lora_rank, :]

    diff = (weight - (U @ Vh).reshape(out_ch, in_ch, kernel_size, kernel_size)).detach()
    extract_weight_A = Vh.reshape(lora_rank, in_ch, kernel_size, kernel_size).detach()
    extract_weight_B = U.reshape(out_ch, lora_rank, 1, 1).detach()
    del U, S, Vh, weight
    return (extract_weight_A, extract_weight_B, diff), 'low rank'


def extract_linear(
        weight: Union[torch.Tensor, nn.Parameter],
        mode='fixed',
        mode_param=0,
        device='cpu',
) -> Tuple[nn.Parameter, nn.Parameter]:
    weight = weight.to(device)
    out_ch, in_ch = weight.shape

    U, S, Vh = linalg.svd(weight)

    if mode == 'fixed':
        lora_rank = mode_param
    elif mode == 'threshold':
        assert mode_param >= 0
        lora_rank = torch.sum(S > mode_param)
    elif mode == 'ratio':
        assert 1 >= mode_param >= 0
        min_s = torch.max(S) * mode_param
        lora_rank = torch.sum(S > min_s)
    elif mode == 'quantile' or mode == 'percentile':
        assert 1 >= mode_param >= 0
        s_cum = torch.cumsum(S, dim=0)
        min_cum_sum = mode_param * torch.sum(S)
        lora_rank = torch.sum(s_cum < min_cum_sum)
    else:
        raise NotImplementedError('Extract mode should be "fixed", "threshold", "ratio" or "quantile"')
    lora_rank = max(1, lora_rank)
    lora_rank = min(out_ch, in_ch, lora_rank)
    if lora_rank >= out_ch / 2:
        return weight, 'full'

    U = U[:, :lora_rank]
    S = S[:lora_rank]
    U = U @ torch.diag(S)
    Vh = Vh[:lora_rank, :]

    diff = (weight - U @ Vh).detach()
    extract_weight_A = Vh.reshape(lora_rank, in_ch).detach()
    extract_weight_B = U.reshape(out_ch, lora_rank).detach()
    del U, S, Vh, weight
    return (extract_weight_A, extract_weight_B, diff), 'low rank'


def extract_diff(
        base_model,
        db_model,
        mode='fixed',
        linear_mode_param=0,
        conv_mode_param=0,
        extract_device='cpu',
        use_bias=False,
        sparsity=0.98,
        small_conv=True,
        linear_only=False,
        extract_unet=True,
        extract_text_encoder=True,
):
    meta = OrderedDict()

    UNET_TARGET_REPLACE_MODULE = [
        "Transformer2DModel",
        "Attention",
        "ResnetBlock2D",
        "Downsample2D",
        "Upsample2D"
    ]
    UNET_TARGET_REPLACE_NAME = [
        "conv_in",
        "conv_out",
        "time_embedding.linear_1",
        "time_embedding.linear_2",
    ]
    if linear_only:
        UNET_TARGET_REPLACE_MODULE = ["Transformer2DModel", "Attention"]
        UNET_TARGET_REPLACE_NAME = [
            "conv_in",
            "conv_out",
        ]

    if not extract_unet:
        UNET_TARGET_REPLACE_MODULE = []
        UNET_TARGET_REPLACE_NAME = []

    TEXT_ENCODER_TARGET_REPLACE_MODULE = ["CLIPAttention", "CLIPMLP"]

    if not extract_text_encoder:
        TEXT_ENCODER_TARGET_REPLACE_MODULE = []

    LORA_PREFIX_UNET = 'lora_unet'
    LORA_PREFIX_TEXT_ENCODER = 'lora_te'

    def make_state_dict(
            prefix,
            root_module: torch.nn.Module,
            target_module: torch.nn.Module,
            target_replace_modules,
            target_replace_names=[]
    ):
        loras = {}
        temp = {}
        temp_name = {}

        for name, module in root_module.named_modules():
            if module.__class__.__name__ in target_replace_modules:
                temp[name] = {}
                for child_name, child_module in module.named_modules():
                    if child_module.__class__.__name__ not in {'Linear', 'LoRACompatibleLinear', 'Conv2d', 'LoRACompatibleConv'}:
                        continue
                    temp[name][child_name] = child_module.weight
            elif name in target_replace_names:
                temp_name[name] = module.weight

        for name, module in tqdm(list(target_module.named_modules())):
            if name in temp:
                weights = temp[name]
                for child_name, child_module in module.named_modules():
                    lora_name = prefix + '.' + name + '.' + child_name
                    lora_name = lora_name.replace('.', '_')
                    layer = child_module.__class__.__name__
                    if layer in {'Linear', 'LoRACompatibleLinear', 'Conv2d', 'LoRACompatibleConv'}:
                        root_weight = child_module.weight
                        if torch.allclose(root_weight, weights[child_name]):
                            continue

                    if layer == 'Linear' or layer == 'LoRACompatibleLinear':
                        weight, decompose_mode = extract_linear(
                            (child_module.weight - weights[child_name]),
                            mode,
                            linear_mode_param,
                            device=extract_device,
                        )
                        if decompose_mode == 'low rank':
                            extract_a, extract_b, diff = weight
                    elif layer == 'Conv2d' or layer == 'LoRACompatibleConv':
                        is_linear = (child_module.weight.shape[2] == 1
                                     and child_module.weight.shape[3] == 1)
                        if not is_linear and linear_only:
                            continue
                        weight, decompose_mode = extract_conv(
                            (child_module.weight - weights[child_name]),
                            mode,
                            linear_mode_param if is_linear else conv_mode_param,
                            device=extract_device,
                        )
                        if decompose_mode == 'low rank':
                            extract_a, extract_b, diff = weight
                        if small_conv and not is_linear and decompose_mode == 'low rank':
                            dim = extract_a.size(0)
                            (extract_c, extract_a, _), _ = extract_conv(
                                extract_a.transpose(0, 1),
                                'fixed', dim,
                                extract_device, True
                            )
                            extract_a = extract_a.transpose(0, 1)
                            extract_c = extract_c.transpose(0, 1)
                            loras[f'{lora_name}.lora_mid.weight'] = extract_c.detach().cpu().contiguous().half()
                            diff = child_module.weight - torch.einsum(
                                'i j k l, j r, p i -> p r k l',
                                extract_c, extract_a.flatten(1, -1), extract_b.flatten(1, -1)
                            ).detach().cpu().contiguous()
                            del extract_c
                    else:
                        continue
                    if decompose_mode == 'low rank':
                        loras[f'{lora_name}.lora_down.weight'] = extract_a.detach().cpu().contiguous().half()
                        loras[f'{lora_name}.lora_up.weight'] = extract_b.detach().cpu().contiguous().half()
                        loras[f'{lora_name}.alpha'] = torch.Tensor([extract_a.shape[0]]).half()
                        if use_bias:
                            diff = diff.detach().cpu().reshape(extract_b.size(0), -1)
                            sparse_diff = make_sparse(diff, sparsity).to_sparse().coalesce()

                            indices = sparse_diff.indices().to(torch.int16)
                            values = sparse_diff.values().half()
                            loras[f'{lora_name}.bias_indices'] = indices
                            loras[f'{lora_name}.bias_values'] = values
                            loras[f'{lora_name}.bias_size'] = torch.tensor(diff.shape).to(torch.int16)
                        del extract_a, extract_b, diff
                    elif decompose_mode == 'full':
                        loras[f'{lora_name}.diff'] = weight.detach().cpu().contiguous().half()
                    else:
                        raise NotImplementedError
            elif name in temp_name:
                weights = temp_name[name]
                lora_name = prefix + '.' + name
                lora_name = lora_name.replace('.', '_')
                layer = module.__class__.__name__

                if layer in {'Linear', 'LoRACompatibleLinear', 'Conv2d', 'LoRACompatibleConv'}:
                    root_weight = module.weight
                    if torch.allclose(root_weight, weights):
                        continue

                if layer == 'Linear' or layer == 'LoRACompatibleLinear':
                    weight, decompose_mode = extract_linear(
                        (root_weight - weights),
                        mode,
                        linear_mode_param,
                        device=extract_device,
                    )
                    if decompose_mode == 'low rank':
                        extract_a, extract_b, diff = weight
                elif layer == 'Conv2d' or layer == 'LoRACompatibleConv':
                    is_linear = (
                            root_weight.shape[2] == 1
                            and root_weight.shape[3] == 1
                    )
                    if not is_linear and linear_only:
                        continue
                    weight, decompose_mode = extract_conv(
                        (root_weight - weights),
                        mode,
                        linear_mode_param if is_linear else conv_mode_param,
                        device=extract_device,
                    )
                    if decompose_mode == 'low rank':
                        extract_a, extract_b, diff = weight
                    if small_conv and not is_linear and decompose_mode == 'low rank':
                        dim = extract_a.size(0)
                        (extract_c, extract_a, _), _ = extract_conv(
                            extract_a.transpose(0, 1),
                            'fixed', dim,
                            extract_device, True
                        )
                        extract_a = extract_a.transpose(0, 1)
                        extract_c = extract_c.transpose(0, 1)
                        loras[f'{lora_name}.lora_mid.weight'] = extract_c.detach().cpu().contiguous().half()
                        diff = root_weight - torch.einsum(
                            'i j k l, j r, p i -> p r k l',
                            extract_c, extract_a.flatten(1, -1), extract_b.flatten(1, -1)
                        ).detach().cpu().contiguous()
                        del extract_c
                else:
                    continue
                if decompose_mode == 'low rank':
                    loras[f'{lora_name}.lora_down.weight'] = extract_a.detach().cpu().contiguous().half()
                    loras[f'{lora_name}.lora_up.weight'] = extract_b.detach().cpu().contiguous().half()
                    loras[f'{lora_name}.alpha'] = torch.Tensor([extract_a.shape[0]]).half()
                    if use_bias:
                        diff = diff.detach().cpu().reshape(extract_b.size(0), -1)
                        sparse_diff = make_sparse(diff, sparsity).to_sparse().coalesce()

                        indices = sparse_diff.indices().to(torch.int16)
                        values = sparse_diff.values().half()
                        loras[f'{lora_name}.bias_indices'] = indices
                        loras[f'{lora_name}.bias_values'] = values
                        loras[f'{lora_name}.bias_size'] = torch.tensor(diff.shape).to(torch.int16)
                    del extract_a, extract_b, diff
                elif decompose_mode == 'full':
                    loras[f'{lora_name}.diff'] = weight.detach().cpu().contiguous().half()
                else:
                    raise NotImplementedError
        return loras

    text_encoder_loras = make_state_dict(
        LORA_PREFIX_TEXT_ENCODER,
        base_model[0], db_model[0],
        TEXT_ENCODER_TARGET_REPLACE_MODULE
    )

    unet_loras = make_state_dict(
        LORA_PREFIX_UNET,
        base_model[2], db_model[2],
        UNET_TARGET_REPLACE_MODULE,
        UNET_TARGET_REPLACE_NAME
    )
    print(len(text_encoder_loras), len(unet_loras))
    # the | will
    return (text_encoder_loras | unet_loras), meta


def get_module(
        lyco_state_dict: Dict,
        lora_name
):
    if f'{lora_name}.lora_up.weight' in lyco_state_dict:
        up = lyco_state_dict[f'{lora_name}.lora_up.weight']
        down = lyco_state_dict[f'{lora_name}.lora_down.weight']
        mid = lyco_state_dict.get(f'{lora_name}.lora_mid.weight', None)
        alpha = lyco_state_dict.get(f'{lora_name}.alpha', None)
        return 'locon', (up, down, mid, alpha)
    elif f'{lora_name}.hada_w1_a' in lyco_state_dict:
        w1a = lyco_state_dict[f'{lora_name}.hada_w1_a']
        w1b = lyco_state_dict[f'{lora_name}.hada_w1_b']
        w2a = lyco_state_dict[f'{lora_name}.hada_w2_a']
        w2b = lyco_state_dict[f'{lora_name}.hada_w2_b']
        t1 = lyco_state_dict.get(f'{lora_name}.hada_t1', None)
        t2 = lyco_state_dict.get(f'{lora_name}.hada_t2', None)
        alpha = lyco_state_dict.get(f'{lora_name}.alpha', None)
        return 'hada', (w1a, w1b, w2a, w2b, t1, t2, alpha)
    elif f'{lora_name}.weight' in lyco_state_dict:
        weight = lyco_state_dict[f'{lora_name}.weight']
        on_input = lyco_state_dict.get(f'{lora_name}.on_input', False)
        return 'ia3', (weight, on_input)
    elif (f'{lora_name}.lokr_w1' in lyco_state_dict
          or f'{lora_name}.lokr_w1_a' in lyco_state_dict):
        w1 = lyco_state_dict.get(f'{lora_name}.lokr_w1', None)
        w1a = lyco_state_dict.get(f'{lora_name}.lokr_w1_a', None)
        w1b = lyco_state_dict.get(f'{lora_name}.lokr_w1_b', None)
        w2 = lyco_state_dict.get(f'{lora_name}.lokr_w2', None)
        w2a = lyco_state_dict.get(f'{lora_name}.lokr_w2_a', None)
        w2b = lyco_state_dict.get(f'{lora_name}.lokr_w2_b', None)
        t1 = lyco_state_dict.get(f'{lora_name}.lokr_t1', None)
        t2 = lyco_state_dict.get(f'{lora_name}.lokr_t2', None)
        alpha = lyco_state_dict.get(f'{lora_name}.alpha', None)
        return 'kron', (w1, w1a, w1b, w2, w2a, w2b, t1, t2, alpha)
    elif f'{lora_name}.diff' in lyco_state_dict:
        return 'full', lyco_state_dict[f'{lora_name}.diff']
    else:
        return 'None', ()


def cp_weight_from_conv(
        up, down, mid
):
    up = up.reshape(up.size(0), up.size(1))
    down = down.reshape(down.size(0), down.size(1))
    return torch.einsum('m n w h, i m, n j -> i j w h', mid, up, down)


def cp_weight(
        wa, wb, t
):
    temp = torch.einsum('i j k l, j r -> i r k l', t, wb)
    return torch.einsum('i j k l, i r -> r j k l', temp, wa)


@torch.no_grad()
def rebuild_weight(module_type, params, orig_weight, scale=1):
    if orig_weight is None:
        return orig_weight
    merged = orig_weight
    if module_type == 'locon':
        up, down, mid, alpha = params
        if alpha is not None:
            scale *= alpha / up.size(1)
        if mid is not None:
            rebuild = cp_weight_from_conv(up, down, mid)
        else:
            rebuild = up.reshape(up.size(0), -1) @ down.reshape(down.size(0), -1)
        merged = orig_weight + rebuild.reshape(orig_weight.shape) * scale
        del up, down, mid, alpha, params, rebuild
    elif module_type == 'hada':
        w1a, w1b, w2a, w2b, t1, t2, alpha = params
        if alpha is not None:
            scale *= alpha / w1b.size(0)
        if t1 is not None:
            rebuild1 = cp_weight(w1a, w1b, t1)
        else:
            rebuild1 = w1a @ w1b
        if t2 is not None:
            rebuild2 = cp_weight(w2a, w2b, t2)
        else:
            rebuild2 = w2a @ w2b
        rebuild = (rebuild1 * rebuild2).reshape(orig_weight.shape)
        merged = orig_weight + rebuild * scale
        del w1a, w1b, w2a, w2b, t1, t2, alpha, params, rebuild, rebuild1, rebuild2
    elif module_type == 'ia3':
        weight, on_input = params
        if not on_input:
            weight = weight.reshape(-1, 1)
        merged = orig_weight + weight * orig_weight * scale
        del weight, on_input, params
    elif module_type == 'kron':
        w1, w1a, w1b, w2, w2a, w2b, t1, t2, alpha = params
        if alpha is not None and (w1b is not None or w2b is not None):
            scale *= alpha / (w1b.size(0) if w1b else w2b.size(0))
        if w1a is not None and w1b is not None:
            if t1:
                w1 = cp_weight(w1a, w1b, t1)
            else:
                w1 = w1a @ w1b
        if w2a is not None and w2b is not None:
            if t2:
                w2 = cp_weight(w2a, w2b, t2)
            else:
                w2 = w2a @ w2b
        rebuild = torch.kron(w1, w2).reshape(orig_weight.shape)
        merged = orig_weight + rebuild * scale
        del w1, w1a, w1b, w2, w2a, w2b, t1, t2, alpha, params, rebuild
    elif module_type == 'full':
        rebuild = params.reshape(orig_weight.shape)
        merged = orig_weight + rebuild * scale
        del params, rebuild

    return merged


def merge(
        base_model,
        lyco_state_dict,
        scale: float = 1.0,
        device='cpu'
):
    UNET_TARGET_REPLACE_MODULE = [
        "Transformer2DModel",
        "Attention",
        "ResnetBlock2D",
        "Downsample2D",
        "Upsample2D"
    ]
    UNET_TARGET_REPLACE_NAME = [
        "conv_in",
        "conv_out",
        "time_embedding.linear_1",
        "time_embedding.linear_2",
    ]
    TEXT_ENCODER_TARGET_REPLACE_MODULE = ["CLIPAttention", "CLIPMLP"]
    LORA_PREFIX_UNET = 'lora_unet'
    LORA_PREFIX_TEXT_ENCODER = 'lora_te'
    merged = 0

    def merge_state_dict(
            prefix,
            root_module: torch.nn.Module,
            lyco_state_dict: Dict[str, torch.Tensor],
            target_replace_modules,
            target_replace_names=[]
    ):
        nonlocal merged
        for name, module in tqdm(list(root_module.named_modules()), desc=f'Merging {prefix}'):
            if module.__class__.__name__ in target_replace_modules:
                for child_name, child_module in module.named_modules():
                    if child_module.__class__.__name__ not in {'Linear', 'LoRACompatibleLinear', 'Conv2d',
                                                               'LoRACompatibleConv'}:
                        continue
                    lora_name = prefix + '.' + name + '.' + child_name
                    lora_name = lora_name.replace('.', '_')

                    result = rebuild_weight(*get_module(
                        lyco_state_dict, lora_name
                    ), getattr(child_module, 'weight'), scale)
                    if result is not None:
                        merged += 1
                        child_module.requires_grad_(False)
                        child_module.weight.copy_(result)
            elif name in target_replace_names:
                lora_name = prefix + '.' + name
                lora_name = lora_name.replace('.', '_')

                result = rebuild_weight(*get_module(
                    lyco_state_dict, lora_name
                ), getattr(module, 'weight'), scale)
                if result is not None:
                    merged += 1
                    module.requires_grad_(False)
                    module.weight.copy_(result)

    if device == 'cpu':
        for k, v in tqdm(list(lyco_state_dict.items()), desc='Converting Dtype'):
            lyco_state_dict[k] = v.float()

    merge_state_dict(
        LORA_PREFIX_TEXT_ENCODER,
        base_model[0],
        lyco_state_dict,
        TEXT_ENCODER_TARGET_REPLACE_MODULE,
        UNET_TARGET_REPLACE_NAME
    )
    merge_state_dict(
        LORA_PREFIX_UNET,
        base_model[2],
        lyco_state_dict,
        UNET_TARGET_REPLACE_MODULE,
        UNET_TARGET_REPLACE_NAME
    )
    print(f'{merged} Modules been merged')
