"""The multimodal circuit generation model: *Circuit Diffusion Transformer* (CirDiT)."""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../../../src/models/transformers/cirdit_multimodal.ipynb.

# %% auto #0
__all__ = ['RotaryMultiheadAttention', 'FeedForwardBlock', 'SelfAttnBlock', 'AdaptiveSelfAttnBlock', 'CrossAttnBlock',
           'CoreTransformer', 'PackingTransformer', 'UnpackingTransformer', 'TimeEmbedding', 'CirDiTConfig', 'CirDiT',
           'UnitaryCLIPPartialNoiseCompilationCirDiTConfig', 'UnitaryCLIPPartialNoiseCompilationCirDiT']

# %% ../../../src/models/transformers/cirdit_multimodal.ipynb #27070f09-14a3-4603-8ea4-0b3f3284eae1
from ...imports import *
from ..config_model import *
from ..position_encoding import RotaryPositionalEmbedding, LearnedPositionalEmbedding
from ..layers import PositionalEncoding

# %% ../../../src/models/transformers/cirdit_multimodal.ipynb #7cf35439-5f9a-48bc-ae72-ffd835c9f053
class RotaryMultiheadAttention(nn.Module):
    """
    MultiheadAttention described in the paper: Attention Is All You Need (https://arxiv.org/abs/1706.03762).
    We add a rotary position encoding (RoPE). 

    The attention core is `F.scaled_dot_attention` from pytorch. 
    Could be switched to `https://github.com/Dao-AILab/flash-attention` or `xFormers`.
    """

    def __init__(self,
                 in_dim: int,
                 embed_dim: int, 
                 num_heads: int, 
                 bias: bool = True, 
                 p_rope: float = 1.0, 
                 max_seq_len: int = 4096, 
                 base_rope: float = 10_000,
                 enable_qk_norm: bool = False) -> None:
        
        super().__init__()

        self.num_heads = num_heads
        self.bias      = bias
        self.head_dim  = embed_dim // num_heads 

        self.q_proj   = nn.Linear(in_dim, embed_dim, bias=bias)
        self.k_proj   = nn.Linear(in_dim, embed_dim, bias=bias)
        self.v_proj   = nn.Linear(in_dim, embed_dim, bias=bias)
        
        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)

        self.enable_qk_norm = enable_qk_norm
        if self.enable_qk_norm:
            self.q_norm = nn.RMSNorm(self.head_dim)
            self.k_norm = nn.RMSNorm(self.head_dim)
        
        self.rope = RotaryPositionalEmbedding(head_dim=self.head_dim, p=p_rope, max_seq_len=max_seq_len, base=base_rope)
        
        self._init_weights()
          
    def _init_weights(self) -> None:
        nn.init.xavier_normal_(self.q_proj.weight)
        nn.init.xavier_normal_(self.k_proj.weight)
        nn.init.xavier_normal_(self.v_proj.weight)
        nn.init.xavier_normal_(self.out_proj.weight)

        if self.bias:
            nn.init.zeros_(self.q_proj.bias)
            nn.init.zeros_(self.k_proj.bias)
            nn.init.zeros_(self.v_proj.bias)
            nn.init.zeros_(self.out_proj.bias)

    
    def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, pos_idx: Optional[torch.Tensor] = None) -> torch.Tensor:
        """
        Assumes batch first. When `pos_idx` is provided we use RoPE, else NOT!

        Shapes:
            query     ... [b, n1, c]
            key/value ... [b, n2, c]
        """

        assert key.shape == value.shape
        
        b, n1, _ = query.shape
        _, n2, _ = key.shape

        q = self.q_proj(query)
        k = self.k_proj(key)
        v = self.v_proj(value)

        q = q.view(b, n1, self.num_heads, self.head_dim)
        k = k.view(b, n2, self.num_heads, self.head_dim)
        v = v.view(b, n2, self.num_heads, self.head_dim)

        if self.enable_qk_norm:
            q = self.q_norm(q)
            k = self.k_norm(k)
        
        if exists(pos_idx):
            q = self.rope(q, pos_idx=pos_idx)
            k = self.rope(k, pos_idx=pos_idx)

        # scaled_dot_product_attention takes [b, num_heads, seq, head_dim]
        q = q.permute((0, 2, 1, 3)) 
        k = k.permute((0, 2, 1, 3)) 
        v = v.permute((0, 2, 1, 3)) 
        
        # see https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html
        attn = F.scaled_dot_product_attention(query=q, 
                                              key=k, 
                                              value=v, 
                                              attn_mask=None, 
                                              dropout_p=0.0,
                                              is_causal=False, 
                                              scale=None, 
                                              #enable_gqa=False
                                             )

        # back to [b, seq, num_heads, head_dim]
        attn = attn.permute((0, 2, 1, 3)) 

        # pack heads together
        attn = attn.reshape(b, n1, self.num_heads * self.head_dim)
        attn = self.out_proj(attn)
        return attn

# %% ../../../src/models/transformers/cirdit_multimodal.ipynb #1cbc7b56-dfc7-4922-b7e4-f3ee94ccd311
class FeedForwardBlock(nn.Module):
    """
    A small dense feed-forward network as used in `transformers`. Assumes channel last.
    Inspired by https://arxiv.org/pdf/2401.11605 and added 
    from https://arxiv.org/pdf/2002.05202 a modification to SiGLU structure.
    """

    def __init__(self, 
                 in_dim: int, 
                 hidden_dim: int, 
                 out_dim: Optional[int] = None, 
                 dropout: float = 0.0) -> None:
        super().__init__()
        out_dim = default(out_dim, in_dim)
        
        self.hidden_dim = hidden_dim
        self.proj_in  = nn.Linear(in_dim, 2*hidden_dim) # factor two for GLU part split
        self.proj_out = nn.Linear(hidden_dim, out_dim) 
        self.act   = nn.SiLU()
        self.drop  = nn.Dropout(dropout)

        self._init_weights()
   
    def _init_weights(self) -> None:
        nn.init.zeros_(self.proj_out.bias)
        # nn.init.xavier_normal_(self.proj_out.weight)
    
    def siglu(self, x: torch.Tensor) -> torch.Tensor:
        x = self.proj_in(x) 
        return x[..., :self.hidden_dim] * self.act(x[..., self.hidden_dim:])

    #@torch.compile
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.siglu(x)
        x = self.drop(x)
        x = self.proj_out(x)
        return x

# %% ../../../src/models/transformers/cirdit_multimodal.ipynb #3b028c09-74ad-4db8-bf7f-c1f1a9b259d5
class SelfAttnBlock(nn.Module):
    """A self-attention block which includes the time condition `t_emb`, see https://arxiv.org/pdf/2312.02139."""
    
    def __init__(self, ch: int, t_emb_size: int, num_heads: int, dropout: float = 0.0, p_rope: float = 1.0, base_rope: float = 10_000) -> None:
        super().__init__()

        self.self_att = RotaryMultiheadAttention(in_dim=ch+t_emb_size, embed_dim=ch, num_heads=num_heads, p_rope=p_rope, base_rope=base_rope)
 
        self.ff        = FeedForwardBlock(in_dim=ch, hidden_dim=2*ch)   
        self.norm_self = nn.RMSNorm(ch)
        self.norm_ff   = nn.RMSNorm(ch)
        self.drop      = nn.Dropout(dropout)

        self._init_weights()
          
    def _init_weights(self) -> None:

        # note a bonus of res-pos-norm is that we can init as identity!
        nn.init.zeros_(self.norm_self.weight)  
        nn.init.zeros_(self.norm_ff.weight)
    
    def forward(self, x: torch.Tensor, t_emb: torch.Tensor, pos_idx: torch.Tensor) -> torch.Tensor:
        """
        Assumes batch first.
        
        Shapes:
            x       ... [b, n, ch]       
            t_emb   ... [b,  1, t_emb_size]
            pos_idx ... [b, n] or [n]
        """
        
        t_emb_self  = t_emb.expand(x.shape[0], x.shape[1], -1)

        # Self-attention part
        self_out = torch.cat([x, t_emb_self], dim=-1)    # concat time tokens
        self_out = self.self_att(query=self_out, key=self_out, value=self_out, pos_idx=pos_idx)
        self_out = self.norm_self(self_out)  
        self_out = self.drop(self_out) + x      

        # Feed-Forward part
        feed_out = self.ff(self_out)
        feed_out = self.norm_ff(feed_out)  
        feed_out = self.drop(feed_out) + self_out                      
        return feed_out

# %% ../../../src/models/transformers/cirdit_multimodal.ipynb #7d503e52-2c26-4efc-b501-1fa01f84a1f0
class AdaptiveSelfAttnBlock(nn.Module):
    """A self-attention block which includes the time condition `t_emb`, see https://arxiv.org/pdf/2312.02139."""
    
    def __init__(self, 
                 ch: int, 
                 mod_ch: int,
                 t_emb_size: int, 
                 num_heads: int, 
                 dropout: float = 0.0,
                 p_rope: float = 1.0, 
                 base_rope: float = 10_000) -> None:
        super().__init__()

        self.self_att = RotaryMultiheadAttention(in_dim=ch+t_emb_size, embed_dim=ch, num_heads=num_heads, p_rope=p_rope, base_rope=base_rope)
 
        self.ff        = FeedForwardBlock(in_dim=ch, hidden_dim=2*ch)   
        self.norm_self = nn.RMSNorm(ch)
        self.norm_ff   = nn.RMSNorm(ch)
        self.drop      = nn.Dropout(dropout)
        
        self.adaRMS_modulation = nn.Linear(mod_ch, 6*ch)

        self._init_weights()
          
    def _init_weights(self) -> None:
        nn.init.zeros_(self.adaRMS_modulation.bias) 
 
    def forward(self, x: torch.Tensor, mod: torch.Tensor, t_emb: torch.Tensor, pos_idx: torch.Tensor) -> torch.Tensor:
        """
        Assumes batch first.
        
        Shapes:
            x       ... [b, n, ch]  
            mod     ... [b, n, mod_ch]       
            t_emb   ... [b,  1, t_emb_size]
            pos_idx ... [b, n] or [n]
        """

        scale_att, shift_att, gate_attn, scale_mlp, shift_mlp, gate_mlp = self.adaRMS_modulation(mod).chunk(6, dim=-1)
        
        t_emb_self = t_emb.expand(x.shape[0], x.shape[1], -1)

        # Self-attention part
        self_out = x * (1.0 + scale_att) + shift_att
        self_out = torch.cat([self_out, t_emb_self], dim=-1)    # concat time tokens
        self_out = self.self_att(query=self_out, key=self_out, value=self_out, pos_idx=pos_idx)
        self_out = self.norm_self(self_out) * gate_attn.tanh()
        self_out = self.drop(self_out) + x      

        # Feed-Forward part
        feed_out = self_out * (1.0 + scale_mlp) + shift_mlp
        feed_out = self.ff(feed_out)
        feed_out = self.norm_ff(feed_out) * gate_mlp.tanh()
        feed_out = self.drop(feed_out) + self_out                      
        return feed_out

# %% ../../../src/models/transformers/cirdit_multimodal.ipynb #6fea48d2-7167-4cfb-8c20-6cc39c4c10f5
class CrossAttnBlock(nn.Module):
    """A cross-attention block which includes the time condition `t_emb`, see https://arxiv.org/pdf/2312.02139"""
    
    def __init__(self, ch: int, t_emb_size: int, num_heads: int, dropout: float = 0.0, p_rope: float = 1.0, base_rope: float = 10_000) -> None:
        super().__init__()
        
        self.self_att  = RotaryMultiheadAttention(in_dim=ch+t_emb_size, embed_dim=ch, num_heads=num_heads, p_rope=p_rope, base_rope=base_rope)
        self.multi_att = RotaryMultiheadAttention(in_dim=ch+t_emb_size, embed_dim=ch, num_heads=num_heads, p_rope=p_rope, base_rope=base_rope)

        self.ff         = FeedForwardBlock(in_dim=ch, hidden_dim=2*ch)   
        self.norm_self  = nn.RMSNorm(ch)
        self.norm_multi = nn.RMSNorm(ch)
        self.norm_ff    = nn.RMSNorm(ch)
        self.drop       = nn.Dropout(dropout)

        self._init_weights()
    
    def _init_weights(self) -> None:
        nn.init.zeros_(self.norm_self.weight)  
        nn.init.zeros_(self.norm_multi.weight)  
        nn.init.zeros_(self.norm_ff.weight)
    
    def forward(self, x: torch.Tensor, c_emb: torch.Tensor, t_emb: torch.Tensor, pos_idx: torch.Tensor) -> torch.Tensor:
        """
        Assumes batch first.
        
        Shapes:
            x       ... [b, n1, ch]       
            c_emb   ... [b, n2, ch]
            t_emb   ... [b,  1, t_emb_size]
            pos_idx ... [b, n1] or [n1]
        """
        
        t_emb_self  = t_emb.expand(    x.shape[0], x.shape[1]                 , -1)
        t_emb_multi = t_emb.expand(c_emb.shape[0], x.shape[1] + c_emb.shape[1], -1)
        
        # Self-attention part
        self_out = torch.cat([x, t_emb_self], dim=-1)    # concat time tokens
        self_out = self.self_att(query=self_out, key=self_out, value=self_out, pos_idx=pos_idx)
        self_out = self.norm_self(self_out)
        self_out = self.drop(self_out) + x      

        # Multimodial-attention part
        multi_out = torch.cat([self_out, c_emb], dim=1)          # concat latents with condition ... [b, n1+n2, ch]

        multi_out = torch.cat([multi_out, t_emb_multi], dim=-1)  # concat time tokens           
        multi_out = self.multi_att(query=multi_out, key=multi_out, value=multi_out, pos_idx=None)

        multi_out, multi_out_gate = multi_out[:, :x.shape[1]], multi_out[:, x.shape[1]:]
        multi_out_gate = multi_out_gate.mean(dim=1, keepdim=True) # ... [b, 1, ch]
        
        multi_out = self.norm_multi(multi_out) * multi_out_gate.tanh()
        multi_out = self.drop(multi_out) + self_out         

        # Feed-Forward part
        feed_out = self.ff(multi_out)
        feed_out = self.norm_ff(feed_out)              
        feed_out = self.drop(feed_out) + multi_out                       
        return feed_out

# %% ../../../src/models/transformers/cirdit_multimodal.ipynb #18b30977-f8d2-4f5d-91f7-f96d4e4edf11
class CoreTransformer(nn.Module):
    """
    The main transformer of the CirDiT model, intakes time (attn-concat) and condition encodings (cross-attn). 
    Applies a RoPE for time dimension.
    """

    def __init__(self,
                 ch: int, 
                 c_emb_size: int,
                 t_emb_size: int,
                 depth: int,
                 num_heads: int, 
                 dropout: float = 0.0, 
                 p_rope: float = 1.0, 
                 base_rope: float = 10_000) -> None:
        super().__init__()

        self.norm = nn.RMSNorm(ch)
        
        self.c_proj = nn.Linear(c_emb_size, ch)          
        self.blocks = nn.ModuleList([
                            CrossAttnBlock(ch=ch, 
                                           t_emb_size=t_emb_size, 
                                           num_heads=num_heads, 
                                           dropout=dropout, 
                                           p_rope=p_rope, 
                                           base_rope=base_rope)
                            for d in range(depth)
                        ])

    def forward(self, x: torch.Tensor, c_emb: torch.Tensor, t_emb: torch.Tensor) -> torch.Tensor:
        """
        Shapes:
            x     ... [b,   t,         ch]
            c_emb ... [b, seq, c_emb_size]
            t_emb ... [b,   1, t_emb_size]
        """
        
        c_emb   = self.c_proj(c_emb)
        pos_idx = torch.arange(x.shape[1], device=x.device, dtype=torch.int32) 

        x = self.norm(x)
        
        for block in self.blocks:
            x = block(x=x, c_emb=c_emb, t_emb=t_emb, pos_idx=pos_idx)

        return x

# %% ../../../src/models/transformers/cirdit_multimodal.ipynb #96923c30-ede7-4d24-a4a8-d80423bece98
class PackingTransformer(ConfigModel):
    """
    The first stage packing/unpacking transformers of the CirDiT model, intakes time (attn-concat). 
    Applies a RoPE for time dimension only, not on spatial dimension.
    """
    
    def __init__(self,
                 ch: int, 
                 t_emb_size: int,
                 depth: int,
                 num_heads: int, 
                 dropout: float = 0.0, 
                 p_rope: float = 1.0, 
                 base_rope: float = 10_000) -> None:
        super().__init__()

        self.norm   = nn.RMSNorm(ch)
        self.blocks = nn.ModuleList([
                    SelfAttnBlock(ch=ch, 
                                  t_emb_size=t_emb_size, 
                                  num_heads=num_heads, 
                                  dropout=dropout, 
                                  p_rope=p_rope, 
                                  base_rope=base_rope)
                    for d in range(depth)
                ])
            
    def forward(self, x: torch.Tensor, t_emb: torch.Tensor, return_penultimate: bool = False) -> torch.Tensor:
        """
        Shapes:
            x     ... [b, s, t, ch]
            t_emb ... [b, 1, t_emb_size]
        """

        b, s, t, ch = x.shape

        # create pos_idx such that they only depend on the time position
        pos_idx = torch.arange(t, device=x.device, dtype=torch.int32).expand(b, s, -1)
        pos_idx = pos_idx.reshape(b, -1)

        # flatten spatial and time into seq
        x = x.reshape(b, s*t, ch)
        x = self.norm(x)

        if return_penultimate:
            for block in self.blocks[:-1]:
                x = block(x=x, t_emb=t_emb, pos_idx=pos_idx)

            penultimate = x
            x = self.blocks[-1](x=x, t_emb=t_emb, pos_idx=pos_idx)   
 
        else:
            for block in self.blocks:
                x = block(x=x, t_emb=t_emb, pos_idx=pos_idx)

        # undo flatten
        x = x.reshape(b, s, t, ch)
        
        if return_penultimate:
            penultimate = penultimate.reshape(b, s, t, ch)
            return x, penultimate
        
        return x

# %% ../../../src/models/transformers/cirdit_multimodal.ipynb #4a16f88c-36fd-4888-a631-4e61f91af3b1
class UnpackingTransformer(ConfigModel):
    """
    The first stage packing/unpacking transformers of the CirDiT model, intakes time (attn-concat). 
    Applies a RoPE for time dimension only, not on spatial dimension.
    """
    
    def __init__(self,
                 ch: int, 
                 mod_ch: int,
                 t_emb_size: int,
                 depth: int,
                 num_heads: int, 
                 dropout: float = 0.0, 
                 p_rope: float = 1.0, 
                 base_rope: float = 10_000) -> None:
        super().__init__()

        self.norm   = nn.RMSNorm(ch)
        self.blocks = nn.ModuleList([
                    AdaptiveSelfAttnBlock(ch=ch, 
                                          mod_ch=mod_ch,
                                          t_emb_size=t_emb_size, 
                                          num_heads=num_heads, 
                                          dropout=dropout, 
                                          p_rope=p_rope, 
                                          base_rope=base_rope)
                        for d in range(depth)
                    ])
            
    def forward(self, x: torch.Tensor, mod: torch.Tensor, t_emb: torch.Tensor) -> torch.Tensor:
        """
        Shapes:
            x     ... [b, s, t, ch]
            t_emb ... [b, 1, t_emb_size]
        """

        b, s, t, ch = x.shape
        *_,  mod_ch = mod.shape
        
        # create pos_idx such that they only depend on the time position
        pos_idx = torch.arange(t, device=x.device, dtype=torch.int32).expand(b, s, -1)
        pos_idx = pos_idx.reshape(b, -1)

        # flatten spatial and time into seq
        x   =   x.reshape(b, s*t, ch)
        mod = mod.reshape(b, s*t, mod_ch).contiguous()

        x = self.norm(x)

        for block in self.blocks:
            x = block(x=x, mod=mod, t_emb=t_emb, pos_idx=pos_idx)

        # undo flatten
        x = x.reshape(b, s, t, ch)
        return x

# %% ../../../src/models/transformers/cirdit_multimodal.ipynb #00e9e5f9-45f5-475d-8daf-537c851b9cb6
class TimeEmbedding(PositionalEncoding):
    """A time embedding layer."""
    
    def __init__(self, d_model: int, 
                 dropout: float = 0.0, 
                 max_len: int = 5000, 
                 freq_factor: float = 10_000.0) -> None:
        super().__init__(d_model=d_model, dropout=dropout, max_len=max_len, freq_factor=freq_factor)          
    
        self.ff = FeedForwardBlock(in_dim=d_model, hidden_dim=2*d_model)  
       
    def forward(self, t: torch.Tensor) -> torch.Tensor:       
        x = self.pe[t]       
        x = self.ff(x)               
        return self.dropout(x)

# %% ../../../src/models/transformers/cirdit_multimodal.ipynb #f0a55971-2132-4b61-b673-e39bed65f679
@dataclass
class CirDiTConfig:  
    clr_dim: int
    ch_packing: int
    ch_core: int
    c_emb_size: int
    t_emb_size: int        
    depth_packing: int
    depth_core: int                 
    num_heads_packing: int
    num_heads_core: int              
    dropout: float 
    p_rope: float 
    base_rope: float

# %% ../../../src/models/transformers/cirdit_multimodal.ipynb #c97230bd-8f12-423f-a4ca-baab7c9ca7c9
class CirDiT(ConfigModel):
    """
    The proposed Circuit Diffusion Transformer (CirDiT).
    """

    def __init__(self,
                 clr_dim: int,
                 ch_packing: int, 
                 ch_core: int,
                 c_emb_size: int,
                 t_emb_size: int, 
                 depth_packing: int,
                 depth_core: int,                 
                 num_heads_packing: int,
                 num_heads_core: int,              
                 dropout: float = 0.0, 
                 p_rope: float = 1.0, 
                 base_rope: float = 10_000) -> None:
        super().__init__()
        
        self.ch_packing = ch_packing
        self.ch_core    = ch_core
        
        self.params_config = CirDiTConfig(clr_dim=clr_dim,
                                          ch_packing=ch_packing, 
                                          ch_core=ch_core,
                                          c_emb_size=c_emb_size,
                                          t_emb_size=t_emb_size,   
                                          depth_packing=depth_packing,
                                          depth_core=depth_core,                 
                                          num_heads_packing=num_heads_packing,
                                          num_heads_core=num_heads_core,              
                                          dropout=dropout, 
                                          p_rope=p_rope, 
                                          base_rope=base_rope)
        
        self.packing = PackingTransformer(ch=ch_packing, 
                                          t_emb_size=t_emb_size, 
                                          depth=depth_packing, 
                                          num_heads=num_heads_packing, 
                                          dropout=dropout, 
                                          p_rope=p_rope, 
                                          base_rope=base_rope)
        
        self.unpacking = UnpackingTransformer(ch=ch_packing, 
                                              mod_ch=ch_core,
                                              t_emb_size=t_emb_size, 
                                              depth=depth_packing, 
                                              num_heads=num_heads_packing, 
                                              dropout=dropout, 
                                              p_rope=p_rope, 
                                              base_rope=base_rope)

        self.core = CoreTransformer(ch=ch_core, 
                                    c_emb_size=c_emb_size, 
                                    t_emb_size=t_emb_size, 
                                    depth=depth_core, 
                                    num_heads=num_heads_core, 
                                    dropout=dropout, 
                                    p_rope=p_rope, 
                                    base_rope=base_rope)

        self.proj_in   = nn.Linear(clr_dim, ch_packing)
        self.proj_out  = nn.Linear(ch_packing, clr_dim)
        self.core_proj = nn.Linear(ch_packing, ch_core)

        self.t_emb         = TimeEmbedding(d_model=t_emb_size, max_len=5000) #here max number of timetseps
        self.qubit_pos_enc = LearnedPositionalEmbedding(dim=ch_packing, max_seq_len=64) #here max number of qubits

        self._init_weights()
   
    def _init_weights(self) -> None:
        nn.init.orthogonal_(self.core_proj.weight)
        nn.init.zeros_(self.core_proj.bias)
        nn.init.zeros_(self.proj_out.bias)
    
    def main_pass(self, x: torch.Tensor, t_emb: torch.Tensor, c_emb: torch.Tensor) -> torch.Tensor:
        b, s, t, _ = x.shape

        x = self.proj_in(x)
        x = self.qubit_pos_enc(x)

        # Pack spatial into tokens
        x_main, x = self.packing(x=x, t_emb=t_emb, return_penultimate=True)

        # Downsample, reduce spatial, ... [b, t, ch_core]
        x_main = x_main.mean(dim=1)   
        x_main = self.core_proj(x_main)  

        # Core transformer
        x_main = self.core(x=x_main, c_emb=c_emb, t_emb=t_emb) - x_main  # subtraction such that if core=ident at init we cancel the signal
        x_main = x_main.unsqueeze(1).expand(b, s, t, self.ch_core) 

        # Unpack tokens into spatial
        x = self.unpacking(x=x, mod=x_main, t_emb=t_emb)
        x = self.proj_out(x)  

        return x
    
    def forward(self, x: torch.Tensor, t: torch.Tensor, c_emb: torch.Tensor, micro_cond: Optional[torch.Tensor] = None) -> torch.Tensor:
        """
        Assumes a `channel_last` embedding of circuits.
        
        Shapes:
            x     ... [b, s, t, ch]  
            t     ... [b]
            c_emb ... [b, seq, c_emb_size]
            micro_cond ... [b]
        """

        t_emb = self.t_emb(t)  #.detach()
        t_emb = t_emb.unsqueeze(1) # to [b, 1, ch]
            
        x = self.main_pass(x, t_emb, c_emb)
        return x

# %% ../../../src/models/transformers/cirdit_multimodal.ipynb #ba7c3c61-2044-4fb2-a00a-eafeaf647664
@dataclass
class UnitaryCLIPPartialNoiseCompilationCirDiTConfig(CirDiTConfig):  
    unitary_encoder_config: dict

# %% ../../../src/models/transformers/cirdit_multimodal.ipynb #a3bda09d-ca1a-4406-b630-08b346e7a27a
class UnitaryCLIPPartialNoiseCompilationCirDiT(CirDiT):
    """Extends `CirDiT` to the multimodal unitary compilation model."""
    
    def __init__(self,                  
                 clr_dim: int,
                 ch_packing: int, 
                 ch_core: int,
                 c_emb_size: int,
                 t_emb_size: int,            
                 depth_packing: int,
                 depth_core: int,                 
                 num_heads_packing: int,
                 num_heads_core: int,              
                 dropout: float = 0.0, 
                 p_rope: float = 1.0, 
                 base_rope: float = 10_000,
                 unitary_encoder_config: Optional[dict] = None, 
                 unitary_encoder: Optional[nn.Module] = None) -> None:
        
        super().__init__(clr_dim=clr_dim, 
                         ch_packing=ch_packing,
                         ch_core=ch_core,
                         c_emb_size=c_emb_size,
                         t_emb_size=t_emb_size,
                         depth_packing=depth_packing,
                         depth_core=depth_core,
                         num_heads_packing=num_heads_packing,
                         num_heads_core=num_heads_core,
                         dropout=dropout,
                         p_rope=p_rope,
                         base_rope=base_rope)

        if exists(unitary_encoder_config): #load a trained encoder
            self.unitary_encoder = ConfigModel.from_config(unitary_encoder_config, device=None, silent=True)
            
        elif exists(unitary_encoder):
            self.unitary_encoder   = unitary_encoder
            unitary_encoder_config = self.unitary_encoder.get_config()
            
            unitary_encoder_config = {"target": unitary_encoder_config["target"],
                                      "params": unitary_encoder_config["params"]}
        
        else: 
            raise RuntimeError("Provide either `unitary_encoder_config` to load a pretrained encoder or a `unitary_encoder` model directly!`")

        self.params_config = UnitaryCLIPPartialNoiseCompilationCirDiTConfig(
                                            clr_dim=clr_dim, 
                                            ch_packing=ch_packing,
                                            ch_core=ch_core,
                                            c_emb_size=c_emb_size,
                                            t_emb_size=t_emb_size,
                                            depth_packing=depth_packing,
                                            depth_core=depth_core,
                                            num_heads_packing=num_heads_packing,
                                            num_heads_core=num_heads_core,
                                            dropout=dropout,
                                            p_rope=p_rope,
                                            base_rope=base_rope,
                                            unitary_encoder_config=unitary_encoder_config
                                        )

        #--------

        self.empty_cond = nn.Parameter(torch.randn((1, 1, c_emb_size)))

        self.t_emb  = TimeEmbedding(d_model=t_emb_size, max_len=5000) #here max number of timetseps
        self.t_emb2 = TimeEmbedding(d_model=t_emb_size, max_len=5000) #here max number of timetseps
    
    def forward(self, 
                x: torch.Tensor, 
                t_h: torch.Tensor, 
                t_w: torch.Tensor, 
                c_emb: torch.Tensor, 
                U: torch.Tensor, 
                rnd: Optional[torch.Tensor] = None) -> torch.Tensor:
        """
        Assumes a channel_last embedding of circuits.
        
        Shapes:
            x     ... [b, s, t, ch]  
            t_h   ... [b]
            t_w   ... [b]
            c_emb ... [b, seq, c_emb_size]
            U     ... [b, 2, N, N]
            rnd   ... [b]
        """
        
        t_emb = self.t_emb(t_h) + self.t_emb2(t_w)
        t_emb = t_emb.unsqueeze(1) # to [b, 1, ch]

        #------
        
        u_emb = self.unitary_encoder(y_emb=c_emb, U=U, penultimate=True).detach() # [batch, seq1+seq2, ch]     

        if not_exists(rnd):
            # one means we dont drop, so U is not all zero
            rnd = 1-torch.isclose(U, torch.zeros_like(U)).all(dim=(1, 2, 3)).type(torch.int64)    
        rnd = rnd.view(-1, 1, 1)
        
        # Note: we ignore text drop and unitary drop, we replace all with a learned uncond token here
        u_emb = u_emb * rnd + (1-rnd) * self.empty_cond.expand(u_emb.shape)

        #------
        
        x = self.main_pass(x, t_emb, u_emb) 
        return x
