"""
TH flow problem
8 cases...
primal/mixed + fc/conv + variational/residual
"""
import torch
import torch.autograd as ag
import torch.nn.functional as F
import matplotlib.pyplot as plt
plt.switch_backend('agg')


def grad(outputs, inputs):
    return ag.grad(outputs, inputs, grad_outputs=torch.ones_like(outputs), 
                   create_graph=True)


def bilinear_interpolate_torch(im, x, y):
    # https://gist.github.com/peteflorence/a1da2c759ca1ac2b74af9a83f69ce20e
    if torch.cuda.is_available():
        dtype = torch.cuda.FloatTensor
        dtype_long = torch.cuda.LongTensor
    else:
        dtype = torch.FloatTensor
        dtype_long = torch.LongTensor
    x0 = torch.floor(x).type(dtype_long)
    x1 = x0 + 1
    
    y0 = torch.floor(y).type(dtype_long)
    y1 = y0 + 1

    x0 = torch.clamp(x0, 0, im.shape[1]-1)
    x1 = torch.clamp(x1, 0, im.shape[1]-1)
    y0 = torch.clamp(y0, 0, im.shape[0]-1)
    y1 = torch.clamp(y1, 0, im.shape[0]-1)
    
    Ia = im[ y0, x0 ][0]
    Ib = im[ y1, x0 ][0]
    Ic = im[ y0, x1 ][0]
    Id = im[ y1, x1 ][0]
    
    wa = (x1.type(dtype)-x) * (y1.type(dtype)-y)
    wb = (x1.type(dtype)-x) * (y-y0.type(dtype))
    wc = (x-x0.type(dtype)) * (y1.type(dtype)-y)
    wd = (x-x0.type(dtype)) * (y-y0.type(dtype))

    return torch.t((torch.t(Ia)*wa)) + torch.t(torch.t(Ib)*wb) \
        + torch.t(torch.t(Ic)*wc) + torch.t(torch.t(Id)*wd)


def primal_residual_fc(model, x, K_grad_ver, K_grad_hor, K, verbose=False):
    """Computes the residules of satisifying PDE at x
    Permeability is also provided: K

    First assume x is on grid

    Args:
        model (Module): u = f(x) pressure network, input is spatial coordinate
        x (Tensor): (N, 2) spatial input x. could be off-grid, vary very pass
        grad_K (Tensor): estimated gradient field
        verbose (bool): If True, print info
    Returns:
        residual: (N, 1)
    """
 
    assert len(K_grad_ver) == len(x)
    x.requires_grad = True
    u = model(x)
    # grad outputs a tuple: (N, 2)
    u_x = grad(u, x)[0]
    
    div1 = K_grad_ver * u_x[:, 0] + K * grad(u_x[:, 0], x)[0][:, 0]
    div2 = K_grad_hor * u_x[:, 1] + K * grad(u_x[:, 1], x)[0][:, 1]
    div = div1 + div2

    if verbose:
        print(div.detach().mean(), div.detach().max(), div.detach().min())
    return (div ** 2).mean()

def neumann_boundary(model, x):
    # bug: u_y! NOT u_x
    x.requires_grad = True
    u = model(x)
    u_ver = grad(u, x)[0][:, 0]
    return (u_ver ** 2).mean()


def neumann_boundary_mixed(model, x):

    # x.requires_grad = True
    y = model(x)
    tau_ver = y[:, 1]
    
    return (tau_ver ** 2).mean()


def primal_variational_fc(model, x, K, verbose=False):
    """Evaulate energy functional. Simple MC. Evaluate on [1:-1, 1:-1] of grid

    Args:
        x (Tensor): colloc points on interior of grid (63 ** 2, 2)
    """
    x.requires_grad = True
    u = model(x)
    u_x = grad(u, x)[0]
    u_x_squared = (u_x ** 2).sum(1)
    energy = (0.5 * K * u_x_squared).mean()
    if verbose:
        print(f'energy: {energy:.6f}')
    return energy


def mixed_residual_fc(model, x, K, verbose=False, rand_colloc=False, fig_dir=None):
    """
    Args:
        x: (N, 2)
        K: (N, 1)

    """
    x.requires_grad = True
    # (N, 3)
    y = model(x)
    u = y[:, 0]
    # (N, 2)
    tau = y[:, [1, 2]]
    # (N, 2)
    u_x = grad(u, x)[0]

    grad_tau_ver = grad(y[:, 1], x)[0][:, 0]
    grad_tau_hor = grad(y[:, 2], x)[0][:, 1]

    if rand_colloc:
        K = bilinear_interpolate_torch(K.unsqueeze(-1), x[:, [1]], x[:, [0]])
        K = K.t()
        # print(f'K interp: {K.shape}')
        # plt.imshow(K[0].detach().cpu().numpy().reshape(65, 65))
        # plt.savefig(fig_dir+'/Kinterp.png')
        # plt.close()


    loss_constitutive = ((K * u_x + tau) ** 2).mean()
    loss_continuity = ((grad_tau_ver + grad_tau_hor) ** 2).mean()

    return loss_constitutive + loss_continuity


"""
ConvNet ============================================
"""

def energy_functional_exp(input, output, sobel_filter):
    r""" sigma = -exp(K * u) * grad(u)

    V(u, K) = \int 0.5 * exp(K*u) * |grad(u)|^2 dx
    """
    grad_h = sobel_filter.grad_h(output)
    grad_v = sobel_filter.grad_v(output)

    return (0.5 * torch.exp(input * output) * (grad_h ** 2 + grad_v ** 2)).mean()


def th_conv_constitutive_constraint(input, output, sobel_filter):
    """sigma = - K * grad(u)

    Args:
        input (Tensor): (1, 3, 65, 65)
        output (Tensor): (1, 1, 65, 65), 
            three channels from 0-2: u, sigma_1, sigma_2
    """
    thermal_conductivity = 0.01
    Enthalpy =	4100
    input = input.unsqueeze(0)
    #qxIN = torch.tensor.copy(output)
    #qxIN.new_zeros((1, 1, 64, 64))
    #qyIN = torch.tensor.copy(output)
    #qyIN.new_zeros((1, 1, 64, 64))
    #print(input.shape)
    #qxIN = output
    #qyIN = output
    qxIN = input[0,0,:,:] 
    qyIN = input[0,1,:,:] 
    #energyIn = input[0,2,:,:]
    qxIN = qxIN.unsqueeze(0)
    qyIN = qyIN.unsqueeze(0)
    #print(qxIN.shape)
    #print(qyIN.shape)
    
    
    grad_temp = thermal_conductivity * sobel_filter.grad_temp(output)	# Delta T in x direction
    #grad_h = sobel_filter.grad_h(output)	# Delta T in x direction
    #grad_v = sobel_filter.grad_v(output)	# Delta T in y direction
    #est_sigma1 = thermal_conductivity * grad_h	# k delta T
    #est_sigma2 = thermal_conductivity * grad_v	# k delta T
    #totalEnergy = thermal_conductivity*grad_temp + input[2, :,:]
    
    #qxH = input[0,:,:].copy()
    #qyH = input[1,:,:].copy()
    #h_grad = Enthalpy*sobel_filter.grad_h(output)*qxIN
    #v_grad = Enthalpy*sobel_filter.grad_v(output)*qyIN
    '''
    h_grad_input_up = Enthalpy*sobel_filter.grad_h_temp_up(qxIN)*sobel_filter.grad_temp(output)
    h_grad_input_down = Enthalpy*sobel_filter.grad_h_temp_down(qxIN)*sobel_filter.grad_temp(output)
    v_grad_input_up = Enthalpy*sobel_filter.grad_v_temp_up(qyIN)*sobel_filter.grad_temp(output)
    v_grad_input_down = Enthalpy*sobel_filter.grad_v_temp_down(qyIN)*sobel_filter.grad_temp(output)
    '''
    #h_grad_input = Enthalpy*qxIN
    #v_grad_input = Enthalpy*qyIN
    ## Manually determine upwind gradient
    '''for i in range(1,63):
         for j in range(1,63):
             if qxIN[0,0,i,j] > 0:
                 h_grad_input[0,0,i,j] = (output[0,0,i,j-1] - output[0,0,i,j])*qxIN[0,0,i,j]
             else:
                 h_grad_input[0,0,i,j] = (output[0,0,i,j+1] - output[0,0,i,j])*qxIN[0,0,i,j]
             
             if qyIN[0,0,i,j] > 0:
                 h_grad_input[0,0,i,j] = (output[0,0,i-1,j] - output[0,0,i,j])*qxIN[0,0,i,j]
             else:
                 h_grad_input[0,0,i,j] = (output[0,0,i+1,j] - output[0,0,i,j])*qxIN[0,0,i,j]
    '''
    #qxH = h_grad_input[0,:,:] * output
    #qyH = v_grad_input[1,:,:] * output
    #qxH = input[0, :,:] * 2 * (output)		# qx times Enthalpy in x
    #qyH = input[1, :,:] * 2 * (output)	# qy times Enthalpy in y
    totalEnergy = (grad_temp) #+ h_grad_input_up + v_grad_input_up + h_grad_input_down + v_grad_input_down) #+ input[0,2, :,:]
    
    #print("shapes: ", est_sigma1.shape)
    #print("shapes: ", qxH.shape)
    
    #qxHMin = qxH #+ est_sigma1
    #qyHMin = qyH #+ est_sigma2
    #totalEnergy = est_sigma1 + input[2, :,:]
    #totalEnergy = qxHMin + qyHMin + input[2, :,:]
    #deltaQxMin = sobel_filter.grad_v(qxHMin)
    #deltaQyMin = sobel_filter.grad_h(qyHMin)
    #totalEnergy = deltaQxMin[0,0,:,:] + deltaQyMin[0,0,:,:] + input[2, :,:]
    
    #print(deltaQxMin[0,0,32,32])
    #print(deltaQyMin[0,0,32,32])
    #print(input[2,32,32])

    #print("input: ", input[0][0][0])
    # input [which permeability field number - batch size][only dimension 1, permeability field only. Can add other things here][values for one permeability field][]

    #return ((output[:, [1]] - est_sigma1) ** 2 
    #    + (output[:, [2]] - est_sigma2) ** 2).mean()
    return ((totalEnergy**2).mean())

def conv_constitutive_constraint_nonlinear(input, output, sobel_filter, beta1, beta2):
    """Nonlinear extension of Darcy's law
        - K * grad_u = sigma + beta1 * sqrt(K) * sigma ** 2 + beta2 * K * sigma ** 3

    Args:
        input: K
        output: u, sigma1, sigma2
    """
    K_u_h = - input * sobel_filter.grad_h(output[:, [0]])
    K_u_v = - input * sobel_filter.grad_v(output[:, [0]])
    sigma = output[:, [1, 2]]
    rhs = sigma + beta1 * torch.sqrt(input) * sigma ** 2 + beta2 * input * sigma ** 3
    return ((K_u_h - rhs[:, [0]])** 2 + (K_u_v - rhs[:, [1]]) ** 2).mean()

def conv_constitutive_constraint_nonlinear_exp(input, output, sobel_filter):
    """Nonlinear extension of Darcy's law
        sigma = - exp(K * u) grad(u)

    Args:
        input: K
        output: u, sigma1, sigma2
    """
    grad_h = sobel_filter.grad_h(output[:, [0]])
    grad_v = sobel_filter.grad_v(output[:, [0]])

    sigma_h = - torch.exp(input * output[:, [0]]) * grad_h
    sigma_v = - torch.exp(input * output[:, [0]]) * grad_v


    return ((output[:, [1]] - sigma_h) ** 2 
        + (output[:, [2]] - sigma_v) ** 2).mean()

def conv_continuity_constraint(output, sobel_filter, use_tb=True):
    """y
    div(sigma) = -f

    Args:

    """
    sigma1_x1 = sobel_filter.grad_h(output[:, [1]])
    sigma2_x2 = sobel_filter.grad_v(output[:, [2]])
    # leave the top and bottom row free, since sigma2_x2 is almost 0,
    # don't want to enforce sigma1_x1 to be also zero.
    if use_tb:
        return ((sigma1_x1 + sigma2_x2) ** 2).mean()
    else:
        return ((sigma1_x1 + sigma2_x2) ** 2)[:, :, 1:-1, :].mean()

def th_conv_boundary_condition(output,temp_data):
    #print("output shape: ", output.shape)
    left_bound, right_bound = output[:, 0, :, 0], output[:, 0, :, -1]
    top_down_flux = output[:, 0, [0, -1], :]
    heat_pump_flux = output[:,0, 7, 4] #output[:, 2, [0, -1], :]
    loss_dirichlet = 0
    # print(f'output: {output.shape}') 
    # print(f'temp_data: {temp_data.shape}') 

    #for i in range(0,64):
    #  for j in range(0,64):
    #    loss_dirichlet += abs(output[:, 0, i, j] - temp_data[i,j])
    #print(top_down_flux)
    #loss_dirichlet = F.mse_loss(left_bound, 0*torch.ones_like(left_bound)) \
    #    + F.mse_loss(top_down_flux, 0*torch.ones_like(top_down_flux)) \
    #    + F.mse_loss(right_bound, 0*torch.ones_like(right_bound)) #\
    loss_dirichlet = F.mse_loss(output[:, :, :, :], temp_data[:,:,:,:])
        
        #+ F.mse_loss(right_bound, 6*torch.ones_like(right_bound))
    loss_neumann = 0#F.mse_loss(heat_pump_flux, 5*torch.ones_like(heat_pump_flux)) 
    #F.mse_loss(top_down_flux, torch.zeros_like(top_down_flux))
    #print(loss_neumann)
    # torch.ones_like(left_bound) is the left boundary condition of value 1. 

    return loss_dirichlet , loss_neumann
