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,
                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.shape[0]  # Number of samples in the batch
        print(batch_size)
        noise_source = torch.randn_like(shape).to(self.device)
        with torch.no_grad():
            for i in tqdm(range(num_steps), ascii=" ▖▘▝▗▚▞█", disable=False):
                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):
            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)


        # 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 = ((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







Traceback (most recent call last):
  File "/home/muinez/OMI-Training-Pipeline/pipelines/architecture-experiments/main.py", line 154, in <module>
    main()
  File "/home/muinez/OMI-Training-Pipeline/pipelines/architecture-experiments/main.py", line 125, in main
    metrics=modelT.train_step2(image.to('cuda'))
  File "/home/muinez/OMI-Training-Pipeline/pipelines/architecture-experiments/OMI/RectifiedFlow.py", line 123, in train_step2
    self.ema_model.update_parameters(self.model)
  File "/home/muinez/.local/lib/python3.10/site-packages/torch/optim/swa_utils.py", line 215, in update_parameters
    p_averaged.detach().copy_(self.avg_fn(p_averaged.detach(), p_model, n_averaged))
  File "/home/muinez/.local/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/home/muinez/.local/lib/python3.10/site-packages/torch/optim/swa_utils.py", line 29, in ema_update
    torch._foreach_lerp_(ema_param_list, current_param_list, 1 - decay)
TypeError: _foreach_lerp_() received an invalid combination of arguments - got (Tensor, Tensor, float), but expected one of:
 * (tuple of Tensors self, tuple of Tensors tensors1, Number weight)
      didn't match because some of the arguments have invalid types: (Tensor, Tensor, float)
 * (tuple of Tensors self, tuple of Tensors tensors1, tuple of Tensors weights)
      didn't match because some of the arguments have invalid types: (Tensor, Tensor, float)