import sys
import os
import threading

import numpy
import torch
from torch.autograd import Variable
from torch.autograd import gradcheck

from .lr_scheduler import *


def charbonier_loss(x, epsilon):
    loss = torch.mean(torch.sqrt(x * x + epsilon * epsilon))
    return loss


def negPSNR_loss(x, epsilon):
    loss = torch.mean(torch.mean(torch.mean(torch.sqrt(x * x + epsilon * epsilon), dim=1), dim=1), dim=1)
    return torch.mean(-torch.log(1.0 / loss) / 100.0)


def tv_loss(x, epsilon):
    loss = torch.mean(torch.sqrt(
        (x[:, :, :-1, :-1] - x[:, :, 1:, :-1]) ** 2 +
        (x[:, :, :-1, :-1] - x[:, :, :-1, 1:]) ** 2 + epsilon * epsilon
    )
    )
    return loss


def gra_adap_tv_loss(flow, image, epsilon):
    w = torch.exp(- torch.sum(torch.abs(image[:, :, :-1, :-1] - image[:, :, 1:, :-1]) +
                              torch.abs(image[:, :, :-1, :-1] - image[:, :, :-1, 1:]), dim=1))
    tv = torch.sum(torch.sqrt((flow[:, :, :-1, :-1] - flow[:, :, 1:, :-1]) ** 2 + (
                flow[:, :, :-1, :-1] - flow[:, :, :-1, 1:]) ** 2 + epsilon * epsilon), dim=1)
    loss = torch.mean(w * tv)
    return loss


def smooth_loss(x, epsilon):
    loss = torch.mean(
        torch.sqrt(
            (x[:, :, :-1, :-1] - x[:, :, 1:, :-1]) ** 2 +
            (x[:, :, :-1, :-1] - x[:, :, :-1, 1:]) ** 2 + epsilon ** 2
        )
    )
    return loss


def motion_sym_loss(offset, epsilon, occlusion=None):
    if occlusion == None:
        # return torch.mean(torch.sqrt( (offset[:,:2,...] + offset[:,2:,...])**2 + epsilon **2))
        return torch.mean(torch.sqrt((offset[0] + offset[1]) ** 2 + epsilon ** 2))
    else:
        # TODO: how to design the occlusion aware offset symmetric loss?
        # return torch.mean(torch.sqrt((offset[:,:2,...] + offset[:,2:,...])**2 + epsilon **2))
        return torch.mean(torch.sqrt((offset[0] + offset[1]) ** 2 + epsilon ** 2))


def part_loss(diffs, offsets, occlusions, images, epsilon, use_negPSNR=False):
    if use_negPSNR:
        pixel_loss = [negPSNR_loss(diff, epsilon) for diff in diffs]
    else:
        pixel_loss = [charbonier_loss(diff, epsilon) for diff in diffs]
    # offset_loss = [tv_loss(offset[0], epsilon) + tv_loss(offset[1], epsilon) for offset in
    #               offsets]

    if offsets[0][0] is not None:
        offset_loss = [gra_adap_tv_loss(offset[0], images[0], epsilon) + gra_adap_tv_loss(offset[1], images[1], epsilon)
                       for offset in
                       offsets]
    else:
        offset_loss = [Variable(torch.zeros(1).cuda())]
    # print(torch.max(occlusions[0]))
    # print(torch.min(occlusions[0]))
    # print(torch.mean(occlusions[0]))

    # occlusion_loss = [smooth_loss(occlusion, epsilon) + charbonier_loss(occlusion - 0.5, epsilon) for occlusion in occlusions]
    # occlusion_loss = [smooth_loss(occlusion, epsilon) + charbonier_loss(occlusion[:, 0, ...] - occlusion[:, 1, ...], epsilon) for occlusion in occlusions]

    sym_loss = [motion_sym_loss(offset, epsilon=epsilon) for offset in offsets]
    # sym_loss = [ motion_sym_loss(offset,occlusion) for offset,occlusion in zip(offsets,occlusions)]
    return pixel_loss, offset_loss, sym_loss
