"""
Training utilities for 2D segmentation models.

Provides a flexible `train` function that accepts a configuration dictionary,
handles data loading, model creation, optimizer setup, and the full training loop.
"""

import os
import datetime
from typing import Dict, Any, Optional, Callable, Union

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from tensorboardX import SummaryWriter
from tqdm import tqdm

# Import necessary components from the library
from mindglow.seg2d.models import UNet
from mindglow.seg2d.models.encoders import get_encoder
from mindglow.seg2d.dataset.data_loaders import create_dataloaders
from mindglow.seg2d.metrics.iou import compute_iou
from mindglow.seg2d.metrics.dice import compute_dice
from mindglow.seg2d.losses import get_loss


def create_model(config: Dict[str, Any]) -> nn.Module:
    """Build encoder + UNet decoder from config."""
    encoder = get_encoder(
        name=config['encoder'],
        pretrained=(config['pretrained'] == 'imagenet'),
        in_channels=config['in_channels'],
        output_stride=config['output_stride'],
    )
    model = UNet(
        encoder=encoder,
        out_channels=config['classes'],
        bilinear=config.get('bilinear', False),
        dropout=config.get('dropout', 0.0),
        norm=config.get('norm', 'bn'),
        activation=config.get('activation', 'relu'),
    )
    return model


def create_optimizer(config: Dict[str, Any], model: nn.Module) -> optim.Optimizer:
    """Create optimizer based on config."""
    opt_name = config['optimizer'].lower()
    lr = config['lr']
    wd = config.get('weight_decay', 0.0)

    if opt_name == 'adam':
        return optim.Adam(model.parameters(), lr=lr, weight_decay=wd)
    elif opt_name == 'adamw':
        return optim.AdamW(model.parameters(), lr=lr, weight_decay=wd)
    elif opt_name == 'sgd':
        return optim.SGD(model.parameters(), lr=lr, momentum=0.9, weight_decay=wd)
    else:
        raise ValueError(f"Unsupported optimizer: {opt_name}")


def save_checkpoint(state: Dict[str, Any], filename: str) -> None:
    """Save a training checkpoint."""
    torch.save(state, filename)
    print(f"Checkpoint saved: {filename}")


def train_one_epoch(
    model: nn.Module,
    loader: DataLoader,
    criterion: Union[nn.Module,Callable],
    optimizer: optim.Optimizer,
    device: torch.device,
    epoch: int,
    writer: SummaryWriter,
    scaler: Optional[torch.cuda.amp.GradScaler] = None
) -> float:
    """Run one training epoch."""
    model.train()
    total_loss = 0.0
    pbar = tqdm(loader, desc=f'Train Epoch {epoch}')

    for batch_idx, batch in enumerate(pbar):
        if isinstance(batch,(list,tuple)):
            images,masks = batch
        else:
            images = batch['image'].to(device)
            masks = batch['mask'].to(device)

        images = images.to(device)
        masks = masks.to(device)

        optimizer.zero_grad()

        if scaler is not None:
            with torch.cuda.amp.autocast():
                outputs = model(images)
                loss = criterion(outputs, masks)
            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()
        else:
            outputs = model(images)
            loss = criterion(outputs, masks)
            loss.backward()
            optimizer.step()

        total_loss += loss.item()
        pbar.set_postfix({'loss': loss.item()})

        if batch_idx % 20 == 0:
            writer.add_scalar('train/batch_loss', loss.item(),
                              epoch * len(loader) + batch_idx)

    avg_loss = total_loss / len(loader)
    writer.add_scalar('train/epoch_loss', avg_loss, epoch)
    return avg_loss


@torch.no_grad()
def validate(
    model: nn.Module,
    loader: DataLoader,
    criterion: Union[nn.Module,Callable],
    device: torch.device,
    epoch: int,
    writer: SummaryWriter,
    num_classes: int
) -> tuple:
    """Run validation and return loss, mIoU, Dice."""
    model.eval()
    total_loss = 0.0
    total_iou = 0.0
    total_dice = 0.0
    pbar = tqdm(loader, desc=f'Val Epoch {epoch}')

    for batch in pbar:
        if isinstance(batch,(list,tuple)):
            images,masks = batch
        else:
            images = batch['image'].to(device)
            masks = batch['mask'].to(device)

            images = images.to(device)
            masks = masks.to(device)

        outputs = model(images)
        loss = criterion(outputs, masks)

        total_loss += loss.item()
        total_iou += compute_iou(outputs, masks, num_classes)
        total_dice += compute_dice(outputs, masks, num_classes)

        pbar.set_postfix({'loss': loss.item()})

    avg_loss = total_loss / len(loader)
    avg_iou = total_iou / len(loader)
    avg_dice = total_dice / len(loader)

    writer.add_scalar('val/epoch_loss', avg_loss, epoch)
    writer.add_scalar('val/mIoU', avg_iou, epoch)
    writer.add_scalar('val/Dice', avg_dice, epoch)

    return avg_loss, avg_iou, avg_dice


def train(config: Dict[str, Any]) -> None:
    """
    Main training routine.

    Args:
        config: A dictionary containing all training parameters.
                Required keys:
                    - dataset: path to dataset folder (with train/ and val/ subfolders)
                    - classes: number of segmentation classes
                    - encoder: backbone name (e.g., 'resnet34')
                    - epochs: number of epochs
                    - batch_size: batch size
                    - lr: learning rate
                    - optimizer: 'adam', 'adamw', or 'sgd'
                    - loss: loss function name (as understood by get_loss)
                Optional keys with defaults:
                    - image_size: 512
                    - in_channels: 3
                    - pretrained: 'imagenet' or 'none' (default 'imagenet')
                    - output_stride: 16 (8,16,32)
                    - bilinear: False
                    - dropout: 0.0
                    - norm: 'bn'
                    - activation: 'relu'
                    - weight_decay: 1e-4
                    - scheduler: 'plateau', 'step', or 'none' (default 'plateau')
                    - amp: False
                    - device: 'cuda' or 'cpu' (auto fallback)
                    - num_workers: 4
                    - seed: 42
                    - exp_name: auto-generated if None
                    - log_dir: './logs'
                    - ckpt_dir: './checkpoints'
                    - resume: path to checkpoint (default None)
    """
    # ===== Setup =====
    torch.manual_seed(config.get('seed', 42))
    torch.cuda.manual_seed(config.get('seed', 42))
    device = torch.device(config.get('device', 'cuda') if torch.cuda.is_available() else 'cpu')
    print(f"Using device: {device}")

    # Experiment name
    exp_name = config.get('exp_name')
    if exp_name is None:
        exp_name = f"{config['encoder']}_{datetime.datetime.now():%Y%m%d_%H%M%S}"
        config['exp_name'] = exp_name

    # Create directories
    log_dir = config.get('log_dir', './logs')
    ckpt_dir = config.get('ckpt_dir', './checkpoints')
    os.makedirs(log_dir, exist_ok=True)
    os.makedirs(ckpt_dir, exist_ok=True)

    writer = SummaryWriter(log_dir=os.path.join(log_dir, exp_name))

    # ===== Data =====
    train_loader, val_loader = create_dataloaders(config)

    # ===== Model =====
    model = create_model(config)
    model = model.to(device)
    print(f"Model: {config['encoder']} + UNet")
    print(f"Trainable parameters: {sum(p.numel() for p in model.parameters() if p.requires_grad):,}")

    # ===== Loss, Optimizer, Scheduler =====
    criterion = get_loss(config['loss'])
    optimizer = create_optimizer(config, model)

    scheduler = None
    sched_name = config.get('scheduler', 'plateau')
    if sched_name == 'plateau':
        scheduler = optim.lr_scheduler.ReduceLROnPlateau(
            optimizer, mode='min', patience=5, factor=0.5
        )
    elif sched_name == 'step':
        scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.1)

    # ===== Mixed precision =====
    scaler = torch.cuda.amp.GradScaler() if config.get('amp', False) else None

    # ===== Resume =====
    start_epoch = 0
    best_iou = 0.0
    resume_path = config.get('resume')
    if resume_path and os.path.isfile(resume_path):
        checkpoint = torch.load(resume_path, map_location=device)
        model.load_state_dict(checkpoint['model_state_dict'])
        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
        start_epoch = checkpoint.get('epoch', -1) + 1
        best_iou = checkpoint.get('best_iou', 0.0)
        print(f"Resumed from {resume_path} (epoch {start_epoch})")
    elif resume_path:
        print(f"Warning: Checkpoint not found at {resume_path}, starting fresh.")

    # ===== Training Loop =====
    print(f"\nStarting training: {exp_name}")
    print(f"Encoder: {config['encoder']}, Classes: {config['classes']}, Epochs: {config['epochs']}\n")

    for epoch in range(start_epoch, config['epochs']):
        train_loss = train_one_epoch(
            model, train_loader, criterion, optimizer, device, epoch, writer, scaler
        )
        val_loss, val_iou, val_dice = validate(
            model, val_loader, criterion, device, epoch, writer, config['classes']
        )

        if scheduler is not None:
            if isinstance(scheduler, optim.lr_scheduler.ReduceLROnPlateau):
                scheduler.step(val_loss)
            else:
                scheduler.step()

        current_lr = optimizer.param_groups[0]['lr']

        print(f"\nEpoch {epoch:3d}/{config['epochs']-1} | "
              f"Train Loss: {train_loss:.4f} | "
              f"Val Loss: {val_loss:.4f} | "
              f"Val mIoU: {val_iou:.4f} | "
              f"Val Dice: {val_dice:.4f} | "
              f"LR: {current_lr:.2e}")

        # Save best model
        if val_iou > best_iou:
            best_iou = val_iou
            best_path = os.path.join(ckpt_dir, f"{exp_name}_best.pth")
            save_checkpoint({
                'epoch': epoch,
                'model_state_dict': model.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
                'best_iou': best_iou,
                'config': config,
            }, best_path)

        # Regular checkpoint every 10 epochs
        if epoch % 10 == 0 and epoch > 0:
            ckpt_path = os.path.join(ckpt_dir, f"{exp_name}_epoch{epoch}.pth")
            save_checkpoint({
                'epoch': epoch,
                'model_state_dict': model.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
                'best_iou': best_iou,
                'config': config,
            }, ckpt_path)

    writer.close()
    print(f"\nTraining complete! Best mIoU: {best_iou:.4f}")
    print(f"Best model saved at: {os.path.join(ckpt_dir, f'{exp_name}_best.pth')}")