from einops.einops import rearrange
import torch
import torch.nn as nn
import torch.nn.functional as F
from romatch.utils.utils import get_gt_warp
import wandb
import romatch
import math

# This is slightly different than regular romatch due to significantly worse corresps
# The confidence loss is quite tricky here //Johan


class RobustLosses(nn.Module):
    def __init__(
        self,
        robust=False,
        center_coords=False,
        scale_normalize=False,
        ce_weight=0.01,
        local_loss=True,
        local_dist=None,
        smooth_mask=False,
        depth_interpolation_mode="bilinear",
        mask_depth_loss=False,
        relative_depth_error_threshold=0.05,
        alpha=1.0,
        c=1e-3,
        epe_mask_prob_th=None,
        cert_only_on_consistent_depth=False,
    ):
        super().__init__()
        if local_dist is None:
            local_dist = {}
        self.robust = robust  # measured in pixels
        self.center_coords = center_coords
        self.scale_normalize = scale_normalize
        self.ce_weight = ce_weight
        self.local_loss = local_loss
        self.local_dist = local_dist
        self.smooth_mask = smooth_mask
        self.depth_interpolation_mode = depth_interpolation_mode
        self.mask_depth_loss = mask_depth_loss
        self.relative_depth_error_threshold = relative_depth_error_threshold
        self.avg_overlap = dict()
        self.alpha = alpha
        self.c = c
        self.epe_mask_prob_th = epe_mask_prob_th
        self.cert_only_on_consistent_depth = cert_only_on_consistent_depth

    def corr_volume_loss(self, mnn: torch.Tensor, corr_volume: torch.Tensor, scale):
        b, h, w, h, w = corr_volume.shape
        inv_temp = 10
        corr_volume = corr_volume.reshape(-1, h * w, h * w)
        nll = -(inv_temp * corr_volume).log_softmax(dim=1) - (inv_temp * corr_volume).log_softmax(dim=2)
        corr_volume_loss = nll[mnn[:, 0], mnn[:, 1], mnn[:, 2]].mean()

        losses = {
            f"gm_corr_volume_loss_{scale}": corr_volume_loss.mean(),
        }
        wandb.log(losses, step=romatch.GLOBAL_STEP)
        return losses

    def regression_loss(self, x2, prob, flow, certainty, scale, eps=1e-8, mode="delta"):
        epe = (flow.permute(0, 2, 3, 1) - x2).norm(dim=-1)
        if scale in self.local_dist:
            prob = prob * (epe < (2 / 512) * (self.local_dist[scale] * scale)).float()
        if scale == 1:
            pck_05 = (epe[prob > 0.99] < 0.5 * (2 / 512)).float().mean()
            wandb.log({"train_pck_05": pck_05}, step=romatch.GLOBAL_STEP)
        if self.epe_mask_prob_th is not None:
            # if too far away from gt, certainty should be 0
            gt_cert = prob * (epe < scale * self.epe_mask_prob_th)
        else:
            gt_cert = prob
        if self.cert_only_on_consistent_depth:
            ce_loss = F.binary_cross_entropy_with_logits(certainty[:, 0][prob > 0], gt_cert[prob > 0])
        else:
            ce_loss = F.binary_cross_entropy_with_logits(certainty[:, 0], gt_cert)
        a = self.alpha[scale] if isinstance(self.alpha, dict) else self.alpha
        cs = self.c * scale
        x = epe[prob > 0.99]
        reg_loss = cs**a * ((x / (cs)) ** 2 + 1**2) ** (a / 2)
        if not torch.any(reg_loss):
            reg_loss = ce_loss * 0.0  # Prevent issues where prob is 0 everywhere
        losses = {
            f"{mode}_certainty_loss_{scale}": ce_loss.mean(),
            f"{mode}_regression_loss_{scale}": reg_loss.mean(),
        }
        wandb.log(losses, step=romatch.GLOBAL_STEP)
        return losses

    def forward(self, corresps, batch):
        scales = list(corresps.keys())
        tot_loss = 0.0
        # scale_weights due to differences in scale for regression gradients and classification gradients
        for scale in scales:
            scale_corresps = corresps[scale]
            (
                scale_certainty,
                flow_pre_delta,
                delta_cls,
                offset_scale,
                scale_gm_corr_volume,
                scale_gm_certainty,
                flow,
                scale_gm_flow,
            ) = (
                scale_corresps["certainty"],
                scale_corresps.get("flow_pre_delta"),
                scale_corresps.get("delta_cls"),
                scale_corresps.get("offset_scale"),
                scale_corresps.get("corr_volume"),
                scale_corresps.get("gm_certainty"),
                scale_corresps["flow"],
                scale_corresps.get("gm_flow"),
            )
            if flow_pre_delta is not None:
                flow_pre_delta = rearrange(flow_pre_delta, "b d h w -> b h w d")
                b, h, w, d = flow_pre_delta.shape
            else:
                # _ = 1
                b, _, h, w = scale_certainty.shape
            gt_warp, gt_prob = get_gt_warp(
                batch["im_A_depth"],
                batch["im_B_depth"],
                batch["T_1to2"],
                batch["K1"],
                batch["K2"],
                H=h,
                W=w,
            )
            x2 = gt_warp.float()
            prob = gt_prob

            if scale_gm_corr_volume is not None:
                gt_warp_back, _ = get_gt_warp(
                    batch["im_B_depth"],
                    batch["im_A_depth"],
                    batch["T_1to2"].inverse(),
                    batch["K2"],
                    batch["K1"],
                    H=h,
                    W=w,
                )
                grid = torch.stack(
                    torch.meshgrid(
                        torch.linspace(-1 + 1 / w, 1 - 1 / w, w),
                        torch.linspace(-1 + 1 / h, 1 - 1 / h, h),
                        indexing="xy",
                    ),
                    dim=-1,
                ).to(gt_warp.device)
                # fwd_bck = F.grid_sample(gt_warp_back.permute(0,3,1,2), gt_warp, align_corners=False, mode = 'bilinear').permute(0,2,3,1)
                # diff = (fwd_bck - grid).norm(dim = -1)
                with torch.no_grad():
                    D_B = torch.cdist(gt_warp.float().reshape(-1, h * w, 2), grid.reshape(-1, h * w, 2))
                    D_A = torch.cdist(grid.reshape(-1, h * w, 2), gt_warp_back.float().reshape(-1, h * w, 2))
                    inds = torch.nonzero(
                        (D_B == D_B.min(dim=-1, keepdim=True).values)
                        * (D_A == D_A.min(dim=-2, keepdim=True).values)
                        * (D_B < 0.01)
                        * (D_A < 0.01)
                    )

                gm_cls_losses = self.corr_volume_loss(inds, scale_gm_corr_volume, scale)
                gm_loss = gm_cls_losses[f"gm_corr_volume_loss_{scale}"]
                tot_loss = tot_loss + gm_loss
            elif scale_gm_flow is not None:
                gm_flow_losses = self.regression_loss(x2, prob, scale_gm_flow, scale_gm_certainty, scale, mode="gm")
                gm_loss = (
                    self.ce_weight * gm_flow_losses[f"gm_certainty_loss_{scale}"]
                    + gm_flow_losses[f"gm_regression_loss_{scale}"]
                )
                tot_loss = tot_loss + gm_loss
            delta_regression_losses = self.regression_loss(x2, prob, flow, scale_certainty, scale)
            reg_loss = (
                self.ce_weight * delta_regression_losses[f"delta_certainty_loss_{scale}"]
                + delta_regression_losses[f"delta_regression_loss_{scale}"]
            )
            tot_loss = tot_loss + reg_loss
        return tot_loss
