import os
import argparse
import datetime
from pathlib import Path

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from tensorboardX import SummaryWriter
import numpy as np
from tqdm import tqdm
import cv2

# Assumes you have installed mindglow in editable mode
from mindglow.seg2d.models import UNet
from mindglow.seg2d.models.encoders import get_encoder
from mindglow.seg2d.dataset.segdataset import SegDataset
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(args):
    """Build Mindglow encoder + UNet decoder."""
    encoder = get_encoder(
        name=args.encoder,
        pretrained=(args.pretrained == 'imagenet'),  # supports 'imagenet' or 'none'
        in_channels=args.in_channels,
        output_stride=args.output_stride,
    )
    model = UNet(
        encoder=encoder,
        out_channels=args.classes,
        bilinear=args.bilinear,
        dropout=args.dropout,
        norm=args.norm,
        activation=args.activation,
    )
    return model


def create_optimizer(args, model):
    """Optimizer factory."""
    if args.optimizer == 'adam':
        return optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
    elif args.optimizer == 'adamw':
        return optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
    elif args.optimizer == 'sgd':
        return optim.SGD(model.parameters(), lr=args.lr, momentum=0.9, weight_decay=args.weight_decay)
    else:
        raise ValueError(f"Unknown optimizer: {args.optimizer}")


def save_checkpoint(state, filename):
    """Save checkpoint."""
    torch.save(state, filename)
    print(f"Checkpoint saved: {filename}")


# =============== Training & Validation Epochs ===============
def train_one_epoch(model, loader, criterion, optimizer, device, epoch, writer, scaler=None):
    model.train()
    total_loss = 0
    pbar = tqdm(loader, desc=f'Train Epoch {epoch}')

    for batch_idx, batch in enumerate(pbar):
        images = batch['image'].to(device)
        masks = batch['mask'].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()})

        # Log batch loss
        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, loader, criterion, device, epoch, writer, num_classes):
    model.eval()
    total_loss = 0
    total_iou = 0
    total_dice = 0

    pbar = tqdm(loader, desc=f'Val Epoch {epoch}')

    for batch in pbar:
        images = batch['image'].to(device)
        masks = batch['mask'].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


# =============== Main ===============
def main():
    parser = argparse.ArgumentParser(description='Mindglow 2D Segmentation Training')
    # Dataset
    parser.add_argument('--dataset', '-d', type=str, required=True,
                        help='Path to dataset folder (contains train/ and val/ subfolders)')
    parser.add_argument('--image-size', type=int, default=512,
                        help='Input image size (resized)')
    parser.add_argument('--classes', type=int, required=True,
                        help='Number of segmentation classes')
    parser.add_argument('--in-channels', type=int, default=3,
                        help='Number of input image channels')

    # Model
    parser.add_argument('--encoder', type=str, default='resnet34',
                        help='Encoder backbone (e.g., resnet50, vgg16_bn, resnext50_32x4d)')
    parser.add_argument('--pretrained', type=str, default='imagenet', choices=['imagenet', 'none'],
                        help='Use ImageNet pretrained weights')
    parser.add_argument('--output-stride', type=int, default=16, choices=[8, 16, 32],
                        help='Encoder output stride (dilation control)')
    parser.add_argument('--bilinear', action='store_true', default=False,
                        help='Use bilinear upsampling instead of transposed conv')
    parser.add_argument('--dropout', type=float, default=0.1,
                        help='Dropout rate')
    parser.add_argument('--norm', type=str, default='bn', choices=['bn', 'in', 'none'],
                        help='Normalization layer')
    parser.add_argument('--activation', type=str, default='relu', choices=['relu', 'leaky', 'swish'],
                        help='Activation function')

    # Training hyperparameters
    parser.add_argument('--epochs', '-e', type=int, default=50,
                        help='Number of epochs')
    parser.add_argument('--batch-size', '-b', type=int, default=8,
                        help='Batch size')
    parser.add_argument('--lr', type=float, default=1e-3,
                        help='Learning rate')
    parser.add_argument('--weight-decay', type=float, default=1e-4,
                        help='Weight decay')
    parser.add_argument('--optimizer', type=str, default='adamw', choices=['adam', 'adamw', 'sgd'],
                        help='Optimizer')
    parser.add_argument('--loss', type=str, default='combined', choices=['ce', 'dice', 'combined'],
                        help='Loss function')
    parser.add_argument('--scheduler', type=str, default='plateau', choices=['plateau', 'step', 'none'],
                        help='Learning rate scheduler')

    # System
    parser.add_argument('--amp', action='store_true', default=False,
                        help='Use mixed precision training')
    parser.add_argument('--device', type=str, default='cuda',
                        help='Device (cuda or cpu)')
    parser.add_argument('--num-workers', type=int, default=4,
                        help='Dataloader workers')
    parser.add_argument('--seed', type=int, default=42,
                        help='Random seed')

    # Experiment
    parser.add_argument('--exp-name', type=str, default=None,
                        help='Experiment name (for logs/checkpoints)')
    parser.add_argument('--log-dir', type=str, default='./logs',
                        help='TensorBoard log directory')
    parser.add_argument('--ckpt-dir', type=str, default='./checkpoints',
                        help='Checkpoint directory')
    parser.add_argument('--resume', type=str, default=None,
                        help='Path to checkpoint to resume training')

    args = parser.parse_args()

    # Setup
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed(args.seed)
    device = torch.device(args.device if torch.cuda.is_available() else 'cpu')
    print(f"Using device: {device}")

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

    # Create directories
    os.makedirs(args.log_dir, exist_ok=True)
    os.makedirs(args.ckpt_dir, exist_ok=True)

    writer = SummaryWriter(log_dir=os.path.join(args.log_dir, args.exp_name))
    # Data
    train_loader, val_loader = create_dataloaders(args)

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

    # Loss, optimizer, scheduler
    criterion = get_loss(args.loss)
    optimizer = create_optimizer(args, model)

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

    # AMP scaler
    scaler = torch.cuda.amp.GradScaler() if args.amp else None

    # Resume training
    start_epoch = 0
    best_iou = 0.0
    if args.resume:
        if os.path.isfile(args.resume):
            checkpoint = torch.load(args.resume, map_location=device)
            model.load_state_dict(checkpoint['model_state_dict'])
            optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
            start_epoch = checkpoint['epoch'] + 1
            best_iou = checkpoint.get('best_iou', 0.0)
            print(f"Resumed from {args.resume} (epoch {start_epoch})")
        else:
            print(f"Warning: No checkpoint found at {args.resume}, starting from scratch.")

    # Training loop
    print(f"\nStarting training: {args.exp_name}")
    print(f"Encoder: {args.encoder}, Classes: {args.classes}, Epochs: {args.epochs}\n")

    for epoch in range(start_epoch, args.epochs):
        # Train
        train_loss = train_one_epoch(
            model, train_loader, criterion, optimizer, device, epoch, writer, scaler
        )

        # Validate
        val_loss, val_iou, val_dice = validate(
            model, val_loader, criterion, device, epoch, writer, args.classes
        )

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

        # Current learning rate
        current_lr = optimizer.param_groups[0]['lr']

        # Print epoch summary
        print(f"\nEpoch {epoch:3d}/{args.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(args.ckpt_dir, f"{args.exp_name}_best.pth")
            save_checkpoint({
                'epoch': epoch,
                'model_state_dict': model.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
                'best_iou': best_iou,
                'args': args,
            }, best_path)

        # Regular checkpoint every 10 epochs
        if epoch % 10 == 0 and epoch > 0:
            checkpoint_path = os.path.join(args.ckpt_dir, f"{args.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,
                'args': args,
            }, checkpoint_path)

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


if __name__ == '__main__':
    main()