# MIT License

# Copyright (c) 2021 Pashupati Hegde.
# Copyright (c) 2023 Magnus Ross.

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import torch
import torch.nn as nn
from torchdiffeq import odeint as odeint_nonadjoint
from torchdiffeq import odeint_adjoint


class ODEfunc(nn.Module):
    def __init__(self, diffeq):
        """
        Defines the ODE function:
            mainly calls layer.build_cache() method to fix the draws from random variables.
        Modified from https://github.com/rtqichen/ffjord/

        @param diffeq: Layer of GPODE/npODE/neuralODE
        """
        super(ODEfunc, self).__init__()
        self.diffeq = diffeq
        self.register_buffer("_num_evals", torch.tensor(0.0))

    def before_odeint(self, rebuild_cache):
        self._num_evals.fill_(0)
        if rebuild_cache:
            self.diffeq.build_cache()

    def num_evals(self):
        return self._num_evals.item()

    def forward(self, t, states):
        self._num_evals += 1
        dy = self.diffeq(t, states)
        return dy


class Flow(nn.Module):
    def __init__(
        self, diffeq, solver="dopri5", atol=1e-6, rtol=1e-6, use_adjoint=False
    ):
        """
        Defines an ODE flow:
            mainly defines forward() method for forward numerical integration of an ODEfunc object
        See https://github.com/rtqichen/torchdiffeq for more information on numerical ODE solvers.

        @param diffeq: Layer of GPODE/npODE/neuralODE
        @param solver: Solver to be used for ODE numerical integration
        @param atol: Absolute tolerence for the solver
        @param rtol: Relative tolerence for the solver
        @param use_adjoint: Use adjoint method for computing loss gradients, calls odeint_adjoint fro torchdiffeq
        """
        super(Flow, self).__init__()
        self.odefunc = ODEfunc(diffeq)
        self.solver = solver
        self.atol = atol
        self.rtol = rtol
        self.use_adjoint = use_adjoint

    def forward(self, x0, ts, return_energy=False):
        """
        Numerical solution of an IVP, and optionally compute divergence term for density transformation computation
        @param x0: Initial state (N,D) tensor x(t_0).
        @param ts: Time sequence of length T, first value is considered as t_0
        @param return_divergence: Bool flag deciding the divergence computation
        @return: xs: (N,T,D) tensor
        """
        odeint = odeint_adjoint if self.use_adjoint else odeint_nonadjoint
        self.odefunc.before_odeint(rebuild_cache=True)
        if return_energy:
            xs = odeint(
                self.odefunc, x0, ts, atol=self.atol, rtol=self.rtol, method=self.solver
            )
            energy = self.odefunc.diffeq.hamiltonian(
                ts, xs.reshape(-1, xs.shape[-1])
            ).reshape(xs.shape[0], xs.shape[1], 1)

            return xs.permute(1, 0, 2), energy.permute(1, 0, 2)  # (N,T,D), # (N,T,1)

        else:
            xs = odeint(
                self.odefunc, x0, ts, atol=self.atol, rtol=self.rtol, method=self.solver
            )
            return xs.permute(1, 0, 2)  # (N,T,D)

    def num_evals(self):
        return self.odefunc.num_evals()

    def kl(self):
        """
        Calls KL() computation from the diffeq layer
        """
        return self.odefunc.diffeq.kl().sum()

    def log_prior(self):
        """
        Calls log_prior() computation from the diffeq layer
        """
        return self.odefunc.diffeq.log_prior().sum()
