# ══════════════════════════════════════════════════════
# RCMTUNetV4-VLM — Pipeline complet v4.0
# Architecture : Segmentation + RAG + VLM + Post-validation
# ══════════════════════════════════════════════════════

import os, re, json, warnings
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange
from PIL import Image
import cv2
from scipy import ndimage
from skimage.measure import label, regionprops
from scipy.ndimage import binary_fill_holes


# ─────────────────────────────────────────────────────
# ARCHITECTURE SEGMENTATION — RCMTUNetV4
# ─────────────────────────────────────────────────────

class ResidualBlock(nn.Module):
    def __init__(self, in_c, out_c):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv3d(in_c, out_c, kernel_size=3, padding=1),
            nn.InstanceNorm3d(out_c),
            nn.LeakyReLU(0.1, inplace=True),
            nn.Conv3d(out_c, out_c, kernel_size=3, padding=1),
            nn.InstanceNorm3d(out_c),
        )
        self.shortcut = (nn.Conv3d(in_c, out_c, kernel_size=1)
                         if in_c != out_c else nn.Identity())
    def forward(self, x):
        return F.leaky_relu(self.conv(x) + self.shortcut(x), 0.1)


class AttentionGate(nn.Module):
    def __init__(self, F_g, F_l, F_int):
        super().__init__()
        self.W_g = nn.Sequential(nn.Conv3d(F_g, F_int, kernel_size=1), nn.InstanceNorm3d(F_int))
        self.W_x = nn.Sequential(nn.Conv3d(F_l, F_int, kernel_size=1), nn.InstanceNorm3d(F_int))
        self.psi = nn.Sequential(nn.Conv3d(F_int, 1, kernel_size=1), nn.Sigmoid())
    def forward(self, g, x):
        g1 = self.W_g(g)
        x1 = self.W_x(x)
        if g1.shape[2:] != x1.shape[2:]:
            g1 = F.interpolate(g1, size=x1.shape[2:], mode="trilinear", align_corners=False)
        return x * self.psi(F.leaky_relu(g1 + x1, 0.1))


class ModalityFusionBlock(nn.Module):
    def __init__(self, in_channels=4, fusion_dim=24, num_heads=4):
        super().__init__()
        self.M = in_channels
        self.modal_proj = nn.ModuleList([nn.Linear(1, fusion_dim) for _ in range(in_channels)])
        self.norm_q  = nn.LayerNorm(fusion_dim)
        self.norm_kv = nn.LayerNorm(fusion_dim)
        self.cross_attn = nn.MultiheadAttention(fusion_dim, num_heads, batch_first=True)
        self.out_proj = nn.ModuleList([nn.Linear(fusion_dim, 1) for _ in range(in_channels)])
    def forward(self, x):
        B   = x.shape[0]
        gap = x.mean(dim=[2, 3, 4])
        tokens = torch.cat([self.modal_proj[m](gap[:, m:m+1]).unsqueeze(1) for m in range(self.M)], dim=1)
        Q  = self.norm_q(tokens[:, 2:3, :])
        KV = self.norm_kv(tokens)
        attn_out, _ = self.cross_attn(Q, KV, KV)
        tokens = tokens.clone()
        tokens[:, 2:3, :] += attn_out
        bias = torch.cat([self.out_proj[m](tokens[:, m, :]).view(B, 1, 1, 1, 1) for m in range(self.M)], dim=1)
        return x + bias


class RegionPrototypeModule(nn.Module):
    def __init__(self, dim, num_regions=3, proto_dim=48):
        super().__init__()
        self.feat_proj  = nn.Linear(dim, proto_dim)
        self.prototypes = nn.Parameter(torch.randn(num_regions, proto_dim))
        self.score_proj = nn.Conv3d(num_regions, num_regions, kernel_size=1)
    def forward(self, x):
        B, C, H, W, D = x.shape
        x_flat = rearrange(x, "b c h w d -> b (h w d) c")
        x_proj = self.feat_proj(x_flat)
        x_norm = F.normalize(x_proj, dim=-1)
        p_norm = F.normalize(self.prototypes, dim=-1)
        scores = torch.einsum("bnd,rd->bnr", x_norm, p_norm)
        scores_3d = rearrange(scores, "b (h w d) r -> b r h w d", h=H, w=W, d=D)
        return self.score_proj(scores_3d), scores[:, :, 2]


class ETBiasedSelfAttention(nn.Module):
    def __init__(self, dim, num_heads=4):
        super().__init__()
        self.H, self.Dh    = num_heads, dim // num_heads
        self.scale         = self.Dh ** -0.5
        self.qkv           = nn.Linear(dim, dim * 3, bias=False)
        self.out_lin       = nn.Linear(dim, dim)
        self.et_bias_scale = nn.Parameter(torch.tensor(0.05))
    def forward(self, x, et_bias=None):
        B, N, D = x.shape
        qkv = self.qkv(x).reshape(B, N, 3, self.H, self.Dh).permute(2, 0, 3, 1, 4)
        Q, K, V = qkv[0], qkv[1], qkv[2]
        attn = (Q @ K.transpose(-2, -1)) * self.scale
        if et_bias is not None:
            et_b = (self.et_bias_scale * torch.softmax(et_bias, dim=-1)).unsqueeze(1).unsqueeze(2)
            attn += et_b
        attn = torch.softmax(attn, dim=-1)
        out  = (attn @ V).transpose(1, 2).reshape(B, N, D)
        return self.out_lin(out)


class ETGuidedTransformerBlock(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.proto_module = RegionPrototypeModule(dim)
        self.gate_feat    = nn.Conv3d(dim, dim, kernel_size=1)
        self.gate_region  = nn.Conv3d(3, dim, kernel_size=1)
        self.norm1        = nn.LayerNorm(dim)
        self.attn         = ETBiasedSelfAttention(dim)
        self.norm2        = nn.LayerNorm(dim)
        self.ffn          = nn.Sequential(nn.Linear(dim, dim * 2), nn.GELU(), nn.Linear(dim * 2, dim))
    def forward(self, x):
        B, C, H, W, D = x.shape
        proto_scores, et_scores = self.proto_module(x)
        gate    = torch.sigmoid(self.gate_feat(x) + self.gate_region(proto_scores))
        x_gated = x * gate
        x_flat  = rearrange(x_gated, "b c h w d -> b (h w d) c")
        x_flat  = x_flat + self.attn(self.norm1(x_flat), et_bias=et_scores)
        x_flat  = x_flat + self.ffn(self.norm2(x_flat))
        return rearrange(x_flat, "b (h w d) c -> b c h w d", h=H, w=W, d=D), proto_scores


class RCMTUNetV4(nn.Module):
    def __init__(self, in_channels=4, out_channels=4, features=(24, 48, 96, 192)):
        super().__init__()
        f = features
        self.fusion     = ModalityFusionBlock(in_channels=in_channels, fusion_dim=f[0], num_heads=4)
        self.enc1       = ResidualBlock(in_channels, f[0])
        self.enc2       = ResidualBlock(f[0], f[1])
        self.enc3       = ResidualBlock(f[1], f[2])
        self.enc4       = ResidualBlock(f[2], f[3])
        self.pool       = nn.MaxPool3d(2)
        self.bottleneck = ETGuidedTransformerBlock(dim=f[3])
        self.up3        = nn.ConvTranspose3d(f[3], f[2], 2, stride=2)
        self.ag3        = AttentionGate(f[2], f[2], f[2]//2)
        self.dec3       = ResidualBlock(f[3], f[2])
        self.up2        = nn.ConvTranspose3d(f[2], f[1], 2, stride=2)
        self.ag2        = AttentionGate(f[1], f[1], f[1]//2)
        self.dec2       = ResidualBlock(f[2], f[1])
        self.up1        = nn.ConvTranspose3d(f[1], f[0], 2, stride=2)
        self.ag1        = AttentionGate(f[0], f[0], f[0]//2)
        self.dec1       = ResidualBlock(f[1], f[0])
        self.final_conv = nn.Conv3d(f[0], out_channels, 1)
        self.deep_sup2  = nn.Conv3d(f[1], out_channels, 1)
        self.deep_sup3  = nn.Conv3d(f[2], out_channels, 1)

    def forward(self, x):
        xf = self.fusion(x)
        e1 = self.enc1(xf)
        e2 = self.enc2(self.pool(e1))
        e3 = self.enc3(self.pool(e2))
        e4 = self.enc4(self.pool(e3))
        b, _ = self.bottleneck(e4)
        def up_and_cat(up_layer, ag, dec, b_feat, enc_feat):
            u = up_layer(b_feat)
            if u.shape[2:] != enc_feat.shape[2:]:
                u = F.interpolate(u, size=enc_feat.shape[2:], mode="trilinear", align_corners=False)
            return dec(torch.cat([u, ag(g=u, x=enc_feat)], dim=1))
        d3 = up_and_cat(self.up3, self.ag3, self.dec3, b,  e3)
        d2 = up_and_cat(self.up2, self.ag2, self.dec2, d3, e2)
        d1 = up_and_cat(self.up1, self.ag1, self.dec1, d2, e1)
        return self.final_conv(d1)