import os
import pickle
from copy import deepcopy
from typing import Optional

import torch
from diffusers.models.activations import GEGLU, GELU
from cross_attn_hook import CrossAttentionExtractionHook
from ffn_hooker import FeedForwardHooker
from norm_attn_hook import NormHooker


# create dummy module for skip connection
class SkipConnection(torch.nn.Module):
    def __init__(self):
        super(SkipConnection, self).__init__()

    def forward(*args, **kwargs):
        return args[1]


def calculate_mask_sparsity(hooker, threshold: Optional[float] = None):
    total_num_lambs = 0
    num_activate_lambs = 0
    binary = getattr(
        hooker, "binary", None
    )  # if binary is not present, it will return None for ff_hooks
    for lamb in hooker.lambs:
        total_num_lambs += lamb.size(0)
        if binary:
            assert threshold is None, "threshold should be None for binary mask"
            num_activate_lambs += lamb.sum().item()
        else:
            assert (
                threshold is not None
            ), "threshold must be provided for non-binary mask"
            num_activate_lambs += (lamb >= threshold).sum().item()
    return total_num_lambs, num_activate_lambs, num_activate_lambs / total_num_lambs


def create_pipeline(
    pipe,
    model_id,
    device,
    torch_dtype,
    save_pt=None,
    lambda_threshold: float = 1,
    binary=True,
    epsilon=0.0,
    masking="binary",
    attn_name="attn",
    return_hooker=False,
    scope=None,
    ratio=None,
):
    """
    create the pipeline and optionally load the saved mask
    """
    pipe.to(device)
    pipe.vae.requires_grad_(False)
    if hasattr(pipe, "unet"):
        pipe.unet.requires_grad_(False)
    else:
        pipe.transformer.requires_grad_(False)
    if save_pt:
        # TODO should merge all the hooks checkpoint into one
        if "ff.pt" in save_pt or "attn.pt" in save_pt:
            save_pts = get_save_pts(save_pt)

            cross_attn_hooker = CrossAttentionExtractionHook(
                pipe,
                model_name=model_id,
                regex=".*",
                dtype=torch_dtype,
                head_num_filter=1,
                masking=masking,  # need to change to binary during inference
                dst=save_pts["attn"],
                epsilon=epsilon,
                attn_name=attn_name,
                binary=binary,
            )
            cross_attn_hooker.add_hooks(init_value=1)

            ff_hooker = FeedForwardHooker(
                pipe,
                regex=".*",
                dtype=torch_dtype,
                masking=masking,
                dst=save_pts["ff"],
                epsilon=epsilon,
                binary=binary,
            )
            ff_hooker.add_hooks(init_value=1)

            if os.path.exists(save_pts["norm"]):
                norm_hooker = NormHooker(
                    pipe,
                    regex=".*",
                    dtype=torch_dtype,
                    masking=masking,
                    dst=save_pts["norm"],
                    epsilon=epsilon,
                    binary=binary,
                )
                norm_hooker.add_hooks(init_value=1)
            else:
                norm_hooker = None

            _ = pipe("abc", num_inference_steps=1)
            cross_attn_hooker.load(device=device, threshold=lambda_threshold)
            ff_hooker.load(device=device, threshold=lambda_threshold)
            if norm_hooker:
                norm_hooker.load(device=device, threshold=lambda_threshold)
            if scope == "local" or scope == "global":
                if isinstance(ratio, float):
                    attn_hooker_ratio = ratio
                    ff_hooker_ratio = ratio
                else:
                    attn_hooker_ratio, ff_hooker_ratio = ratio[0], ratio[1]

                    if norm_hooker:
                        if len(ratio) < 3:
                            raise ValueError("Need to provide ratio for norm layer")
                        norm_hooker_ratio = ratio[2]

                cross_attn_hooker.binarize(scope, attn_hooker_ratio)
                ff_hooker.binarize(scope, ff_hooker_ratio)
                if norm_hooker:
                    norm_hooker.binarize(scope, norm_hooker_ratio)
            hookers = [cross_attn_hooker, ff_hooker]
            if norm_hooker:
                hookers.append(norm_hooker)

    if return_hooker:
        return pipe, hookers
    else:
        return pipe


def linear_layer_pruning(module, lamb):
    heads_to_keep = torch.nonzero(lamb).squeeze()
    if len(heads_to_keep.shape) == 0:
        # if only one head is kept, or none
        heads_to_keep = heads_to_keep.unsqueeze(0)

    modules_to_remove = [module.to_k, module.to_q, module.to_v]
    new_heads = int(lamb.sum().item())

    if new_heads == 0:
        return SkipConnection()

    for module_to_remove in modules_to_remove:
        # get head dimension
        inner_dim = module_to_remove.out_features // module.heads
        # place holder for the rows to keep
        rows_to_keep = torch.zeros(
            module_to_remove.out_features,
            dtype=torch.bool,
            device=module_to_remove.weight.device,
        )

        for idx in heads_to_keep:
            rows_to_keep[idx * inner_dim : (idx + 1) * inner_dim] = True

        # overwrite the inner projection with masked projection
        module_to_remove.weight.data = module_to_remove.weight.data[rows_to_keep, :]
        if module_to_remove.bias is not None:
            module_to_remove.bias.data = module_to_remove.bias.data[rows_to_keep]
        module_to_remove.out_features = int(sum(rows_to_keep).item())

    # Also update the output projection layer if available, (for FLUXSingleAttnProcessor2_0)
    # with column masking, dim 1
    if getattr(module, "to_out", None) is not None:
        module.to_out[0].weight.data = module.to_out[0].weight.data[:, rows_to_keep]
        module.to_out[0].in_features = int(sum(rows_to_keep).item())

    # update parameters in the attention module
    module.inner_dim = module.inner_dim // module.heads * new_heads
    try:
        module.query_dim = module.query_dim // module.heads * new_heads
        module.inner_kv_dim = module.inner_kv_dim // module.heads * new_heads
    except:
        pass
    module.cross_attention_dim = module.cross_attention_dim // module.heads * new_heads
    module.heads = new_heads
    return module


def ffn_linear_layer_pruning(module, lamb):
    lambda_to_keep = torch.nonzero(lamb).squeeze()
    if len(lambda_to_keep) == 0:
        return SkipConnection()

    num_lambda = len(lambda_to_keep)

    if isinstance(module.net[0], GELU):
        # linear layer weight remove before activation
        module.net[0].proj.weight.data = module.net[0].proj.weight.data[
            lambda_to_keep, :
        ]
        module.net[0].proj.out_features = num_lambda
        if module.net[0].proj.bias is not None:
            module.net[0].proj.bias.data = module.net[0].proj.bias.data[lambda_to_keep]

        update_act = GELU(module.net[0].proj.in_features, num_lambda)
        update_act.proj = module.net[0].proj
        module.net[0] = update_act
    elif isinstance(module.net[0], GEGLU):
        output_feature = module.net[0].proj.out_features
        module.net[0].proj.weight.data = torch.cat(
            [
                module.net[0].proj.weight.data[: output_feature // 2, :][
                    lambda_to_keep, :
                ],
                module.net[0].proj.weight.data[output_feature // 2 :][
                    lambda_to_keep, :
                ],
            ],
            dim=0,
        )
        module.net[0].proj.out_features = num_lambda * 2
        if module.net[0].proj.bias is not None:
            module.net[0].proj.bias.data = torch.cat(
                [
                    module.net[0].proj.bias.data[: output_feature // 2][lambda_to_keep],
                    module.net[0].proj.bias.data[output_feature // 2 :][lambda_to_keep],
                ]
            )

        update_act = GEGLU(module.net[0].proj.in_features, num_lambda * 2)
        update_act.proj = module.net[0].proj
        module.net[0] = update_act

    # proj weight after activation
    module.net[2].weight.data = module.net[2].weight.data[:, lambda_to_keep]
    module.net[2].in_features = num_lambda

    return module


# create SparsityLinear module
class SparsityLinear(torch.nn.Module):
    def __init__(self, in_features, out_features, lambda_to_keep, num_lambda):
        super(SparsityLinear, self).__init__()
        self.linear = torch.nn.Linear(in_features, num_lambda)
        self.out_features = out_features
        self.lambda_to_keep = lambda_to_keep

    def forward(self, x):
        x = self.linear(x)
        output = torch.zeros(
            x.size(0), self.out_features, device=x.device, dtype=x.dtype
        )
        output[:, self.lambda_to_keep] = x
        return output


def norm_layer_pruning(module, lamb):
    """
    Pruning the layer normalization layer for FLUX model
    """
    lambda_to_keep = torch.nonzero(lamb).squeeze()
    if len(lambda_to_keep) == 0:
        return SkipConnection()

    num_lambda = len(lambda_to_keep)

    # get num_features
    in_features = module.linear.in_features
    out_features = module.linear.out_features

    linear = SparsityLinear(in_features, out_features, lambda_to_keep, num_lambda)
    linear.linear.weight.data = module.linear.weight.data[lambda_to_keep]
    linear.linear.bias.data = module.linear.bias.data[lambda_to_keep]
    module.linear = linear
    return module


def get_save_pts(save_pt):
    if "ff.pt" in save_pt:
        ff_save_pt = deepcopy(save_pt)  # avoid in-place operation
        attn_save_pt = save_pt.split(os.sep)
        attn_save_pt[-1] = attn_save_pt[-1].replace("ff", "attn")
        attn_save_pt_output = os.sep.join(attn_save_pt)
        attn_save_pt[-1] = attn_save_pt[-1].replace("attn", "norm")
        norm_save_pt = os.sep.join(attn_save_pt)

        return {
            "ff": ff_save_pt,
            "attn": attn_save_pt_output,
            "norm": norm_save_pt,
        }
    else:
        attn_save_pt = deepcopy(save_pt)
        ff_save_pt = save_pt.split(os.sep)
        ff_save_pt[-1] = ff_save_pt[-1].replace("attn", "ff")
        ff_save_pt_output = os.sep.join(ff_save_pt)
        ff_save_pt[-1] = ff_save_pt[-1].replace("ff", "norm")
        norm_save_pt = os.sep.join(attn_save_pt)

        return {
            "ff": ff_save_pt_output,
            "attn": attn_save_pt,
            "norm": norm_save_pt,
        }


def save_img(pipe, g_cpu, steps, prompt, save_path):
    image = pipe(prompt, generator=g_cpu, num_inference_steps=steps)
    image["images"][0].save(save_path)