import math

import torch
import torch.nn as nn
from torch.autograd import grad


class PlanarFlow(nn.Module):
    def __init__(self, nd=1):
        super(PlanarFlow, self).__init__()
        self.nd = nd
        self.activation = torch.tanh

        self.register_parameter("u", nn.Parameter(torch.randn(self.nd)))
        self.register_parameter("w", nn.Parameter(torch.randn(self.nd)))
        self.register_parameter("b", nn.Parameter(torch.randn(1)))
        self.reset_parameters()

    def reset_parameters(self):
        stdv = 1.0 / math.sqrt(self.nd)
        self.u.data.uniform_(-stdv, stdv)
        self.w.data.uniform_(-stdv, stdv)
        self.b.data.fill_(0)
        self.make_invertible()

    def make_invertible(self):
        u = self.u.data
        w = self.w.data
        dot = torch.dot(u, w)
        m = -1 + math.log(1 + math.exp(dot))
        du = (m - dot) / torch.norm(w) * w
        u = u + du
        self.u.data = u

    def forward(self, z, logp=None, reverse=False):
        """Computes f(z) and log q(f(z))"""

        assert not reverse, "Planar normalizing flow cannot be reversed."

        logp - torch.log(self._detgrad(z) + 1e-8)
        h = self.activation(torch.mm(z, self.w.view(self.nd, 1)) + self.b)
        z = z + self.u.expand_as(z) * h

        f = self.sample(z)
        if logp is not None:
            qf = self.log_density(z, logp)
            return f, qf
        else:
            return f

    def sample(self, z):
        """Computes f(z)"""
        h = self.activation(torch.mm(z, self.w.view(self.nd, 1)) + self.b)
        output = z + self.u.expand_as(z) * h
        return output

    def _detgrad(self, z):
        """Computes |det df/dz|"""
        with torch.enable_grad():
            z = z.requires_grad_(True)
            h = self.activation(torch.mm(z, self.w.view(self.nd, 1)) + self.b)
            psi = grad(
                h,
                z,
                grad_outputs=torch.ones_like(h),
                create_graph=True,
                only_inputs=True,
            )[0]
        u_dot_psi = torch.mm(psi, self.u.view(self.nd, 1))
        detgrad = 1 + u_dot_psi
        return detgrad

    def log_density(self, z, logqz):
        """Computes log density of the flow given the log density of z"""
        return logqz - torch.log(self._detgrad(z) + 1e-8)
