from typing import Any, Dict, List, Optional

try:
    from spatial_correlation_sampler import SpatialCorrelationSampler
except ModuleNotFoundError:
    from ptlflow.utils.correlation import (
        IterSpatialCorrelationSampler as SpatialCorrelationSampler,
    )
import torch
import torch.nn as nn
import torch.nn.functional as F

from ptlflow.utils.registry import register_model
from .warp import WarpingLayer
from ..base_model.base_model import BaseModel


class FeatureExtractor(nn.Module):
    def __init__(self):
        super(FeatureExtractor, self).__init__()

        leaky_relu = nn.LeakyReLU(0.1, inplace=True)

        self.convs = nn.ModuleList(
            [
                nn.Sequential(nn.Conv2d(3, 32, 7, 1, 3), leaky_relu),
                nn.Sequential(
                    nn.Conv2d(32, 32, 3, 2, 1),
                    leaky_relu,
                    nn.Conv2d(32, 32, 3, 1, 1),
                    leaky_relu,
                    nn.Conv2d(32, 32, 3, 1, 1),
                    leaky_relu,
                ),
                nn.Sequential(
                    nn.Conv2d(32, 64, 3, 2, 1),
                    leaky_relu,
                    nn.Conv2d(64, 64, 3, 1, 1),
                    leaky_relu,
                ),
                nn.Sequential(
                    nn.Conv2d(64, 96, 3, 2, 1),
                    leaky_relu,
                    nn.Conv2d(96, 96, 3, 1, 1),
                    leaky_relu,
                ),
                nn.Sequential(nn.Conv2d(96, 128, 3, 2, 1), leaky_relu),
                nn.Sequential(nn.Conv2d(128, 192, 3, 2, 1), leaky_relu),
            ]
        )

    def forward(self, images: torch.Tensor) -> List[torch.Tensor]:
        features = []

        x = images.view(-1, *images.shape[2:])
        for i, conv in enumerate(self.convs):
            x = conv(x)
            if i > 1:
                features.append(x.view(*images.shape[:2], *x.shape[1:]))

        return features[::-1]


class Matching(nn.Module):
    def __init__(self, level: int, num_levels: int = 4, div_flow: float = 20.0) -> None:
        super(Matching, self).__init__()

        corr_stride = [1, 1, 1, 2][level]
        flow_kernel_size = [3, 3, 5, 5][level]
        self.mult = [div_flow / 2 ** (num_levels - i + 1) for i in range(num_levels)][
            level
        ]

        self.leaky_relu = nn.LeakyReLU(0.1, inplace=True)

        if level == 0:
            self.up_flow = None
        else:
            self.up_flow = nn.ConvTranspose2d(2, 2, 4, 2, 1, bias=False, groups=2)

        if level < 3:
            self.up_corr = None
        else:
            self.up_corr = nn.ConvTranspose2d(49, 49, 4, 2, 1, bias=False, groups=49)

        self.flow_net = nn.Sequential(
            nn.Conv2d(49, 128, 3, 1, 1),
            self.leaky_relu,
            nn.Conv2d(128, 128, 3, 1, 1),
            self.leaky_relu,
            nn.Conv2d(128, 96, 3, 1, 1),
            self.leaky_relu,
            nn.Conv2d(96, 64, 3, 1, 1),
            self.leaky_relu,
            nn.Conv2d(64, 32, 3, 1, 1),
            self.leaky_relu,
            nn.Conv2d(32, 2, flow_kernel_size, 1, flow_kernel_size // 2),
        )

        self.corr = SpatialCorrelationSampler(
            kernel_size=1,
            patch_size=7,
            padding=0,
            stride=corr_stride,
            dilation_patch=corr_stride,
        )

        self.warp = WarpingLayer()

    def forward(
        self, feats: torch.Tensor, flow: Optional[torch.Tensor]
    ) -> torch.Tensor:
        warped_feat2 = feats[:, 1]
        if flow is not None:
            flow = self.up_flow(flow)
            warped_feat2 = self.warp(
                feats[:, 1], flow, feats.shape[-2], feats.shape[-1], 1.0 / self.mult
            )

        corr = self.leaky_relu(self.corr(feats[:, 0], warped_feat2))
        corr = corr.view(corr.shape[0], -1, corr.shape[3], corr.shape[4])
        corr = corr / feats.shape[2]
        if self.up_corr is not None:
            corr = self.up_corr(corr)
        new_flow = self.flow_net(corr)
        if flow is not None:
            new_flow = flow + new_flow
        return new_flow


class SubPixel(nn.Module):
    def __init__(self, level: int, num_levels: int = 4, div_flow: float = 20.0) -> None:
        super(SubPixel, self).__init__()

        inputs_dims = [386, 258, 194, 130][level]
        flow_kernel_size = [3, 3, 5, 5][level]
        self.mult = [div_flow / 2 ** (num_levels - i + 1) for i in range(num_levels)][
            level
        ]

        self.leaky_relu = nn.LeakyReLU(0.1, inplace=True)

        self.feat_net = nn.Sequential(
            nn.Conv2d(inputs_dims, 128, 3, 1, 1),
            self.leaky_relu,
            nn.Conv2d(128, 128, 3, 1, 1),
            self.leaky_relu,
            nn.Conv2d(128, 96, 3, 1, 1),
            self.leaky_relu,
            nn.Conv2d(96, 64, 3, 1, 1),
            self.leaky_relu,
            nn.Conv2d(64, 32, 3, 1, 1),
            self.leaky_relu,
        )

        self.flow_net = nn.Conv2d(32, 2, flow_kernel_size, 1, flow_kernel_size // 2)

        self.warp = WarpingLayer()

    def forward(self, feats: torch.Tensor, flow: torch.Tensor) -> torch.Tensor:
        feat_warped = self.warp(
            feats[:, 1], flow, feats.shape[-2], feats.shape[-1], 1.0 / self.mult
        )
        x = torch.cat([feats[:, 0], feat_warped, flow], dim=1)
        x = self.feat_net(x)
        new_flow = self.flow_net(x)
        new_flow = flow + new_flow
        return new_flow, x


class Regularization(nn.Module):
    def __init__(self, level: int, num_levels: int = 4, div_flow: float = 20.0) -> None:
        super(Regularization, self).__init__()

        inputs_dims = [195, 131, 99, 67][level]
        flow_kernel_size = [3, 3, 5, 5][level]
        self.mult = [div_flow / 2 ** (num_levels - i + 1) for i in range(num_levels)][
            level
        ]

        self.leaky_relu = nn.LeakyReLU(0.1, inplace=True)

        if level < 2:
            self.feat_conv = nn.Sequential()
        else:
            self.feat_conv = nn.Sequential(
                nn.Conv2d(inputs_dims - 3, 128, 1, 1, 0), self.leaky_relu
            )
            inputs_dims = 131

        self.feat_net = nn.Sequential(
            nn.Conv2d(inputs_dims, 128, 3, 1, 1),
            self.leaky_relu,
            nn.Conv2d(128, 128, 3, 1, 1),
            self.leaky_relu,
            nn.Conv2d(128, 64, 3, 1, 1),
            self.leaky_relu,
            nn.Conv2d(64, 64, 3, 1, 1),
            self.leaky_relu,
            nn.Conv2d(64, 32, 3, 1, 1),
            self.leaky_relu,
            nn.Conv2d(32, 32, 3, 1, 1),
            self.leaky_relu,
        )

        if level < 2:
            self.dist = nn.Conv2d(32, flow_kernel_size**2, 3, 1, 1)
        else:
            self.dist = nn.Sequential(
                nn.Conv2d(
                    32,
                    flow_kernel_size**2,
                    (flow_kernel_size, 1),
                    1,
                    (flow_kernel_size // 2, 0),
                ),
                nn.Conv2d(
                    flow_kernel_size**2,
                    flow_kernel_size**2,
                    (1, flow_kernel_size),
                    1,
                    (0, flow_kernel_size // 2),
                ),
            )

        self.unfold = nn.Unfold(flow_kernel_size, padding=flow_kernel_size // 2)

        self.warp = WarpingLayer()

    def forward(
        self, images: torch.Tensor, feats: torch.Tensor, flow: torch.Tensor
    ) -> torch.Tensor:
        img2_warped = self.warp(
            images[:, 1], flow, images.shape[-2], images.shape[-1], 1.0 / self.mult
        )
        img_diff_norm = torch.norm(
            images[:, 0] - img2_warped[:, 1], p=2, dim=1, keepdim=True
        )

        flow_mean = flow.view(*flow.shape[:2], -1).mean(dim=-1)[..., None, None]
        flow_nomean = flow - flow_mean
        feat = self.feat_conv(feats[:, 0])
        x = torch.cat([img_diff_norm, flow_nomean, feat], dim=1)
        x = self.feat_net(x)
        dist = self.dist(x)
        dist = dist.square().neg()
        dist = (dist - dist.max(dim=1, keepdim=True)[0]).exp()
        div = dist.sum(dim=1, keepdim=True)

        reshaped_flow_x = self.unfold(flow[:, :1])
        reshaped_flow_x = reshaped_flow_x.view(
            *reshaped_flow_x.shape[:2], *flow.shape[2:4]
        )
        flow_smooth_x = (reshaped_flow_x * dist).sum(dim=1, keepdim=True) / div

        reshaped_flow_y = self.unfold(flow[:, 1:2])
        reshaped_flow_y = reshaped_flow_y.view(
            *reshaped_flow_y.shape[:2], *flow.shape[2:4]
        )
        flow_smooth_y = (reshaped_flow_y * dist).sum(dim=1, keepdim=True) / div

        flow = torch.cat([flow_smooth_x, flow_smooth_y], dim=1)

        return flow, x


class PseudoSubpixel(nn.Module):
    def __init__(self) -> None:
        super(PseudoSubpixel, self).__init__()

        self.up_flow = nn.ConvTranspose2d(2, 2, 4, 2, 1, bias=False, groups=2)

        self.flow_net = nn.Sequential(
            nn.ConvTranspose2d(32, 32, 4, 2, 1), nn.Conv2d(32, 2, 7, 1, 3)
        )

    def forward(self, sub_feat: torch.Tensor, flow: torch.Tensor) -> torch.Tensor:
        return self.up_flow(flow) + self.flow_net(sub_feat)


class PseudoRegularization(nn.Module):
    def __init__(self) -> None:
        super(PseudoRegularization, self).__init__()

        self.feat_net = nn.Sequential(
            nn.ConvTranspose2d(32, 32, 4, 2, 1),
            nn.Conv2d(32, 49, (7, 1), 1, (3, 0)),
            nn.Conv2d(49, 49, (1, 7), 1, (0, 3)),
        )

        self.unfold = nn.Unfold(7, padding=3)

    def forward(self, reg_feat: torch.Tensor, flow: torch.Tensor) -> torch.Tensor:
        dist = self.feat_net(reg_feat)
        dist = dist.square().neg()
        dist = (dist - dist.max(dim=1, keepdim=True)[0]).exp()
        div = dist.sum(dim=1, keepdim=True)

        reshaped_flow_x = self.unfold(flow[:, :1])
        reshaped_flow_x = reshaped_flow_x.view(
            *reshaped_flow_x.shape[:2], *flow.shape[2:4]
        )
        flow_smooth_x = (reshaped_flow_x * dist).sum(dim=1, keepdim=True) / div

        reshaped_flow_y = self.unfold(flow[:, 1:2])
        reshaped_flow_y = reshaped_flow_y.view(
            *reshaped_flow_y.shape[:2], *flow.shape[2:4]
        )
        flow_smooth_y = (reshaped_flow_y * dist).sum(dim=1, keepdim=True) / div

        flow = torch.cat([flow_smooth_x, flow_smooth_y], dim=1)

        return flow


class LiteFlowNet2(BaseModel):
    pretrained_checkpoints = {
        "sintel": "https://github.com/hmorimitsu/ptlflow/releases/download/weights1/liteflownet2-sintel-1e1eb282.ckpt"
    }

    def __init__(
        self,
        div_flow: float = 20.0,
        use_pseudo_regularization: bool = False,
        **kwargs,
    ):
        super(LiteFlowNet2, self).__init__(
            loss_fn=None,
            output_stride=32,
            **kwargs,
        )

        self.div_flow = div_flow
        self.use_pseudo_regularization = use_pseudo_regularization

        self.num_levels = 4

        self.feature_net = FeatureExtractor()
        self.matching_nets = nn.ModuleList(
            [
                Matching(i, self.num_levels, self.div_flow)
                for i in range(self.num_levels)
            ]
        )
        self.subpixel_nets = nn.ModuleList(
            [
                SubPixel(i, self.num_levels, self.div_flow)
                for i in range(self.num_levels)
            ]
        )
        self.regularization_nets = nn.ModuleList(
            [
                Regularization(i, self.num_levels, self.div_flow)
                for i in range(self.num_levels)
            ]
        )

        if self.use_pseudo_regularization:
            self.pseudo_subpixel = PseudoSubpixel()
            self.pseudo_regularization = PseudoRegularization()
            self.up_flow = nn.ConvTranspose2d(2, 2, 4, 2, 1, bias=False, groups=2)
        else:
            self.up_flow = nn.ConvTranspose2d(2, 2, 8, 4, 2, bias=False, groups=2)

    def forward(self, inputs: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
        # The original implementation uses different BGR means for im1 and im2
        # Here, we use the same for both to make it simpler
        images, image_resizer = self.preprocess_images(
            inputs["images"],
            bgr_add=[-0.454253, -0.434631, -0.411618],
            bgr_mult=1.0,
            bgr_to_rgb=True,
            resize_mode="interpolation",
            interpolation_mode="bilinear",
            interpolation_align_corners=False,
        )

        feats_pyr = self.feature_net(images)
        images_pyr = self._create_images_pyr(images, feats_pyr)

        flow_preds = []
        flow = None

        for i in range(self.num_levels):
            flow = self.matching_nets[i](feats_pyr[i], flow)
            flow, sub_feat = self.subpixel_nets[i](feats_pyr[i], flow)
            flow, reg_feat = self.regularization_nets[i](
                images_pyr[i], feats_pyr[i], flow
            )
            flow_preds.append(flow)

        if self.use_pseudo_regularization:
            flow = self.pseudo_subpixel(sub_feat, flow)
            flow = self.pseudo_regularization(reg_feat, flow)
            flow = self.up_flow(flow)
        else:
            flow = self.up_flow(flow)
        flow = flow * self.div_flow
        flow = self.postprocess_predictions(flow, image_resizer, is_flow=True)

        outputs = {}
        if self.training:
            outputs["flow_preds"] = flow_preds
            outputs["flows"] = flow[:, None]
        else:
            outputs["flows"] = flow[:, None]
        return outputs

    def _create_images_pyr(
        self, images: torch.Tensor, feats_pyr: List[torch.Tensor]
    ) -> List[torch.Tensor]:
        batch_size = images.shape[0]
        images = images.view(-1, *images.shape[2:]).detach()
        images_pyr = [
            F.interpolate(
                images,
                size=feats_pyr[i].shape[-2:],
                mode="bilinear",
                align_corners=False,
            )
            for i in range(len(feats_pyr))
        ]
        images_pyr = [im.view(batch_size, -1, *im.shape[1:]) for im in images_pyr]
        return images_pyr


class LiteFlowNet2PseudoReg(LiteFlowNet2):
    pretrained_checkpoints = {
        "kitti": "https://github.com/hmorimitsu/ptlflow/releases/download/weights1/liteflownet2-kitti-da069fca.ckpt"
    }

    def __init__(
        self,
        div_flow: float = 20,
        use_pseudo_regularization: bool = True,
        **kwargs,
    ):
        super().__init__(
            div_flow,
            use_pseudo_regularization,
            **kwargs,
        )


@register_model
class liteflownet2(LiteFlowNet2):
    pass


@register_model
class liteflownet2_pseudoreg(LiteFlowNet2PseudoReg):
    pass
