import torch
import torch.nn as nn
import gin

#! GATE-FUSIONS MODIFICATIONS ---------------------------------
#! in_dim = 1024

from enum import Enum
import math

class ResidualGateType(Enum):
    simple_residual_one_gate = 0,
    cross_gate = 1,
    gate_sum = 2,
    alpha_res_cat = 3,
    alpha_res_sum = 4

# TODO: add GatedFusionresidual class to gin

def init_bias_for_sigmoid(linear: nn.Linear, value: float) -> None:
    nn.init.zeros_(linear.weight)
    nn.init.constant_(linear.bias, math.log(value / (1.0 - value)))

def init_residual_projs(linear: nn.Linear, scale: float) -> None:
    nn.init.xavier_uniform_(linear.weight, gain=scale)
    nn.init.zeros_(linear.bias)

RESIDUAL_GATES = {
    ResidualGateType.alpha_res_sum,
    ResidualGateType.alpha_res_cat
}

@gin.configurable
class GatedFusionResidual(nn.Module):
    """Learnable gated fusion of image and text embeddings.

    V1 - Simple residual gating with 1 common gate: 
    V2 - Cross residual gating with 2 cross-gates
    V3 - Gate + Simple Sum of feats x & y
    V4 - Alpha-weighted residual sum (per sample)
    V5 - Alpha-weighted residual concat (per sample)
    """

    def __init__(self, gate_type: ResidualGateType, 
                 init_gate: float = 0.7, in_dim = 1024, 
                 init_res_weight: float = 0.1, residual_proj_scale: float = 0.1,  
                 baseline_mode: bool = False, 
                 ) -> None:
        super().__init__()
        # alpha is in logit space: sigmoid(alpha) = init_gate
        init_alpha = torch.log(torch.tensor(init_gate / (1.0 - init_gate)))
        self.alpha = nn.Parameter(init_alpha)
        
        # alphas for separated cases
        if gate_type == ResidualGateType.cross_gate:
            init_alpha_img_cross_gate = torch.log(torch.tensor(init_gate / (1.0 - init_gate)))
            init_alpha_text_cross_gate = torch.log(torch.tensor(init_gate / (1.0 - init_gate)))
            self.alpha_img = nn.Parameter(init_alpha_img_cross_gate)
            self.alpha_text = nn.Parameter(init_alpha_text_cross_gate)
        
        # weight for sum and cat residual 
        if gate_type in RESIDUAL_GATES:
            self.final_cat_residual_proj = nn.Linear(in_dim * 2, in_dim)
            self.weight_net_for_sum = nn.Linear(in_dim, 1)
            self.weight_net_for_cat = nn.Linear(in_dim * 2, 1)
            init_bias_for_sigmoid(self.weight_net_for_sum, value=init_res_weight)
            init_bias_for_sigmoid(self.weight_net_for_cat, value=init_res_weight)
            init_residual_projs(self.final_cat_residual_proj, scale=residual_proj_scale)

        self.gate_type = gate_type
        self.baseline_mode = baseline_mode

    def FuseSRGF(self,
        img_feat: torch.Tensor,
        text_feat: torch.Tensor | None,
    ) -> torch.Tensor:
        gate = torch.sigmoid(self.alpha)
        img_res = img_feat * gate + img_feat
        text_res = text_feat * (1 - gate) + text_feat
        fused_vec = img_res + text_res
        return fused_vec
    
    def FuseRCGF(self,
        img_feat: torch.Tensor,
        text_feat: torch.Tensor | None,
    ) -> torch.Tensor:
        gate_img = torch.sigmoid(self.alpha_img)
        gate_text = torch.sigmoid(self.alpha_text)
        z_img = img_feat + gate_text * img_feat
        z_text = text_feat + gate_img * text_feat
        fused_vec = z_img + z_text
        return fused_vec
    
    def FuseGSUM(self,
        img_feat: torch.Tensor,
        text_feat: torch.Tensor | None,
    ) -> torch.Tensor:
        gate = torch.sigmoid(self.alpha)
        fuzed_vec = img_feat + text_feat + gate * img_feat + (1.0 - gate) * text_feat
        return fuzed_vec
    
    def FuseARGFSum(
        self,
        img_feat: torch.Tensor,
        text_feat: torch.Tensor | None,
    ) -> torch.Tensor:
        gate = torch.sigmoid(self.alpha)
        residual = img_feat + text_feat
        res_weight = torch.sigmoid(self.weight_net_for_sum(residual))
        fuzed_vec = gate * img_feat + (1.0 - gate) * text_feat + res_weight * residual
        return fuzed_vec

    def FuseARGFCat(
        self,
        img_feat: torch.Tensor,
        text_feat: torch.Tensor | None,
    ) -> torch.Tensor:
        gate = torch.sigmoid(self.alpha)
        cat_vec = torch.cat([img_feat, text_feat], dim=-1)
        residual = self.final_cat_residual_proj(cat_vec)
        res_weight = torch.sigmoid(self.weight_net_for_cat(cat_vec))
        fuzed_vec = gate * img_feat + (1.0 - gate) * text_feat + res_weight * residual
        return fuzed_vec

    def forward(
        self,
        img_feat: torch.Tensor,
        text_feat: torch.Tensor | None,
    ) -> torch.Tensor:
        if text_feat is None or self.baseline_mode:
            return img_feat

        if self.gate_type == ResidualGateType.simple_residual_one_gate:
            fused_vec = self.FuseSRGF(img_feat=img_feat, text_feat=text_feat)
        if self.gate_type == ResidualGateType.cross_gate:
            fused_vec = self.FuseRCGF(img_feat=img_feat, text_feat=text_feat)
        if self.gate_type == ResidualGateType.gate_sum:
            fused_vec = self.FuseGSUM(img_feat=img_feat, text_feat=text_feat)
        if self.gate_type == ResidualGateType.alpha_res_sum:
            fused_vec = self.FuseARGFSum(img_feat=img_feat, text_feat=text_feat)
        if self.gate_type == ResidualGateType.alpha_res_cat:
            fused_vec = self.FuseARGFCat(img_feat=img_feat, text_feat=text_feat)

        return fused_vec
    
    @property
    def gate_value(self) -> float:
        """Current gate value (image weight). 1.0 = text ignored."""
        if self.baseline_mode:
            return 1.0
        return torch.sigmoid(self.alpha).item()