import pytorch_lightning as pl
import torch
import torch.nn as nn
import timm
from torchmetrics import Accuracy
from utils.focal_loss import FocalLoss

class MorphDetectorPL(pl.LightningModule):
    """
    PyTorch Lightning module for Morph Detection.
    Dual-stream: Vision Transformer + EfficientNet backbone, fused with shallow classifier head.
    Uses FocalLoss for hard example mining.
    """
    def __init__(
        self,
        vit_model_name: str = 'vit_base_patch16_224',
        effnet_model_name: str = 'efficientnet_b0',
        pretrained: bool = True,
        lr: float = 1e-4,
        dropout: float = 0.2,
        focal_alpha: float = 0.25,
        focal_gamma: float = 2.0,
    ):
        super().__init__()
        # save hyperparameters
        self.save_hyperparameters()

        # Global feature extractor
        self.vit = timm.create_model(
            vit_model_name,
            pretrained=pretrained,
            num_classes=0,
        )
        # Local feature extractor
        self.effnet = timm.create_model(
            effnet_model_name,
            pretrained=pretrained,
            num_classes=0,
        )
        # Fusion classifier head
        fused_dim = self.vit.num_features + self.effnet.num_features
        self.classifier = nn.Sequential(
            nn.Linear(fused_dim, 256),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(256, 1),
            nn.Sigmoid(),
        )

        # Loss and metrics
        self.criterion = FocalLoss(alpha=focal_alpha, gamma=focal_gamma)
        self.train_acc = Accuracy()
        self.val_acc = Accuracy()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x: (B, C, H, W)
        feat1 = self.vit(x)
        feat2 = self.effnet(x)
        fused = torch.cat([feat1, feat2], dim=1)
        return self.classifier(fused)

    def training_step(self, batch, batch_idx):
        imgs, labels = batch
        labels = labels.float().unsqueeze(1)
        logits = self(imgs)
        loss = self.criterion(logits, labels)
        preds = (logits > 0.5).int()
        self.train_acc.update(preds, labels.int())
        self.log('train_loss', loss, on_step=True, on_epoch=True)
        self.log('train_acc', self.train_acc, on_step=False, on_epoch=True)
        return loss

    def validation_step(self, batch, batch_idx):
        imgs, labels = batch
        labels = labels.float().unsqueeze(1)
        logits = self(imgs)
        loss = self.criterion(logits, labels)
        preds = (logits > 0.5).int()
        self.val_acc.update(preds, labels.int())
        self.log('val_loss', loss, prog_bar=True)
        self.log('val_acc', self.val_acc, prog_bar=True)

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=self.hparams.lr)
        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
            optimizer, mode='min', factor=0.5, patience=3
        )
        return {
            'optimizer': optimizer,
            'lr_scheduler': {
                'scheduler': scheduler,
                'monitor': 'val_loss',
            }
        }