
from __future__ import absolute_import, division, print_function

import numpy as np
import time

import torch
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader
from tensorboardX import SummaryWriter

import json

from utils import *
from kitti_utils import *
from layers import *

import datasets
import networks
from IPython import embed
#import open3d as o3d


class Trainer:
    def __init__(self, options):
        self.opt = options
        self.log_path = os.path.join(self.opt.log_dir, self.opt.model_name)

        # checking height and width are multiples of 32
        assert self.opt.height % 32 == 0, "'height' must be a multiple of 32"
        assert self.opt.width % 32 == 0, "'width' must be a multiple of 32"

        self.models = {}
        self.parameters_to_train = []

        self.device = torch.device("cpu" if self.opt.no_cuda else "cuda")

        self.num_scales = len(self.opt.scales)
        self.num_input_frames = len(self.opt.frame_ids)
        self.num_pose_frames = 2 if self.opt.pose_model_input == "pairs" else self.num_input_frames

        assert self.opt.frame_ids[0] == 0, "frame_ids must start with 0"

        self.use_pose_net = not (self.opt.use_stereo and self.opt.frame_ids == [0])

        if self.opt.use_stereo:
            self.opt.frame_ids.append("s")

        self.models["encoder"] = networks.ResnetEncoder(
        nInputChannels = 3, pretrained = self.opt.weights_init)
        self.models["encoder"].to(self.device)
        self.parameters_to_train += list(self.models["encoder"].parameters())

        self.models["depth"] = networks.DepthDecoder(
            self.models["encoder"].num_ch_enc, self.opt.scales)
        self.models["depth"].to(self.device)
        self.parameters_to_train += list(self.models["depth"].parameters())



        self.models["flow"] = networks.FlowDecoder(self.models["encoder"].num_ch_enc)
        self.models["flow"].to(self.device)
        self.parameters_to_train += list(self.models["flow"].parameters())


        if self.use_pose_net:
            #__init__(self,  pretrained, num_input_images=1):
            self.models["pose_encoder"] = networks.PoseEncoder(
                pretrained = self.opt.weights_init,
                num_input_images = self.num_pose_frames)
            self.models["pose_encoder"].to(self.device)
            self.parameters_to_train += list(self.models["pose_encoder"].parameters())

            #__init__(self, num_ch_enc, num_input_features, num_frames_to_predict_for=None, stride=1):
            self.models["pose"] = networks.PoseDecoder(
                self.models["pose_encoder"].num_ch_enc,
                num_input_features = 1,
                num_frames_to_predict_for = 2)
            self.models["pose"].to(self.device)
            self.parameters_to_train += list(self.models["pose"].parameters())



        if self.opt.predictive_mask:
            assert self.opt.disable_automasking, \
                "When using predictive_mask, please disable automasking with --disable_automasking"

            # Our implementation of the predictive masking baseline has the the same architecture
            # as our depth decoder. We predict a separate mask for each source frame.
            self.models["predictive_mask"] = networks.DepthDecoder(
                self.models["encoder"].num_ch_enc, self.opt.scales,
                num_output_channels=(len(self.opt.frame_ids) - 1))
            self.models["predictive_mask"].to(self.device)
            self.parameters_to_train += list(self.models["predictive_mask"].parameters())


        self.model_optimizer = optim.Adam(self.parameters_to_train, self.opt.learning_rate)
        self.model_lr_scheduler = optim.lr_scheduler.StepLR(
            self.model_optimizer, self.opt.scheduler_step_size, 0.1)

        if self.opt.load_weights_folder is not None:
            self.load_model()

        print("Training model named:\n  ", self.opt.model_name)
        print("Models and tensorboard events files are saved to:\n  ", self.opt.log_dir)
        print("Training is using:\n  ", self.device)

        # data
        datasets_dict = {"kitti": datasets.KITTIRAWDataset,
                         "kitti_odom": datasets.KITTIOdomDataset}
        self.dataset = datasets_dict[self.opt.dataset]

        fpath = os.path.join(os.path.dirname(__file__), "splits", self.opt.split, "{}_files.txt")

        train_filenames = readlines(fpath.format("train"))
        val_filenames = readlines(fpath.format("val"))
        img_ext = '.png' if self.opt.png else '.jpg'

        num_train_samples = len(train_filenames)
        self.num_total_steps = num_train_samples // self.opt.batch_size * self.opt.num_epochs

        train_dataset = self.dataset(
            self.opt.data_path, train_filenames, self.opt.height, self.opt.width,
            self.opt.frame_ids, 4, is_train=True, img_ext=img_ext)
        self.train_loader = DataLoader(
            train_dataset, self.opt.batch_size, True,
            num_workers=self.opt.num_workers, pin_memory=True, drop_last=True)
        val_dataset = self.dataset(
            self.opt.data_path, val_filenames, self.opt.height, self.opt.width,
            self.opt.frame_ids, 4, is_train=False, img_ext=img_ext)
        self.val_loader = DataLoader(
            val_dataset, self.opt.batch_size, True,
            num_workers=self.opt.num_workers, pin_memory=True, drop_last=True)
        self.val_iter = iter(self.val_loader)

        self.writers = {}
        for mode in ["train", "val"]:
            self.writers[mode] = SummaryWriter(os.path.join(self.log_path, mode))

        if not self.opt.no_ssim:
            self.ssim = SSIM()
            self.ssim.to(self.device)

        self.backproject_depth = {}
        self.project_3d = {}
        for scale in self.opt.scales:
            h = self.opt.height // (2 ** scale)
            w = self.opt.width // (2 ** scale)

            self.backproject_depth[scale] = BackprojectDepth(self.opt.batch_size, h, w)
            self.backproject_depth[scale].to(self.device)

            self.project_3d[scale] = Project3D(self.opt.batch_size, h, w)
            self.project_3d[scale].to(self.device)

        self.depth_metric_names = [
            "de/abs_rel", "de/sq_rel", "de/rms", "de/log_rms", "da/a1", "da/a2", "da/a3"]

        print("Using split:\n  ", self.opt.split)
        print("There are {:d} training items and {:d} validation items\n".format(
            len(train_dataset), len(val_dataset)))

        self.save_opts()

    def set_train(self):
        """Convert all models to training mode
        """
        for m in self.models.values():
            m.train()

    def set_eval(self):
        """Convert all models to testing/evaluation mode
        """
        for m in self.models.values():
            m.eval()

    def train(self):
        """Run the entire training pipeline
        """
        self.epoch = 0
        self.step = 0
        self.start_time = time.time()
        for self.epoch in range(self.opt.num_epochs):
            self.run_epoch()
            if (self.epoch + 1) % self.opt.save_frequency == 0:
                self.save_model()

    def run_epoch(self):
        """Run a single epoch of training and validation"""
        self.model_lr_scheduler.step()  
        self.set_train()

        for batch_idx, inputs in enumerate(self.train_loader):
            outputs, losses = self.process_batch(inputs)

            losses["loss"].backward()
            self.model_optimizer.step()


            duration = time.time() - before_op_time

            # log less frequently after the first 2000 steps to save time & disk space
            early_phase = batch_idx % self.opt.log_frequency == 0 and self.step < 2000
            late_phase = self.step % 2000 == 0

            if early_phase or late_phase:
                self.log_time(batch_idx, duration, losses["loss"].cpu().data)

                if "depth_gt" in inputs:
                    self.compute_depth_losses(inputs, outputs, losses)

                self.log("train", inputs, outputs, losses)
                self.val()

            self.step += 1

    
    def process_batch(self, inputs):
        for key, ipt in inputs.items():
            inputs[key] = ipt.to(self.device)

        features = self.models["encoder"](inputs["color_aug", 0, 0])
        outputs = self.models["depth"](features)

        flow_fw = self.models["flow"](features)
        flow_bw = self.models["flow"](features)

        if self.opt.predictive_mask:
            outputs["predictive_mask"] = self.models["predictive_mask"](features)

        for scale in self.opt.scales:
            disp = outputs[("disp", scale)]
            if not self.opt.v1_multiscale:
                disp = F.interpolate(disp, [self.opt.height, self.opt.width], mode="bilinear", align_corners=False)
            outputs[("disp", 0, scale)] = disp
            if 1 in self.opt.frame_ids:
                outputs[("disp", 1, scale)] = disp

        self.generate_images_pred(inputs, outputs)
        losses = self.compute_losses(inputs, outputs)

        flow_loss_weight = 1.0

        flow_loss = self.compute_flow_loss(inputs, flow_fw, flow_bw)
        losses["flow_loss"] = flow_loss

        # Compute speed loss
        speed_loss = self.compute_speed_loss(inputs, outputs)
        losses["total_loss"] = losses["loss"] + flow_loss_weight * flow_loss + speed_loss

        return outputs, losses


    def compute_speed_loss(self, inputs, outputs):
        speed_loss = 0

        for i, frame_id in enumerate(self.opt.frame_ids[1:]):
            if frame_id == "s":
                continue

            # Retrieve IMU data
            imu_data_0 = inputs[("imu", 0, 0)]
            imu_data_1 = inputs[("imu", frame_id, 0)]

            # Compute average speed from IMU data (assuming vf is the forward speed)
            avg_speed_imu = (imu_data_0[0] + imu_data_1[0]) / 2

            # Get the pose distance from the PoseDecoder
            pose_transform = outputs[("cam_T_cam", 0, frame_id)]

            # Calculate the translation distance
            translation_distance = torch.norm(pose_transform[:, :3, 3], dim=1)

            # Compute the expected distance covered in 0.1 seconds
            expected_distance = avg_speed_imu * 0.1

            # Ensure expected_distance has the same shape as translation_distance
            expected_distance = expected_distance.view(-1, 1)

            # Compute the speed loss as the L2 loss between expected and actual distances
            speed_loss += F.mse_loss(translation_distance, expected_distance)

        return speed_loss / (len(self.opt.frame_ids) - 1)


    def predict_poses(self, inputs, features):
        """Predict poses between input frames for monocular sequences."""
        outputs = {}
        if self.num_pose_frames == 2:
            # In this setting, we compute the pose to each source frame via a
            # separate forward pass through the pose network.

            # select what features the pose network takes as input
            if self.opt.pose_model_type == "shared":
                pose_feats = {f_i: features[f_i] for f_i in self.opt.frame_ids}
            else:
                pose_feats = {f_i: inputs["color_aug", f_i, 0] for f_i in self.opt.frame_ids}

            for f_i in self.opt.frame_ids[1:]:
                if f_i != "s":
                    # To maintain ordering we always pass frames in temporal order
                    if f_i < 0:
                        pose_inputs = [pose_feats[f_i], pose_feats[0]]
                    else:
                        pose_inputs = [pose_feats[0], pose_feats[f_i]]

                    if self.opt.pose_model_type == "separate_resnet":
                        pose_inputs = [self.models["pose_encoder"](torch.cat(pose_inputs, 1))]
                    elif self.opt.pose_model_type == "posecnn":
                        pose_inputs = torch.cat(pose_inputs, 1)

                    axisangle, translation = self.models["pose"](pose_inputs)
                    outputs[("axisangle", 0, f_i)] = axisangle
                    outputs[("translation", 0, f_i)] = translation

                    # Invert the matrix if the frame id is negative
                    outputs[("cam_T_cam", 0, f_i)] = transformation_from_parameters(
                        axisangle[:, 0], translation[:, 0], invert=(f_i < 0))
                    
                    # Also store the inverse transformation for cam_T_cam (f_i, 0)
                    outputs[("cam_T_cam", f_i, 0)] = transformation_from_parameters(
                        axisangle[:, 0], translation[:, 0], invert=(f_i >= 0))

        else:
            # Here we input all frames to the pose net (and predict all poses) together
            if self.opt.pose_model_type in ["separate_resnet", "posecnn"]:
                pose_inputs = torch.cat(
                    [inputs[("color_aug", i, 0)] for i in self.opt.frame_ids if i != "s"], 1)

            elif self.opt.pose_model_type == "shared":
                pose_inputs = [features[i] for i in self.opt.frame_ids if i != "s"]

            axisangle, translation = self.models["pose"](pose_inputs)

            for i, f_i in enumerate(self.opt.frame_ids[1:]):
                if f_i != "s":
                    outputs[("axisangle", 0, f_i)] = axisangle
                    outputs[("translation", 0, f_i)] = translation
                    outputs[("cam_T_cam", 0, f_i)] = transformation_from_parameters(
                        axisangle[:, i], translation[:, i])
                    
                    # Also store the inverse transformation for cam_T_cam (f_i, 0)
                    outputs[("cam_T_cam", f_i, 0)] = transformation_from_parameters(
                        axisangle[:, i], translation[:, i], invert=True)

        return outputs

    def val(self):
        """Validate the model on a single minibatch
        """
        self.set_eval()
        try:
            inputs = next(self.val_iter)  # Correct method to get the next batch
        except StopIteration:
            self.val_iter = iter(self.val_loader)
            inputs = next(self.val_iter)  # Correct method to get the next batch

        with torch.no_grad():
            outputs, losses = self.process_batch(inputs)

            if "depth_gt" in inputs:
                self.compute_depth_losses(inputs, outputs, losses)

            self.log("val", inputs, outputs, losses)
            del inputs, outputs, losses

        self.set_train()

     def generate_images_pred(self, inputs, outputs):
        for scale in self.opt.scales:
            disp_0 = outputs[("disp", 0, scale)]
            if self.opt.v1_multiscale:
                source_scale = scale
            else:
                disp_0 = F.interpolate(
                    disp_0, [self.opt.height, self.opt.width], mode="bilinear", align_corners=False)
                source_scale = 0

            _, depth_0 = disp_to_depth(disp_0, self.opt.min_depth, self.opt.max_depth)
            outputs[("depth", 0, scale)] = depth_0

            if 1 in self.opt.frame_ids:
                disp_1 = outputs[("disp", 1, scale)]
                if not self.opt.v1_multiscale:
                    disp_1 = F.interpolate(
                        disp_1, [self.opt.height, self.opt.width], mode="bilinear", align_corners=False)

                _, depth_1 = disp_to_depth(disp_1, self.opt.min_depth, self.opt.max_depth)
                outputs[("depth", 1, scale)] = depth_1

                # Generate reprojected images from frame 0 to frame 1
                T_01 = outputs[("cam_T_cam", 0, 1)]
                cam_points_0 = self.backproject_depth[source_scale](depth_0, inputs[("inv_K", source_scale)])
                pix_coords_0_to_1 = self.project_3d[source_scale](cam_points_0, inputs[("K", source_scale)], T_01)
                outputs[("sample_0_to_1", scale)] = pix_coords_0_to_1
          
                pred_0_from_1 = F.grid_sample(
                    inputs[("color", 1, source_scale)],
                            padding_mode="border")
                outputs[("color_0_from_1", scale)] = pred_0_from_1

                # Generate reprojected images from frame 1 to frame 0
                T_10 = outputs[("cam_T_cam", 1, 0)]
                cam_points_1 = self.backproject_depth[source_scale](depth_1, inputs[("inv_K", source_scale)])
                pix_coords_1_to_0 = self.project_3d[source_scale](cam_points_1, inputs[("K", source_scale)], T_10)
                outputs[("sample_1_to_0", scale)] = pix_coords_1_to_0

                pred_1_from_0 = F.grid_sample(
                    inputs[("color", 0, source_scale)],
                    padding_mode="border")
                outputs[("color_1_from_0", scale)] = pred_1_from_0

            # Save the disparity map for each frame
            outputs[("disp", 0, scale)] = disp_0
            if 1 in self.opt.frame_ids:
                outputs[("disp", 1, scale)] = disp_1
                outputs[("color", 1, scale)] = inputs[("color", 1, source_scale)]  # Ensure color data is stored

    def compute_reprojection_loss(self, pred, target):
        """Computes reprojection loss between a batch of predicted and target images
        """
        if pred.size() != target.size():
            pred = F.interpolate(pred, size=target.shape[2:], mode='bilinear', align_corners=False)

        l1_loss = abs_diff.mean(1, True)

        if self.opt.no_ssim:
            reprojection_loss = l1_loss
        else:
            ssim_loss = self.ssim(pred, target).mean(1, True)
            reprojection_loss = 0.85 * ssim_loss + 0.15 * l1_loss

        return reprojection_loss

    def depth_to_point_cloud(self, depth, K):

        batch_size, _, height, width = depth.size()
        i, j = torch.meshgrid(torch.arange(height), torch.arange(width), indexing='ij')
        j = j.flatten().float().to(depth.device)
        ones = torch.ones_like(i)
        pixel_coords = torch.stack([j, i, ones], dim=0).unsqueeze(0).repeat(batch_size, 1, 1)

        K = K[:, :3, :3]

        cam_points = torch.matmul(torch.inverse(K), pixel_coords) * depth.view(batch_size, 1, -1)
        return cam_points.view(batch_size, 3, height, width)


    def icp_matching(self, source_points, target_points):

        source = o3d.geometry.PointCloud()
        target = o3d.geometry.PointCloud()
        source_points_np = source_points.detach().cpu().numpy().reshape(-1, 3)
        target_points_np = target_points.detach().cpu().numpy().reshape(-1, 3)
        source.points = o3d.utility.Vector3dVector(source_points_np)
        target.points = o3d.utility.Vector3dVector(target_points_np)

        threshold = 0.02  
        trans_init = np.eye(4)
        reg_p2p = o3d.pipelines.registration.registration_icp(
            source, target, threshold, trans_init,
            o3d.pipelines.registration.TransformationEstimationPointToPoint()
        )
        return reg_p2p.transformation

    def compute_icp_loss(self, inputs, outputs):

        K = inputs[("K", 0)]
        depth0 = outputs[("depth", 0, 0)]
        depth1 = outputs[("depth", 1, 0)]

        points0 = self.depth_to_point_cloud(depth0, K)
        points1 = self.depth_to_point_cloud(depth1, K)

        points0_flat = points0.permute(0, 2, 3, 1).reshape(-1, 3)
        points1_flat = points1.permute(0, 2, 3, 1).reshape(-1, 3)

        icp_transformation = self.icp_matching(points0_flat, points1_flat)

        pose_network = outputs[("cam_T_cam", 0, 1)]
        icp_pose = torch.from_numpy(icp_transformation).float().to(pose_network.device)
        icp_pose = icp_pose.unsqueeze(0).repeat(pose_network.size(0), 1, 1)
        
        print("pose_network",pose_network)
        print("icp_pose",icp_pose)

        if icp_pose.shape != pose_network.shape:
            return torch.tensor(10.0, device=pose_network.device)

        loss = F.mse_loss(pose_network, icp_pose)
        return loss

    def compute_edge_loss(self, disp, target):

        target_gray = 0.2989 * target[:, 0, :, :] + 0.5870 * target[:, 1, :, :] + 0.1140 * target[:, 2, :, :]
        target_gray = target_gray.unsqueeze(1) 

        # Define Sobel filters
        sobel_x = torch.tensor([[[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]]], dtype=torch.float32).unsqueeze(0).to(self.device)
        sobel_y = torch.tensor([[[-1, -2, -1], [0, 0, 0], [1, 2, 1]]], dtype=torch.float32).unsqueeze(0).to(self.device)

        # Add an extra dimension to the Sobel filters
        sobel_x = sobel_x.expand(1, 1, 3, 3)
        sobel_y = sobel_y.expand(1, 1, 3, 3)

        target_grad_x = F.conv2d(target_gray, sobel_x, padding=1)
        target_grad_y = F.conv2d(target_gray, sobel_y, padding=1)
        disp_grad_x = F.conv2d(disp, sobel_x, padding=1)
        disp_grad_y = F.conv2d(disp, sobel_y, padding=1)

        edge_loss = torch.mean(torch.abs(target_grad_x - disp_grad_x) + torch.abs(target_grad_y - disp_grad_y))
        return edge_loss


    def gaussian_kernel(self,size: int, sigma: float):
        """Create a Gaussian kernel."""
        coords = torch.arange(size).to(torch.float32) - size // 2
        grid = coords.repeat(size, 1)
        kernel = torch.exp(-(grid ** 2 + grid.t() ** 2) / (2 * sigma ** 2))
        kernel /= kernel.sum()
        return kernel

    def gaussian_filter(input, kernel):
        """Apply Gaussian filter to input."""
        channels = input.shape[1]
        kernel = kernel.expand(channels, 1, -1, -1).to(input.device)
        return F.conv2d(input, kernel, padding=kernel.shape[-1]//2, groups=channels)
   
    def get_image_gradients(self,image):

        sobel_x = torch.tensor([[[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]]], dtype=torch.float32).unsqueeze(0).to(image.device)
        sobel_y = torch.tensor([[[-1, -2, -1], [0, 0, 0], [1, 2, 1]]], dtype=torch.float32).unsqueeze(0).to(image.device)
        sobel_x = sobel_x.expand(1, 1, 3, 3)
        sobel_y = sobel_y.expand(1, 1, 3, 3)
        image_gray = 0.2989 * image[:, 0, :, :] + 0.5870 * image[:, 1, :, :] + 0.1140 * image[:, 2, :, :]
        image_gray = image_gray.unsqueeze(1)

        grad_x = F.conv2d(image_gray, sobel_x, padding=1)
        grad_y = F.conv2d(image_gray, sobel_y, padding=1)
        
        return grad_x, grad_y

    def get_edge_aware_smoothness_loss(self, disp, img):

        disp = F.interpolate(disp, size=(img.size(2), img.size(3)), mode='bilinear', align_corners=False)

        grad_disp_x = torch.abs(disp[:, :, :, :-1] - disp[:, :, :, 1:])
        grad_disp_y = torch.abs(disp[:, :, :-1, :] - disp[:, :, 1:, :])

        grad_img_x = torch.mean(torch.abs(img[:, :, :, :-1] - img[:, :, :, 1:]), 1, keepdim=True)
        grad_img_y = torch.mean(torch.abs(img[:, :, :-1, :] - img[:, :, 1:, :]), 1, keepdim=True)

        sobel_x = torch.tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=torch.float32).unsqueeze(0).unsqueeze(0).to(self.device)
        sobel_y = torch.tensor([[-1, -2, -1], [0, 0, 0], [1, 2, 1]], dtype=torch.float32).unsqueeze(0).unsqueeze(0).to(self.device)

        if grad_img_x.size() != grad_img_sobel_x.size():
            grad_img_sobel_x = F.interpolate(grad_img_sobel_x, size=grad_img_x.shape[2:], mode='bilinear', align_corners=False)
        if grad_img_y.size() != grad_img_sobel_y.size():
            grad_img_sobel_y = F.interpolate(grad_img_sobel_y, size=grad_img_y.shape[2:], mode='bilinear', align_corners=False)

        weight_x = torch.exp(-grad_img_combined_x)
        weight_y = torch.exp(-grad_img_combined_y)

        if weight_x.size()[2:] != grad_disp_x.size()[2:]:
            weight_x = F.interpolate(weight_x, size=grad_disp_x.shape[2:], mode='bilinear', align_corners=False)
        if weight_y.size()[2:] != grad_disp_y.size()[2:]:
            weight_y = F.interpolate(weight_y, size=grad_disp_y.shape[2:], mode='bilinear', align_corners=False)

        weighted_grad_disp_x = grad_disp_x * weight_x
        weighted_grad_disp_y = grad_disp_y * weight_y

        smoothness_x = weighted_grad_disp_x.mean()
        smoothness_y = weighted_grad_disp_y.mean()

        return smoothness_x + smoothness_y



    def compute_affine_invariant_loss(self, inputs, outputs, scale):
        """Compute the affine invariant loss"""
        source_scale = 0 if not self.opt.v1_multiscale else scale

        disp = outputs[("disp", scale)]
        color_0 = inputs[("color", 0, scale)]
        target = inputs[("color", 0, source_scale)]

        losses = []

        for frame_id in self.opt.frame_ids[1:]:
            T = outputs[("cam_T_cam", 0, frame_id)]
            
            # Ensure the depth map size matches the camera intrinsics size
            depth = F.interpolate(disp, [self.opt.height // (2 ** source_scale), self.opt.width // (2 ** source_scale)], mode='bilinear', align_corners=False)
            
            cam_points = self.backproject_depth[source_scale](depth, inputs[("inv_K", source_scale)])
            pix_coords = self.project_3d[source_scale](cam_points, inputs[("K", source_scale)], T)

            pred = F.grid_sample(inputs[("color", frame_id, source_scale)], pix_coords, padding_mode="border")
            
            # Apply affine transformation to align the images
            affine_matrix = self.get_affine_matrix(inputs[("color", frame_id, source_scale)], target, pred.size(0))
            aligned_pred = self.apply_affine_transform(pred, affine_matrix)
            aligned_target = self.apply_affine_transform(target, affine_matrix)

            affine_loss = self.compute_reprojection_loss(aligned_pred, aligned_target)
            losses.append(affine_loss)

        return torch.cat(losses, 1).mean()

    def get_affine_matrix(self, img1, img2, batch_size):
        """Compute the affine transformation matrix to align img1 to img2"""
        # Use feature matching or other geometric methods to compute the affine transformation matrix
        # This is a placeholder function, implement the actual method
        affine_matrix = torch.eye(3, device=img1.device).unsqueeze(0).repeat(batch_size, 1, 1)
        return affine_matrix

    def apply_affine_transform(self, img, affine_matrix):
        """Apply the affine transformation to the image"""
        grid = F.affine_grid(affine_matrix[:, :2], img.size(), align_corners=False)
        transformed_img = F.grid_sample(img, grid, align_corners=False)
        return transformed_img
    
    def warp_image(self, img, flow):
        B, C, H, W = img.size()
        grid_x, grid_y = torch.meshgrid(torch.arange(H), torch.arange(W), indexing='ij')  
        grid = torch.stack((grid_y, grid_x), 2).float()  
        grid = grid.unsqueeze(0).repeat(B, 1, 1, 1).to(img.device)

        flow = F.interpolate(flow, size=(H, W), mode='bilinear', align_corners=False)
        flow = flow.permute(0, 2, 3, 1)  

        grid[:, :, :, 0] = 2.0 * grid[:, :, :, 0] / (W - 1) - 1.0
        grid[:, :, :, 1] = 2.0 * grid[:, :, :, 1] / (H - 1) - 1.0

        warped_img = F.grid_sample(img, grid, padding_mode="border")
        return warped_img

    def compute_flow_loss(self, inputs, flow_fw, flow_bw):
        img_0 = inputs[("color", 0, 0)]
        img_1 = inputs[("color", 1, 0)]

        img_0_warped = self.warp_image(img_1, flow_fw)
        img_1_warped = self.warp_image(img_0, flow_bw)

        loss_fw = self.compute_reprojection_loss(img_0_warped, img_0)
        loss_bw = self.compute_reprojection_loss(img_1_warped, img_1)

        flow_consistency_loss = torch.mean(torch.abs(flow_fw + self.warp_image(flow_bw, flow_fw)))

        return loss_fw.mean() + loss_bw.mean() + 0.1 * flow_consistency_loss

   
    def compute_losses(self, inputs, outputs):
        """Compute the reprojection and smoothness losses for a minibatch"""
        losses = {}
        total_loss = 0

        for scale in self.opt.scales:
            loss = 0

            if self.opt.v1_multiscale:
                source_scale = scale
            else:
                source_scale = 0

            disp = outputs[("disp", scale)]
            color_0 = inputs[("color", 0, scale)]
            color_1 = inputs[("color", 1, scale)]

            pred_0_from_1 = outputs[("color_0_from_1", scale)]
            reprojection_losses.append(self.compute_reprojection_loss(pred_0_from_1, color_0))

            pred_1_from_0 = outputs[("color_1_from_0", scale)]
            back_reprojection_losses.append(self.compute_reprojection_loss(pred_1_from_0, color_1))

            reprojection_losses = torch.cat(reprojection_losses, 1)
            back_reprojection_losses = torch.cat(back_reprojection_losses, 1)

            if not self.opt.disable_automasking:
                identity_reprojection_losses = []
                for frame_id in self.opt.frame_ids[1:]:
                    if frame_id != "s":

                        pred = inputs[("color", frame_id, source_scale)]
                        identity_reprojection_losses.append(self.compute_reprojection_loss(pred, inputs[("color", 0, source_scale)]))

                identity_reprojection_losses = torch.cat(identity_reprojection_losses, 1)

            # Adjust sizes if needed
            #if identity_reprojection_loss.size() != reprojection_loss.size():
            if identity_reprojection_loss.dim() > 2 and identity_reprojection_loss.size() != reprojection_loss.size():
                identity_reprojection_loss = F.interpolate(identity_reprojection_loss, size=reprojection_loss.shape[2:], mode='bilinear', align_corners=False)
            if back_reprojection_loss.size() != reprojection_loss.size():
                back_reprojection_loss = F.interpolate(back_reprojection_loss, size=reprojection_loss.shape[2:], mode='bilinear', align_corners=False)

            if not self.opt.disable_automasking:
                identity_reprojection_loss += torch.randn(identity_reprojection_loss.shape, device=self.device) * 0.00001
                combined = torch.cat((identity_reprojection_loss, reprojection_loss, back_reprojection_loss), dim=1)
            else:
                combined = torch.cat((reprojection_loss, back_reprojection_loss), dim=1)

            if combined.shape[1] == 1:
                to_optimise = combined
            else:
                to_optimise, idxs = torch.min(combined, dim=1)

            if not self.opt.disable_automasking:
                outputs["identity_selection/{}".format(scale)] = (idxs > identity_reprojection_loss.shape[1] - 1).float()

            loss += to_optimise.mean()

            mean_disp = disp.mean(2, True).mean(3, True)
            norm_disp = disp / (mean_disp + 1e-7)

            total_loss += loss
            losses["loss/{}".format(scale)] = loss

        total_loss /= self.num_scales
        losses["loss"] = total_loss
        return losses
