import sys
import os

import sys
import  threading
import torch
from torch.autograd import Variable
from lr_scheduler import *
from torch.autograd import gradcheck

import numpy




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