import numpy as np
import torch
import torch.nn.functional as F
import yaml
from comfy.utils import ProgressBar
from tqdm import tqdm

from .modules.dense_motion import DenseMotionNetwork
from .modules.reconstruction_module import ReconstructionModule
from .modules.segmentation_module import SegmentationModule
from .modules.util import AntiAliasInterpolation2d
from .sync_batchnorm.replicate import DataParallelWithCallback


def face_parse_seg(
    source_image, face_parser_model, in_size=(512, 512), out_size=(64, 64)
):
    seg = F.interpolate(source_image, size=in_size)
    seg = (seg - face_parser_model.mean) / face_parser_model.std
    seg = torch.softmax(face_parser_model(seg)[0], dim=1)
    seg = F.interpolate(seg, size=out_size)
    return seg


class PartSwapGenerator(ReconstructionModule):
    def __init__(self, blend_scale=1, first_order_motion_model=False, **kwargs):
        super(PartSwapGenerator, self).__init__(**kwargs)
        if blend_scale == 1:
            self.blend_downsample = lambda x: x
        else:
            self.blend_downsample = AntiAliasInterpolation2d(1, blend_scale)

        if first_order_motion_model:
            self.dense_motion_network = DenseMotionNetwork(
                block_expansion=64,
                num_blocks=5,
                max_features=1024,
                num_kp=10,
                num_channels=3,
                estimate_occlusion_map=True,
                scale_factor=0.25,
            )
        else:
            self.dense_motion_network = None

    def forward(
        self,
        source_image,
        target_image,
        seg_target,
        seg_source,
        blend_mask,
        use_source_segmentation=False,
    ):
        # Encoding of source image
        enc_source = self.first(source_image)
        for i in range(len(self.down_blocks)):
            enc_source = self.down_blocks[i](enc_source)

        # Encoding of target image
        enc_target = self.first(target_image)
        for i in range(len(self.down_blocks)):
            enc_target = self.down_blocks[i](enc_target)

        output_dict = {}
        # Compute flow field for source image
        if self.dense_motion_network is None:
            segment_motions = self.segment_motion(seg_target, seg_source)
            segment_motions = segment_motions.permute(0, 1, 4, 2, 3)
            mask = seg_target["segmentation"].unsqueeze(2)
            deformation = (segment_motions * mask).sum(dim=1)
            deformation = deformation.permute(0, 2, 3, 1)
        else:
            motion = self.dense_motion_network.forward_partswap(
                source_image=source_image, seg_target=seg_target, seg_source=seg_source
            )
            deformation = motion["deformation"]

        # Deform source encoding according to the motion
        enc_source = self.deform_input(enc_source, deformation)

        if self.estimate_visibility:
            if self.dense_motion_network is None:
                visibility = seg_source["segmentation"][:, 1:].sum(
                    dim=1, keepdim=True
                ) * (
                    1
                    - seg_target["segmentation"][:, 1:]
                    .sum(dim=1, keepdim=True)
                    .detach()
                )
                visibility = 1 - visibility
            else:
                visibility = motion["visibility"]

            if (
                enc_source.shape[2] != visibility.shape[2]
                or enc_source.shape[3] != visibility.shape[3]
            ):
                visibility = F.interpolate(
                    visibility, size=enc_source.shape[2:], mode="bilinear"
                )
            enc_source = enc_source * visibility

        blend_mask = self.blend_downsample(blend_mask)
        # If source segmentation is provided use it should be deformed before blending
        if use_source_segmentation:
            blend_mask = self.deform_input(blend_mask, deformation)

        out = enc_target * (1 - blend_mask) + enc_source * blend_mask

        out = self.bottleneck(out)

        for i in range(len(self.up_blocks)):
            out = self.up_blocks[i](out)

        out = self.final(out)
        out = F.sigmoid(out)

        output_dict["prediction"] = out

        return output_dict


def partswap_inference(
    swap_indices: list[int],
    source_image,
    target_video,
    reconstruction_module: PartSwapGenerator,
    segmentation_module: SegmentationModule,
    use_source_seg: bool,
    face_parser_model=None,
    hard_edges=False,
    cpu=False,
) -> list:
    with torch.no_grad():
        predictions = []
        seg_targets = []
        if not cpu:
            source_image = source_image.cuda()
            target_video = target_video.cuda()
            print("Using GPU")

        seg_source = segmentation_module(source_image)
        # print(f"{seg_source['segmentation'].shape}")

        num_frames = target_video.shape[2]
        pbar = ProgressBar(num_frames)
        for frame_idx in tqdm(range(num_frames)):
            target_frame = target_video[:, :, frame_idx]
            seg_target = segmentation_module(target_frame)
            seg_targets.append(seg_target)

            if face_parser_model is not None:
                blend_mask = face_parse_seg(
                    source_image if use_source_seg else target_frame,
                    face_parser_model,
                    (512, 512),
                    (64, 64),
                )

            else:
                blend_mask = (
                    seg_source["segmentation"]
                    if use_source_seg
                    else seg_target["segmentation"]
                )

            blend_mask = blend_mask[:, swap_indices].sum(dim=1, keepdim=True)
            if hard_edges:
                blend_mask = (blend_mask > 0.5).type(blend_mask.type())
            # print(f"{blend_mask.shape=}")
            out = reconstruction_module(
                source_image,
                target_frame,
                seg_source=seg_source,
                seg_target=seg_target,
                blend_mask=blend_mask,
                use_source_segmentation=use_source_seg,
            )

            predictions.append(
                np.transpose(out["prediction"].data.cpu().numpy(), [0, 2, 3, 1])[0]
            )

            pbar.update_absolute(frame_idx, num_frames)

        return seg_source, seg_targets, predictions


def partial_state_dict_load(module, state_dict):
    own_state = module.state_dict()
    for name, param in state_dict.items():
        if name not in own_state:
            continue

        if isinstance(param, torch.nn.Parameter):
            # backwards compatibility for serialized parameters
            param = param.data
        own_state[name].copy_(param)


def load_reconstruction_module(module, checkpoint):
    if "generator" in checkpoint:
        partial_state_dict_load(module, checkpoint["generator"])
    else:
        module.load_state_dict(checkpoint["reconstruction_module"])


def load_segmentation_module(module, checkpoint):
    if "kp_detector" in checkpoint:
        partial_state_dict_load(module, checkpoint["kp_detector"])
        module.state_dict()["affine.weight"].copy_(
            checkpoint["kp_detector"]["jacobian.weight"]
        )
        module.state_dict()["affine.bias"].copy_(
            checkpoint["kp_detector"]["jacobian.bias"]
        )
        module.state_dict()["shift.weight"].copy_(
            checkpoint["kp_detector"]["kp.weight"]
        )
        module.state_dict()["shift.bias"].copy_(checkpoint["kp_detector"]["kp.bias"])
        if "semantic_seg.weight" in checkpoint["kp_detector"]:
            module.state_dict()["segmentation.weight"].copy_(
                checkpoint["kp_detector"]["semantic_seg.weight"]
            )
            module.state_dict()["segmentation.bias"].copy_(
                checkpoint["kp_detector"]["semantic_seg.bias"]
            )
        else:
            print("Segmentation part initialized at random.")
    else:
        module.load_state_dict(checkpoint["segmentation_module"])


def load_partswap_checkpoint(
    config_path: str,
    checkpoint_path: str,
    blend_scale: float,
    use_fomm: bool,
    cpu=False,
):
    with open(config_path) as f:
        config = yaml.full_load(f)

    reconstruction_module = PartSwapGenerator(
        blend_scale=blend_scale,
        first_order_motion_model=use_fomm,
        **config["model_params"]["reconstruction_module_params"],
        **config["model_params"]["common_params"],
    )

    segmentation_module = SegmentationModule(
        **config["model_params"]["segmentation_module_params"],
        **config["model_params"]["common_params"],
    )

    if cpu:
        checkpoint = torch.load(checkpoint_path, map_location=torch.device("cpu"))
    else:
        segmentation_module.cuda()
        checkpoint = torch.load(checkpoint_path)

    load_reconstruction_module(reconstruction_module, checkpoint)
    load_segmentation_module(segmentation_module, checkpoint)

    if not cpu:
        reconstruction_module = DataParallelWithCallback(reconstruction_module)
        segmentation_module = DataParallelWithCallback(segmentation_module)

    reconstruction_module.eval()
    segmentation_module.eval()

    return reconstruction_module, segmentation_module