import torch
import torch.nn as nn
import timm

class AuxModel(nn.Module):
    """
    Architecture matches the 'BiomassModel' from the new aux training code.
    Outputs: [NDVI, Height]
    """
    def __init__(self, model_name):
        super().__init__()
        self.backbone = timm.create_model(model_name, pretrained=False, num_classes=0)
        img_features = self.backbone.num_features

        # AUXILIARY HEAD ONLY (NDVI, Height)
        self.aux_head = nn.Sequential(
            nn.Linear(img_features, 256),
            nn.BatchNorm1d(256),
            nn.ReLU(),
            nn.Linear(256, 2) 
        )

    def forward(self, image):
        feat = self.backbone(image)
        aux_out = self.aux_head(feat)
        return aux_out


class BiomassModel(nn.Module):
    """
    Main Stage 2 Model.
    Expects concatenated Image features + Encoded Tabular features.
    """
    def __init__(self, model_name, img_size=384):
        super().__init__()
        self.backbone = timm.create_model(model_name, pretrained=False, num_classes=0)
        
        # Calculate backbone features dynamically
        with torch.no_grad():
            dummy_input = torch.randn(1, 3, img_size, img_size)
            img_features = self.backbone(dummy_input).shape[1]
    
        self.tabular_encoder = nn.Sequential(
            nn.Linear(2, 64),
            nn.BatchNorm1d(64),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(64, 128),
            nn.BatchNorm1d(128),
            nn.ReLU(),
            nn.Dropout(0.3),
        )
        
        fusion_dim = img_features + 128
        self.fusion = nn.Sequential(
            nn.Linear(fusion_dim, 512),
            nn.BatchNorm1d(512),
            nn.ReLU(),
            nn.Dropout(0.4),
            nn.Linear(512, 256),
            nn.BatchNorm1d(256),
            nn.ReLU(),
            nn.Dropout(0.3),
        )
        
        self.head_green = nn.Linear(256, 1)
        self.head_dead = nn.Linear(256, 1)
        self.head_clover = nn.Linear(256, 1)
        self.head_gdm = nn.Linear(256, 1)
        self.head_total = nn.Linear(256, 1)
    
    def forward(self, image, tabular):
        img_features = self.backbone(image)
        tab_features = self.tabular_encoder(tabular)
        combined = torch.cat([img_features, tab_features], dim=1)
        fused = self.fusion(combined)
        
        out_green = self.head_green(fused)
        out_dead = self.head_dead(fused)
        out_clover = self.head_clover(fused)
        out_gdm = self.head_gdm(fused)
        out_total = self.head_total(fused)
        
        outputs = torch.cat([out_green, out_dead, out_clover, out_gdm, out_total], dim=1)
        return outputs