from __future__ import annotations
import torch
import torch.nn as nn
import torch.nn.functional as F


class GradReverseFn(torch.autograd.Function):
    @staticmethod
    def forward(ctx,x,lambd): ctx.lambd=float(lambd); return x.view_as(x)
    @staticmethod
    def backward(ctx,g): return -ctx.lambd*g,None

def grad_reverse(x,lambd): return GradReverseFn.apply(x,lambd)


class Tower(nn.Module):
    def __init__(self,din:int,dout:int,dropout:float=.12):
        super().__init__(); self.net=nn.Sequential(nn.LayerNorm(din),nn.Linear(din,dout),nn.GELU(),nn.Dropout(dropout),nn.Linear(dout,dout),nn.GELU())
    def forward(self,x): return self.net(x)


class RavenFusion(nn.Module):

    BRANCHES=("semantic","camera","spectral","patch","specialist")
    def __init__(self,dims:dict,num_ai_domains:int=0,num_real_domains:int=0,latent_dim:int=64):
        super().__init__()
        outdims={"semantic":96,"camera":48,"spectral":48,"patch":48,"specialist":24}
        self.towers=nn.ModuleDict({k:Tower(int(dims[k]),outdims[k]) for k in self.BRANCHES})
        self.branch_heads=nn.ModuleDict({k:nn.Linear(outdims[k],1) for k in self.BRANCHES})
        total=sum(outdims.values()); self.pre_vib=nn.Sequential(nn.LayerNorm(total),nn.Linear(total,128),nn.GELU(),nn.Dropout(.15))
        self.mu=nn.Linear(128,latent_dim); self.logvar=nn.Linear(128,latent_dim)
        qdim=int(dims["quality"]); self.router=nn.Sequential(nn.LayerNorm(qdim),nn.Linear(qdim,48),nn.GELU(),nn.Linear(48,len(self.BRANCHES)))
        self.anomaly_head=nn.Sequential(nn.LayerNorm(5),nn.Linear(5,24),nn.GELU(),nn.Linear(24,1))
        self.global_head=nn.Sequential(nn.LayerNorm(latent_dim),nn.Linear(latent_dim,48),nn.GELU(),nn.Dropout(.12),nn.Linear(48,1))
        self.mixer=nn.Linear(3,1)
        with torch.no_grad():
            self.mixer.weight.copy_(torch.tensor([[1.0,1.0,.8]])); self.mixer.bias.zero_()
        self.ai_domain_head=nn.Sequential(nn.Linear(latent_dim,48),nn.GELU(),nn.Linear(48,num_ai_domains)) if num_ai_domains>1 else None
        self.real_domain_head=nn.Sequential(nn.Linear(latent_dim,48),nn.GELU(),nn.Linear(48,num_real_domains)) if num_real_domains>1 else None

    def _latent(self,h,sample:bool):
        p=self.pre_vib(h); mu=self.mu(p); lv=torch.clamp(self.logvar(p),-6,4)
        if sample and self.training: z=mu+torch.randn_like(mu)*torch.exp(.5*lv)
        else: z=mu
        kl=-.5*torch.mean(1+lv-mu.pow(2)-lv.exp())
        return z,mu,lv,kl

    def forward(self,features:dict,quality:torch.Tensor,anchors:torch.Tensor,luma:torch.Tensor,domain_lambda:float=0.0,sample_vib:bool=True):
        parts={k:self.towers[k](features[k]) for k in self.BRANCHES}; h=torch.cat([parts[k] for k in self.BRANCHES],1); z,mu,lv,kl=self._latent(h,sample_vib)
        branch_logits=torch.cat([self.branch_heads[k](parts[k]) for k in self.BRANCHES],1)


        route_quality=quality.clone()
        if route_quality.shape[1]>=12: route_quality[:,10:12]=0.0
        route_logits=self.router(route_quality)

        darkness=torch.clamp((.25-luma)/.25,0,1).unsqueeze(1)
        prior=torch.tensor([-.60,.50,.38,.30,.12],dtype=route_logits.dtype,device=route_logits.device).unsqueeze(0)
        gates=F.softmax(route_logits+darkness*prior,dim=1); gated=(gates*branch_logits).sum(1,keepdim=True)
        global_logit=self.global_head(z); anomaly_logit=self.anomaly_head(anchors)

        components=torch.cat([global_logit,gated,anomaly_logit],1); logit=self.mixer(components).squeeze(1)
        ai_domain=self.ai_domain_head(grad_reverse(z,domain_lambda)) if self.ai_domain_head is not None and domain_lambda>0 else None
        real_domain=self.real_domain_head(grad_reverse(z,domain_lambda)) if self.real_domain_head is not None and domain_lambda>0 else None
        return {"logit":logit,"branch_logits":branch_logits,"gates":gates,"components":components,"latent":mu,"kl":kl,"ai_domain_logits":ai_domain,"real_domain_logits":real_domain}