import torch

class AdaLN(torch.nn.Module):
    def __init__(self,conditionC,channels,skip=False):
        super().__init__()
        self.norm=torch.nn.GroupNorm(1,channels)
        self.proj1=torch.nn.Conv2d(conditionC,channels,kernel_size=1)
        self.proj2=torch.nn.Conv2d(conditionC,channels,kernel_size=1)
        self.skip=skip
    
    def forward(self,input, condition):
        w=self.proj1(condition)
        b=self.proj2(condition)
        if self.skip:
            return input
        out=self.norm(input)*w+b
        return out

class LayerScale(torch.nn.Module):
    def __init__(self,channels,init_values = 1e-5):
        super().__init__()
        self.gamma = torch.nn.Parameter(init_values * torch.ones(channels))

    def forward(self, x):
        return x * self.gamma.view(1,-1,1,1)
    
class RoPEAttention(torch.nn.Module):
    """
    Rotary Position Embedding (RoPE) Attention module.

    This module implements a multi-head attention mechanism for vision transformers with Rotary Position Embeddings.
    It applies RoPE to the query and key projections before computing attention, uses QK norm,
    Flash Attention v2 and scales alpha to allow interpolation and extrapolation to new resolutions

    Args:
        channels (int): Number of input and output channels.
        heads (int): Number of attention heads.
        baseSize (int): The base resolution RoPE will be scaled around.
    """
    def __init__(self, channels, heads, baseSize):
        super().__init__()
        self.channels = channels
        self.heads = heads
        self.baseSize=baseSize
        self.headDim = channels // heads

        self.qProj = torch.nn.Conv2d(channels, channels, 1)
        self.kProj = torch.nn.Conv2d(channels, channels, 1)
        self.vProj = torch.nn.Conv2d(channels, channels, 1)
        self.outProj = torch.nn.Conv2d(channels, channels, 1)

        self.qLayerNorm=torch.nn.GroupNorm(1,heads)
        self.kLayerNorm=torch.nn.GroupNorm(1,heads)

    def createRotationMatrix(self, height, width, channels):
        scalingH=height/self.baseSize if height>self.baseSize else 1.0
        scalingW=width/self.baseSize if width>self.baseSize else 1.0
        t = torch.arange(height * width, dtype=torch.float32)
        tX = (t % width).float()
        tY = torch.div(t, width, rounding_mode='floor').float()

        freqsX = 1.0 / ((100.0*scalingW) ** (torch.arange(0, channels, 2).float() / channels))
        freqsY = 1.0 / ((100.0*scalingH) ** (torch.arange(0, channels, 2).float() / channels))
        freqsX = torch.outer(tX, freqsX)
        freqsY = torch.outer(tY, freqsY)

        freqsCisX = torch.polar(torch.ones_like(freqsX), freqsX)
        freqsCisY = torch.polar(torch.ones_like(freqsY), freqsY)

        freqsCis = torch.cat([freqsCisX, freqsCisY], dim=-1)
        freqsCis = freqsCis.view(height, width, channels)

        return freqsCis

    def rotate(self, q, k):
        batchSize, channels, height, width = q.shape
        freqsCis = self.createRotationMatrix(height, width, channels).to(q.device)

        q = q.permute(0, 2, 3, 1).contiguous()
        k = k.permute(0, 2, 3, 1).contiguous()

        qRot = q * freqsCis.real - q.roll(shifts=1, dims=-1) * freqsCis.imag
        kRot = k * freqsCis.real - k.roll(shifts=1, dims=-1) * freqsCis.imag

        return qRot.permute(0, 3, 1, 2), kRot.permute(0, 3, 1, 2)

    def forward(self, x):
        batchSize, channels, height, width = x.shape
        
        q = self.qProj(x)
        k = self.kProj(x)
        v = self.vProj(x)

        q, k = self.rotate(q, k)

        q = q.view(batchSize, self.heads, self.headDim, height * width).transpose(2, 3)
        k = k.view(batchSize, self.heads, self.headDim, height * width).transpose(2, 3)
        v = v.view(batchSize, self.heads, self.headDim, height * width).transpose(2, 3).contiguous()

        q=self.qLayerNorm(q)
        k=self.kLayerNorm(k)

        #Temporarily converting to bf16 to use Flash Attention v2
        with torch.nn.attention.sdpa_kernel(torch.nn.attention.SDPBackend.FLASH_ATTENTION), torch.amp.autocast('cuda',enabled=True,dtype=torch.bfloat16):
            attnOutput = torch.nn.functional.scaled_dot_product_attention(q, k, v)
        attnOutput = attnOutput.float().transpose(2, 3).contiguous().view(batchSize, channels, height, width)

        return self.outProj(attnOutput)

class DiTBlock(torch.nn.Module):
    def __init__(self,channels,heads,conditionC):
        super().__init__()
        self.norm1=AdaLN(conditionC,channels)
        self.norm2=torch.nn.GroupNorm(1,channels,eps=1e-6)
        self.attn=RoPEAttention(channels,heads,8)
            
        self.mlp=torch.nn.ModuleDict({
            "fc1":torch.nn.Linear(channels,channels*4),
            "fc2":torch.nn.Linear(channels*4,channels)
        })
        self.ls1=LayerScale(channels,init_values=0.0)
        self.ls2=LayerScale(channels,init_values=0.0)

    def forward(self,x,condition):
        input=x
        x=self.norm1(x,condition)
        x=self.attn(x) 
        x=self.ls1(x)+input

        input=x
        x=self.norm2(x)
        x=x.permute(0,2,3,1)
        x=self.mlp.fc1(x)
        x=torch.nn.functional.gelu(x)
        x=self.mlp.fc2(x)
        x=x.permute(0,3,1,2)
        x=self.ls2(x)+input
        return x
    
class DiT(torch.nn.Module):
    """
    Diffusion Transformer (DiT) model.

    This class implements a Diffusion Transformer, which is a transformer-based architecture
    designed for image generation. It applies a series of DiT blocks to process the input, 
    incorporating time step and condition information.

    Args:
        channels (int): Number of channels in each transformer block.
        nBlocks (int): Number of DiT blocks in the model.
        nHeads (int): Number of attention heads
        inC (int): Number of input channels.
        outC (int): Number of output channels.
        conditionC (int): Number of conditioning channels
        patchSize (int, optional): Size of patches for patchification. Defaults to 1. Increasing this
        will speed up training and lower VRAM requirements, at the expense of generation quality
        and potential artifacts. The final convolution layer attempts to mitigate this effect
    """
    def __init__(self,channels,nBlocks,inC,outC,conditionC,nHeads=8,patchSize=1):
        super().__init__()
        self.patchify=torch.nn.PixelUnshuffle(patchSize)
        self.inProj=torch.nn.Conv2d(inC*patchSize**2,channels,kernel_size=1)
        self.blocks=torch.nn.ModuleList()
        for i in range(nBlocks):
            self.blocks.append(DiTBlock(channels,nHeads,conditionC))
        self.outProj=torch.nn.Conv2d(channels,outC*patchSize**2,kernel_size=1)
        self.unpatchify=torch.nn.PixelShuffle(patchSize)
        self.finalConv=torch.nn.Conv2d(outC,outC,kernel_size=3,padding='same')

    def forward(self,x,condition):
        #print
        x=self.patchify(x)
        #condition=self.patchify(condition)
        #print(condition.shape)
        x=self.inProj(x)
        for block in self.blocks:
            x=block(x,condition)
        x=self.outProj(x)
        x=self.unpatchify(x)
        x=self.finalConv(x)
        return x


import torch
from tqdm.auto import tqdm 

class RectifiedFlow():
    def __init__(self, model, total_time, lr=3e-4, device='cpu', optimizer=None, ema_decay=0.0):
        self.model = model
        if ema_decay > 0.0:
            self.ema_model = torch.optim.swa_utils.AveragedModel(
                self.model,
                multi_avg_fn=torch.optim.swa_utils.get_ema_multi_avg_fn(ema_decay)
            )
        else:
            self.ema_model = self.model

        self.total_time = total_time
        self.device = device

        if optimizer is None:
            self.optimizer = torch.optim.AdamW(self.model.parameters(), lr=lr)
        else:
            self.optimizer = optimizer

    def state_dict(self):
        return {
            'model': self.model.state_dict(),
            'ema': self.ema_model.state_dict(),
            'optimizer': self.optimizer.state_dict()
        }
    
    def load_state_dict(self, state_dict):
        self.model.load_state_dict(state_dict['model'])
        self.ema_model.load_state_dict(state_dict['ema'])
        self.optimizer.load_state_dict(state_dict['optimizer'])
    
    def add_noise(self, images, time_steps):
        noise_source = torch.randn_like(images).to(self.device)
        noisy_images = (1 - time_steps) * images + time_steps * noise_source
        return noisy_images, noise_source
    
    def predict_velocity(self, noisy_images, time_steps, condition=None):
        batch_size = noisy_images.shape[0]
        if time_steps.shape[0] == 1 and time_steps.shape[0] != batch_size:
            time_steps = time_steps.repeat(batch_size, *([1] * (time_steps.dim() - 1)))
        
        if condition is None:
            condition_input = time_steps
        else:
            if condition.shape[0] == 1 and condition.shape[0] != batch_size:
                condition = condition.repeat(batch_size, *([1] * (condition.dim() - 1)))
            condition_input = torch.cat([time_steps, condition], dim=1)

        if torch.is_inference_mode_enabled():
            velocity_prediction = self.ema_model(noisy_images, condition_input)
        else:
            velocity_prediction = self.model(noisy_images, condition_input)
            
        return velocity_prediction
    
    def generate_samples(self, total_steps, shape=None, condition=None, sampling_step_function=None):
        # sampling_step_function is a function with the same signature as euler_sampling_step, returning the sample at t-1.
        # It can be used to implement CFG and schedulers
        if sampling_step_function is None:
            sampling_step_function = self.euler_sampling_step
        
        current_sample = torch.randn(shape).to(self.device)
        with torch.no_grad():
            for step_index in tqdm(range(total_steps), ascii=" ▖▘▝▗▚▞█", disable=False):
                current_sample = sampling_step_function(
                    step_index, shape, current_sample, total_steps, condition
                )
        
        return current_sample
    
    def generate_samples2(self, total_steps, shape=None, condition=None, sampling_step_function=None):
        num_steps = 5
        batch_size = shape[0]  # Number of samples in the batch
        noise_source = torch.randn(shape).to(self.device)
        with torch.no_grad():
            for i in tqdm(range(num_steps), ascii=" ▖▘▝▗▚▞█", disable=False):
                noise_source = self.euler_sampling_step(
                    i, shape, noise_source, 5, condition
                )
                continue
                time_step_value = 1 - i / num_steps  # Linearly decreasing time step from 1 to 0

                # Step 3: Create tensors for 'time_step' and 'delta_t_tensor' to match the dimensions of the model input.
                # These tensors have shape [batch_size, 1, 1, 1] to broadcast correctly during computations.
                time_step = torch.full((batch_size, 1, 1, 1), time_step_value).to(self.device)

                # Step 4: Use the model to predict the velocity (noise) at the current sample and time step.
                # The 'predict_velocity' function outputs the estimated noise to be removed.
                noise_source = self.predict_velocity(noise_source, time_step, condition=condition)
        
        return noise_source
    
    def train_step2(self, real_data, condition=None, validation_time_step=None, classifier=None):
        num_steps = 5
        batch_size = real_data.shape[0]  # Number of samples in the batch
        noise_source = torch.randn_like(real_data).to(self.device)
        for i in range(num_steps):
            extra_noise = torch.randn_like(real_data).to(self.device)
            noise_source = 0.1 * real_data + 0.9 * noise_source
            noise_source = self.euler_sampling_step(
                    i, real_data.shape, noise_source, 5, condition
                )
            continue
            time_step_value = 1 - i / num_steps  # Linearly decreasing time step from 1 to 0

            # Step 3: Create tensors for 'time_step' and 'delta_t_tensor' to match the dimensions of the model input.
            # These tensors have shape [batch_size, 1, 1, 1] to broadcast correctly during computations.
            time_step = torch.full((batch_size, 1, 1, 1), time_step_value).to(self.device)

            # Step 4: Use the model to predict the velocity (noise) at the current sample and time step.
            # The 'predict_velocity' function outputs the estimated noise to be removed.
            noise_source = self.predict_velocity(noise_source, time_step, condition=condition)

        loss = ((noise_source - real_data) ** 2).mean(dim=(2, 3)).mean()

        if validation_time_step is None:
            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()
            if hasattr(self, 'ema_model') and self.ema_model is not self.model:
                self.ema_model.update_parameters(self.model)

        return {'loss': loss.detach()}
    
    def train_step(self, real_data, condition=None, validation_time_step=None, classifier=None):
        batch_size = real_data.shape[0]
        time_steps = torch.rand(batch_size, 1, 1, 1).to(self.device)
        
        if validation_time_step is not None:
            time_steps = torch.full((batch_size, 1, 1, 1), validation_time_step).to(self.device)

        noisy_data, random_noise = self.add_noise(real_data, time_steps)
        velocity_prediction = self.predict_velocity(noisy_data, time_steps, condition=condition)
        loss = ((random_noise - real_data - velocity_prediction) ** 2).mean(dim=(2, 3)).mean()

        if validation_time_step is None:
            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()
            if hasattr(self, 'ema_model') and self.ema_model is not self.model:
                self.ema_model.update_parameters(self.model)

        return {'loss': loss.detach()}

    def euler_sampling_step(self, step_index, input_shape, current_sample, total_steps, condition):
        """
        Perform a single sampling step using Euler's method.

        Args:
            step_index (int): The current step index in the sampling process.
            input_shape (torch.Size): The shape of the input tensor (batch size, channels, height, width).
            current_sample (torch.Tensor): The current sample at time t.
            total_steps (int): The total number of sampling steps (N).
            condition (torch.Tensor, optional): Optional conditioning information.
                           
        Returns:
            torch.Tensor: The next sample at time t - Δt.
        """
        # Step 1: Compute the current time step 't' as a scalar value decreasing from 1 to 0.
        time_step_value = 1 - step_index / total_steps  # Linearly decreasing time step from 1 to 0
        
        # Step 2: Compute the time increment 'Δt' for each step.
        delta_t = 1 / total_steps  # Constant time decrement per step
        
        batch_size = input_shape[0]  # Number of samples in the batch

        # Step 3: Create tensors for 'time_step' and 'delta_t_tensor' to match the dimensions of the model input.
        # These tensors have shape [batch_size, 1, 1, 1] to broadcast correctly during computations.
        time_step = torch.full((batch_size, 1, 1, 1), time_step_value).to(self.device)
        delta_t_tensor = torch.full((batch_size, 1, 1, 1), delta_t).to(self.device)

        # Step 4: Use the model to predict the velocity (noise) at the current sample and time step.
        # The 'predict_velocity' function outputs the estimated noise to be removed.
        velocity_prediction = self.predict_velocity(current_sample, time_step, condition=condition)

        # Step 5: Update the sample using Euler's method.
        # The next sample is obtained by subtracting the product of 'delta_t_tensor' and 'velocity_prediction' from 'current_sample'.
        # This simulates the reverse diffusion, incrementally denoising the sample.
        next_sample = current_sample - delta_t_tensor * velocity_prediction

        # Step 6: Return the updated sample for the next iteration.
        # This sample will be used as 'current_sample' in the next step of the loop.
        return next_sample









change code so it will add random noise levels like so that each pixel has its own noise value