"""R^3 diffusion methods."""
import numpy as np
from scipy.special import gamma
import torch


class R3Diffuser:
    """VP-SDE diffuser class for translations."""

    def __init__(self, r3_conf):
        """
        Args:
            min_b: starting value in variance schedule.
            max_b: ending value in variance schedule.
        """
        self._r3_conf = r3_conf
        self.min_b = r3_conf.min_b
        self.max_b = r3_conf.max_b

    def _scale(self, x):
        return x * self._r3_conf.coordinate_scaling

    def _unscale(self, x):
        return x / self._r3_conf.coordinate_scaling

    def b_t(self, t):
        if np.any(t < 0) or np.any(t > 1):
            raise ValueError(f'Invalid t={t}')
        return self.min_b + t*(self.max_b - self.min_b)

    def diffusion_coef(self, t):
        """Time-dependent diffusion coefficient."""
        return np.sqrt(self.b_t(t))

    def drift_coef(self, x, t):
        """Time-dependent drift coefficient."""
        return -1/2 * self.b_t(t) * x

    def sample_ref(self, n_samples: float=1, motif_mask: np.array=None, scale: float=None):
        if scale is not None:
            # due to 3 sigma
            #noise = np.random.normal(size=(n_samples, 3), scale=scale)
            noise = np.random.normal(size=(n_samples, 3))*scale
        else:
            noise = np.random.normal(size=(n_samples, 3))
        noise = self.apply_motif_mask(motif_mask, noise)
        return noise

    def marginal_b_t(self, t):
        return t*self.min_b + (1/2)*(t**2)*(self.max_b-self.min_b)

    def calc_trans_0(self, score_t, x_t, t, use_torch=True):
        beta_t = self.marginal_b_t(t)
        beta_t = beta_t[..., None, None]
        exp_fn = torch.exp if use_torch else np.exp
        cond_var = 1 - exp_fn(-beta_t)
        return (score_t * cond_var + x_t) / exp_fn(-1/2*beta_t)

    def forward(self, x_t_1: np.ndarray, t: float, num_t: int):
        """Samples marginal p(x(t) | x(t-1)).

        Args:
            x_0: [..., n, 3] initial positions in Angstroms.
            t: continuous time in [0, 1].

        Returns:
            x_t: [..., n, 3] positions at time t in Angstroms.
            score_t: [..., n, 3] score at time t in scaled Angstroms.
        """
        if not np.isscalar(t):
            raise ValueError(f'{t} must be a scalar.')
        x_t_1 = self._scale(x_t_1)
        b_t = torch.tensor(self.marginal_b_t(t) / num_t).to(x_t_1.device)
        print("------ r3_diffuser, 67 -------")
        z_t_1 = torch.tensor(np.random.normal(size=x_t_1.shape)).to(x_t_1.device)
        x_t = torch.sqrt(1 - b_t) * x_t_1 + torch.sqrt(b_t) * z_t_1
        return x_t

    def distribution(self, x_t, score_t, t, mask, dt):
        x_t = self._scale(x_t)
        g_t = self.diffusion_coef(t)
        f_t = self.drift_coef(x_t, t)
        std = g_t * np.sqrt(dt)
        mu = x_t - (f_t - g_t**2 * score_t) * dt
        if mask is not None:
            mu *= mask[..., None]
        return mu, std

    def apply_motif_mask_3d(self, motif_seq_mask, sampled_rots):
        B, num_motif, seq_len = motif_seq_mask.shape
        sequence_mask = np.ones((B, seq_len))
        result_noise = np.zeros_like(sampled_rots)
        for i in range(num_motif):
            sequence_mask *= sequence_mask*(~motif_seq_mask[:,i,:]) #B, seq_len
            dino = motif_seq_mask[:,i,:].sum(axis=-1)[:,np.newaxis]
            up = (motif_seq_mask[:,i,:][:,:,np.newaxis]*sampled_rots).sum(axis=1) # B,dim
            value = (up/dino)[:,np.newaxis,:] # B,1,dim
            result_noise += value * motif_seq_mask[:,i,:][:,:,np.newaxis]
        result_noise += sampled_rots * sequence_mask[:,:, np.newaxis]
        return result_noise

    def apply_motif_mask(self, motif_seq_mask, sampled_rots):
        sequence_mask = np.ones((motif_seq_mask.shape[-1]))
        result_noise = np.zeros_like(sampled_rots)
        for i in range(motif_seq_mask.shape[0]):
            sequence_mask = sequence_mask*(~motif_seq_mask[i,:])
            result_noise += (motif_seq_mask[i,:][:, np.newaxis])*(((np.sum(sampled_rots*motif_seq_mask[i,:][:, np.newaxis], 0))/np.sum(motif_seq_mask[i,:]))[np.newaxis, :])
        result_noise = sequence_mask[:, np.newaxis]*sampled_rots + result_noise
        return result_noise

    def forward_marginal(self, x_0: np.ndarray, t: float, motif_seq_mask: np.ndarray, rot_cause_trans: np.ndarray):
        """Samples marginal p(x(t) | x(0)).

        Args:
            x_0: [..., n, 3] initial positions in Angstroms.
            t: continuous time in [0, 1].

        Returns:
            x_t: [..., n, 3] positions at time t in Angstroms.
            score_t: [..., n, 3] score at time t in scaled Angstroms.
        """
        if not np.isscalar(t):
            raise ValueError(f'{t} must be a scalar.')
        x_0 = self._scale(x_0)
        
        alpha = -1/2*self.marginal_b_t(t)
        beta = (1 - np.exp(-self.marginal_b_t(t)))

        noise = np.random.normal(
            #loc=np.exp(alpha) * x_0,
            loc = 0,
            scale=np.sqrt(beta),
            size=x_0.shape
        )
        noise_mask = self.apply_motif_mask(motif_seq_mask, noise)

        x_t = np.exp(alpha) * x_0 + noise_mask
        score_t = self.score(x_t, x_0, t)
        x_t = self._unscale(x_t)
        return x_t, score_t

    def score_scaling(self, t: float):
        return 1 / np.sqrt(self.conditional_var(t))

    def reverse(
            self,
            *,
            x_t: np.ndarray,
            score_t: np.ndarray,
            t: float,
            dt: float,
            mask: np.ndarray=None,
            center: bool=True,
            noise_scale: float=1.0,
            motif_mask: np.array=None,
            rot_increase_trans: np.array=None
        ):
        """Simulates the reverse SDE for 1 step

        Args:
            x_t: [..., 3] current positions at time t in angstroms.
            score_t: [..., 3] rotation score at time t.
            t: continuous time in [0, 1].
            dt: continuous step size in [0, 1].
            mask: True indicates which residues to diffuse.

        Returns:
            [..., 3] positions at next step t-1.
        """
        if not np.isscalar(t):
            raise ValueError(f'{t} must be a scalar.')
        x_t -= rot_increase_trans
        x_t = self._scale(x_t)
        g_t = self.diffusion_coef(t)
        f_t = self.drift_coef(x_t, t)

        noise = np.random.normal(size=score_t.shape)
        #noise = self.apply_motif_mask_3d(motif_mask, noise)
        z = noise_scale * noise

        perturb = (f_t - g_t**2 * score_t) * dt + g_t * np.sqrt(dt) * z
        perturb = self.apply_motif_mask_3d(motif_mask, perturb)
        if mask is not None:
            perturb *= mask[..., None]
        else:
            mask = np.ones(x_t.shape[:-1])
        x_t_1 = x_t - perturb
        if center:
            com = np.sum(x_t_1, axis=-2) / np.sum(mask, axis=-1)[..., None]
            x_t_1 -= com[..., None, :]
        x_t_1 = self._unscale(x_t_1)
        return x_t_1+rot_increase_trans
        #return x_t_1

    def conditional_var(self, t, use_torch=False):
        """Conditional variance of p(xt|x0).

        Var[x_t|x_0] = conditional_var(t)*I

        """
        if use_torch:
            return 1 - torch.exp(-self.marginal_b_t(t))
        return 1 - np.exp(-self.marginal_b_t(t))

    def score(self, x_t, x_0, t, use_torch=False, scale=False):
        if use_torch:
            exp_fn = torch.exp
        else:
            exp_fn = np.exp
        if scale:
            x_t = self._scale(x_t)
            x_0 = self._scale(x_0)
        return -(x_t - exp_fn(-1/2*self.marginal_b_t(t)) * x_0) / self.conditional_var(t, use_torch=use_torch)
