import torch
import torch.nn as nn
import numpy as np
import torch.nn.functional as F


def mean_flat(x):
    """
    Take the mean over all non-batch dimensions.
    """
    return torch.mean(x, dim=list(range(1, len(x.size()))))

def sum_flat(x):
    """
    Take the mean over all non-batch dimensions.
    """
    return torch.sum(x, dim=list(range(1, len(x.size()))))

# ------------------------
# T2I Loss Function
# ------------------------
class SILoss:
    def __init__(
            self,
            prediction='v',
            path_type="linear",
            weighting="uniform",
            encoders=[], 
            accelerator=None, 
            latents_scale=None, 
            latents_bias=None,
            ):
        self.prediction = prediction
        self.weighting = weighting
        self.path_type = path_type
        self.encoders = encoders
        self.accelerator = accelerator
        self.latents_scale = latents_scale
        self.latents_bias = latents_bias

    def interpolant(self, t):
        if self.path_type == "linear":
            alpha_t = 1 - t
            sigma_t = t
            d_alpha_t = -1
            d_sigma_t =  1
        elif self.path_type == "cosine":
            alpha_t = torch.cos(t * np.pi / 2)
            sigma_t = torch.sin(t * np.pi / 2)
            d_alpha_t = -np.pi / 2 * torch.sin(t * np.pi / 2)
            d_sigma_t =  np.pi / 2 * torch.cos(t * np.pi / 2)
        else:
            raise NotImplementedError()

        return alpha_t, sigma_t, d_alpha_t, d_sigma_t

    def __call__(self, model, images, model_kwargs=None, zs=None, proj_s=None, denoise_only=False):
        if model_kwargs == None:
            model_kwargs = {}
        # sample timesteps
        if self.weighting == "uniform":
            time_input = torch.rand((images.shape[0], 1, 1, 1))
        elif self.weighting == "lognormal":
            # sample timestep according to log-normal distribution of sigmas following EDM
            rnd_normal = torch.randn((images.shape[0], 1 ,1, 1))
            sigma = rnd_normal.exp()
            if self.path_type == "linear":
                time_input = sigma / (1 + sigma)
            elif self.path_type == "cosine":
                time_input = 2 / np.pi * torch.atan(sigma)
                
        time_input = time_input.to(device=images.device, dtype=images.dtype)
        
        noises = torch.randn_like(images)
        alpha_t, sigma_t, d_alpha_t, d_sigma_t = self.interpolant(time_input)
            
        model_input = alpha_t * images + sigma_t * noises
        if self.prediction == 'v':
            model_target = d_alpha_t * images + d_sigma_t * noises
        else:
            raise NotImplementedError() 
        try:
            model_output, zs_tilde, proj_tilde  = model(model_input, time_input.flatten(), return_act=True, **model_kwargs)
        except:
            model_output, zs_tilde = model(model_input, time_input.flatten(), return_act=True, **model_kwargs)
        denoising_loss = mean_flat((model_output - model_target) ** 2)
        
        # zs: list, each element is a list of attentions for each layer, attn shape: [bsz, n_heads, n_tokens, n_tokens]
        # for base model [256, 12, 256, 256]; zs [1, 12, 256, 12, 256, 256]
        # zs_tilde: list of models attn [4, 256, 12, 256, 256]
        
        if not denoise_only:
            zs = zs[0] # [12, 256, 12, 256, 256]
        zs_ = [zs_tilde[0].shape]
        print()
        proj_loss = 0.
        attn_loss = 0.

        # attention alignment: default : 4567 from SiT-XL/2 <- 89X11 from dinov2-b
        if not denoise_only:
            for i in range(4, 8):
                zsi= (zs_tilde[i])[:, :12, :, :].reshape(-1, 256)
                zs_j = zs[i + 4].reshape(-1, 256).softmax(dim=-1)
                attn_loss += -(zs_j * torch.log(zsi.softmax(dim=-1))).sum(dim=-1).mean()
            attn_loss /= 4.0
            attn_loss = attn_loss

        # feature alignment loss; default depth 8
        if not denoise_only:
            bsz = proj_s[0].shape[0]
            for i, (z, z_tilde) in enumerate(zip(proj_s, proj_tilde)):
                for j, (z_j, z_tilde_j) in enumerate(zip(z, z_tilde)):
                    z_tilde_j = torch.nn.functional.normalize(z_tilde_j, dim=-1) 
                    z_j = torch.nn.functional.normalize(z_j, dim=-1) 
                    proj_loss += mean_flat(-(z_j * z_tilde_j).sum(dim=-1))
            proj_loss /= (len(proj_s) * bsz)

        return denoising_loss, attn_loss, proj_loss
