# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/02_normalizers.ipynb.

# %% auto 0
__all__ = ['GuidanceTfm', 'BaseNormGuidance', 'TNormGuidance', 'FullNormGuidance']

# %% ../nbs/02_normalizers.ipynb 3
'''Code for blog post:
    https://enzokro.dev/blog/posts/2022-11-15-guidance-expts-1
'''
import math
import torch

from cf_guidance.schedules import name2schedule


# %% ../nbs/02_normalizers.ipynb 4
class GuidanceTfm:
    "Baseline Classifier-free Guidance for Difussion."
    name = "CFGuidance"
    def __init__(self, schedule_name, *args, **kwargs):
        schedule_kls = name2schedule[schedule_name]
        self.schedules = schedule_kls()
        
    def encode(self, u, t, idx=None):
        "Applies guidance on `u` and `t` with optional pre/post processing."
        self.pre_proc(u, t, idx)
        self.guide(u, t, idx)
        self.post_proc(u, t, idx)
        return self.pred
    
    def guide(self, u, t, idx=None):
        "Mixes latents `u` and `t` based on guidance schedule for `g`."
        self.pred = u + (self.scheduler('g', idx) * (t - u))

    def pre_proc (self, u, t, idx=None): pass
    def post_proc(self, u, t, idx=None): pass
    
    def scheduler(self, name, idx):
        "Gets the scheduled value for parameter `name` at timestep `idx`."
        return self.schedules.get(name)[idx]
    
    def update_sched_kwargs(self, other_kwargs):
        self.schedules.update_sched_kwargs(other_kwargs)
    def set_num_steps(self, num_steps):
        self.schedules.set_num_steps(num_steps)
    def set_guidance_schedule(self):
        self.schedules.set_guidance_schedule()
    
    def __call__(self, *args, **kwargs):
        return self.encode(*args, **kwargs)
    
    
class BaseNormGuidance(GuidanceTfm):
    "Scales the noise prediction by its overall norm."
    name = "BaseNormGuidance"
    def post_proc(self, u, t, idx=None):
        self.pred = self.pred * (torch.linalg.norm(u) / torch.linalg.norm(self.pred))
        
        
class TNormGuidance(GuidanceTfm):
    "Scales the latent mix of `t - u`"
    name = "TNormGuidance"
    def guide(self, u, t, idx=None):
        self.pred = u + (self.scheduler('g', idx) * (t - u)) / torch.linalg.norm(t - u) * torch.linalg.norm(u)
        
        
class FullNormGuidance(TNormGuidance, BaseNormGuidance):
    "Applies both Base and T-Norm on the noise prediction."
    name = "FullNormGuidance"
    pass
