"""
Train Cloud Removal Model using SEN12MS-CR Dataset
Huấn luyện model Deep Learning để khử mây từ ảnh Sentinel-2

Dataset: SEN12MS-CR (Sentinel-12 Multi-Seasonal Cloud Removal)
- Input: S2 cloudy images (ảnh Sentinel-2 bị mây)
- Target: S2 clean images (ảnh Sentinel-2 sạch)
- Optional: S1 SAR data (radar data không bị ảnh hưởng bởi mây)
"""

import os
import sys
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from pathlib import Path
import matplotlib.pyplot as plt
from tqdm import tqdm
import math
from torch.utils.tensorboard import SummaryWriter

# Add winter_dataset to path
sys.path.insert(0, str(Path(__file__).parent / "winter_dataset"))
from sen12ms_cr_dataLoader import SEN12MSCRDataset, Seasons, S1Bands, S2Bands


# ============ DATASET WRAPPER ============

class CloudRemovalDataset(Dataset):
    """
    PyTorch Dataset wrapper cho SEN12MS-CR
    Input: S2 cloudy + S1 (optional)
    Target: S2 clean
    """
    
    def __init__(self, base_dir, season=Seasons.WINTER, use_s1=True, 
                 s2_bands=S2Bands.ALL, normalize=True):
        """
        Args:
            base_dir: Đường dẫn đến thư mục chứa dữ liệu
            season: Mùa (SPRING, SUMMER, FALL, WINTER)
            use_s1: Có sử dụng dữ liệu S1 (radar) không
            s2_bands: Các band S2 cần dùng
            normalize: Normalize dữ liệu về [0, 1]
        """
        self.dataset = SEN12MSCRDataset(base_dir)
        self.season = season
        self.use_s1 = use_s1
        self.s2_bands = s2_bands
        self.normalize = normalize
        
        # Lấy tất cả scene và patch IDs
        season_ids = self.dataset.get_season_ids(season)
        
        # Tạo list of (scene_id, patch_id) pairs
        self.samples = []
        for scene_id, patch_ids in season_ids.items():
            for patch_id in patch_ids:
                self.samples.append((scene_id, patch_id))
        
        # Get band count
        n_s2_bands = len(s2_bands.value) if hasattr(s2_bands, 'value') else len(s2_bands)
        
        print(f"[DATASET] Loaded {len(self.samples)} samples from {season.value}")
        print(f"[DATASET] Use S1: {use_s1}, S2 bands: {n_s2_bands}")
    
    def __len__(self):
        return len(self.samples)
    
    def __getitem__(self, idx):
        scene_id, patch_id = self.samples[idx]
        
        # Load triplet: S1, S2 clean, S2 cloudy
        s1, s2_clean, s2_cloudy, bounds = self.dataset.get_s1s2s2cloudy_triplet(
            self.season,
            scene_id,
            patch_id,
            s1_bands=S1Bands.ALL if self.use_s1 else S1Bands.NONE,
            s2_bands=self.s2_bands,
            s2cloudy_bands=self.s2_bands
        )
        
        # Normalize to [0, 1] if needed
        if self.normalize:
            s2_clean = s2_clean.astype(np.float32) / 10000.0  # S2 values are in [0, 10000]
            s2_cloudy = s2_cloudy.astype(np.float32) / 10000.0
            if self.use_s1:
                # S1 values need different normalization (dB scale)
                s1 = (s1.astype(np.float32) + 30) / 50.0  # Normalize from [-30, 20] to [0, 1]
                s1 = np.clip(s1, 0, 1)
        
        # Convert to torch tensors
        s2_clean = torch.from_numpy(s2_clean).float()
        s2_cloudy = torch.from_numpy(s2_cloudy).float()
        
        # Input: S2 cloudy + S1 (if enabled)
        if self.use_s1:
            s1 = torch.from_numpy(s1).float()
            input_data = torch.cat([s2_cloudy, s1], dim=0)
        else:
            input_data = s2_cloudy
        
        return input_data, s2_clean


# ============ U-NET ARCHITECTURE ============

class DoubleConv(nn.Module):
    """(Conv2d -> BatchNorm -> ReLU) x 2"""
    
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.double_conv = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )
    
    def forward(self, x):
        return self.double_conv(x)


class UNet(nn.Module):
    """
    U-Net architecture cho cloud removal
    Input: S2 cloudy (+ S1 optional) [B, C_in, H, W]
    Output: S2 clean [B, C_out, H, W]
    """
    
    def __init__(self, in_channels, out_channels, features=[64, 128, 256, 512]):
        super().__init__()
        self.encoder = nn.ModuleList()
        self.decoder = nn.ModuleList()
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        
        # Encoder (downsampling)
        for feature in features:
            self.encoder.append(DoubleConv(in_channels, feature))
            in_channels = feature
        
        # Bottleneck
        self.bottleneck = DoubleConv(features[-1], features[-1] * 2)
        
        # Decoder (upsampling)
        for feature in reversed(features):
            self.decoder.append(
                nn.ConvTranspose2d(feature * 2, feature, kernel_size=2, stride=2)
            )
            self.decoder.append(DoubleConv(feature * 2, feature))
        
        # Final output layer
        self.final_conv = nn.Conv2d(features[0], out_channels, kernel_size=1)
    
    def forward(self, x):
        skip_connections = []
        
        # Encoder
        for encode in self.encoder:
            x = encode(x)
            skip_connections.append(x)
            x = self.pool(x)
        
        # Bottleneck
        x = self.bottleneck(x)
        
        # Decoder
        skip_connections = skip_connections[::-1]
        
        for idx in range(0, len(self.decoder), 2):
            x = self.decoder[idx](x)  # Upsample
            skip_connection = skip_connections[idx // 2]
            
            # Handle size mismatch
            if x.shape != skip_connection.shape:
                x = nn.functional.interpolate(x, size=skip_connection.shape[2:])
            
            concat_skip = torch.cat((skip_connection, x), dim=1)
            x = self.decoder[idx + 1](concat_skip)  # Double conv
        
        return self.final_conv(x)


# ============ TRAINING FUNCTIONS ============

def calculate_psnr(img1, img2, max_value=1.0):
    """
    Tính PSNR (Peak Signal-to-Noise Ratio) giữa 2 ảnh
    
    Args:
        img1: Tensor [B, C, H, W] hoặc [C, H, W]
        img2: Tensor [B, C, H, W] hoặc [C, H, W]
        max_value: Giá trị pixel tối đa (thường là 1.0 cho ảnh đã normalize)
    
    Returns:
        psnr_value: PSNR tính bằng dB
    """
    mse = torch.mean((img1 - img2) ** 2)
    if mse == 0:
        return float('inf')
    psnr = 20 * math.log10(max_value) - 10 * torch.log10(mse)
    return psnr.item()


def calculate_baseline_psnr(dataloader, device):
    """
    Tính Cận dưới (Baseline PSNR) - PSNR giữa ảnh cloudy và ảnh clean
    Đây là giá trị tham chiếu để đánh giá model có tốt hơn không làm gì không.
    
    Args:
        dataloader: Validation dataloader
        device: 'cuda' hoặc 'cpu'
    
    Returns:
        baseline_psnr: PSNR trung bình giữa input cloudy và target clean
    """
    total_psnr = 0
    num_batches = 0
    
    print("\n📏 Đang tính Cận dưới (Baseline PSNR)...")
    print("   Baseline = PSNR(s2_cloudy, s2_clean)")
    
    with torch.no_grad():
        for inputs, targets in tqdm(dataloader, desc="Calculating Baseline"):
            inputs = inputs.to(device)
            targets = targets.to(device)
            
            # Lấy phần S2 cloudy từ input
            # Input có thể là [S2_cloudy] hoặc [S2_cloudy, S1]
            # Output channels = target channels
            num_s2_bands = targets.shape[1]
            s2_cloudy = inputs[:, :num_s2_bands, :, :]
            
            # Tính PSNR giữa cloudy và clean
            psnr = calculate_psnr(s2_cloudy, targets)
            total_psnr += psnr
            num_batches += 1
    
    baseline_psnr = total_psnr / num_batches
    print(f"\n✅ Cận dưới (Baseline PSNR): {baseline_psnr:.2f} dB")
    print(f"   → Đây là 'thanh thước đo' - model phải vượt qua giá trị này!")
    
    return baseline_psnr


def train_epoch(model, dataloader, criterion, optimizer, device):
    """Train for one epoch"""
    model.train()
    total_loss = 0
    
    pbar = tqdm(dataloader, desc="Training")
    for batch_idx, (inputs, targets) in enumerate(pbar):
        inputs = inputs.to(device)
        targets = targets.to(device)
        
        # Forward pass
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        
        # Backward pass
        loss.backward()
        optimizer.step()
        
        total_loss += loss.item()
        pbar.set_postfix({'loss': loss.item()})
    
    return total_loss / len(dataloader)


def validate(model, dataloader, criterion, device):
    """Validate model và tính PSNR"""
    model.eval()
    total_loss = 0
    total_psnr = 0
    
    with torch.no_grad():
        for inputs, targets in tqdm(dataloader, desc="Validation"):
            inputs = inputs.to(device)
            targets = targets.to(device)
            
            outputs = model(inputs)
            loss = criterion(outputs, targets)
            total_loss += loss.item()
            
            # Tính PSNR giữa output (fake_s2) và target (s2_clean)
            psnr = calculate_psnr(outputs, targets)
            total_psnr += psnr
    
    avg_loss = total_loss / len(dataloader)
    avg_psnr = total_psnr / len(dataloader)
    
    return avg_loss, avg_psnr


def visualize_results(model, dataset, device, num_samples=3):
    """Visualize cloud removal results"""
    model.eval()
    
    fig, axes = plt.subplots(num_samples, 3, figsize=(15, 5 * num_samples))
    
    with torch.no_grad():
        for i in range(num_samples):
            idx = np.random.randint(0, len(dataset))
            input_data, target = dataset[idx]
            
            input_data = input_data.unsqueeze(0).to(device)
            output = model(input_data)
            
            # Convert to numpy
            input_rgb = input_data[0, :3, :, :].cpu().numpy().transpose(1, 2, 0)
            target_rgb = target[:3, :, :].cpu().numpy().transpose(1, 2, 0)
            output_rgb = output[0, :3, :, :].cpu().numpy().transpose(1, 2, 0)
            
            # Clip to [0, 1]
            input_rgb = np.clip(input_rgb * 3, 0, 1)  # Enhance for visualization
            target_rgb = np.clip(target_rgb * 3, 0, 1)
            output_rgb = np.clip(output_rgb * 3, 0, 1)
            
            if num_samples == 1:
                axes[0].imshow(input_rgb)
                axes[0].set_title("Input (Cloudy)")
                axes[0].axis('off')
                
                axes[1].imshow(output_rgb)
                axes[1].set_title("Output (Predicted)")
                axes[1].axis('off')
                
                axes[2].imshow(target_rgb)
                axes[2].set_title("Target (Clean)")
                axes[2].axis('off')
            else:
                axes[i, 0].imshow(input_rgb)
                axes[i, 0].set_title(f"Sample {i+1}: Input (Cloudy)")
                axes[i, 0].axis('off')
                
                axes[i, 1].imshow(output_rgb)
                axes[i, 1].set_title(f"Sample {i+1}: Output (Predicted)")
                axes[i, 1].axis('off')
                
                axes[i, 2].imshow(target_rgb)
                axes[i, 2].set_title(f"Sample {i+1}: Target (Clean)")
                axes[i, 2].axis('off')
    
    plt.tight_layout()
    return fig


# ============ MAIN TRAINING SCRIPT ============

def train_cloud_removal_model(
    data_dir="winter_dataset",
    use_s1=True,
    batch_size=8,
    num_epochs=50,
    learning_rate=1e-4,
    device="cuda" if torch.cuda.is_available() else "cpu",
    save_dir="model_train",
    status_callback=None
):
    """
    Train cloud removal model
    
    Args:
        data_dir: Thư mục chứa dữ liệu SEN12MS-CR
        use_s1: Có sử dụng S1 radar data không
        batch_size: Batch size
        num_epochs: Số epochs
        learning_rate: Learning rate
        device: 'cuda' hoặc 'cpu'
        save_dir: Thư mục lưu model
        status_callback: Callback function(epoch, train_loss, val_loss, val_psnr, baseline_psnr)
    """
    
    print("=" * 70)
    print("🌥️  CLOUD REMOVAL MODEL TRAINING")
    print("=" * 70)
    print(f"Data directory: {data_dir}")
    print(f"Use S1 (SAR): {use_s1}")
    print(f"Device: {device}")
    print(f"Batch size: {batch_size}")
    print(f"Epochs: {num_epochs}")
    print(f"Learning rate: {learning_rate}")
    print("=" * 70)
    
    # Create dataset
    print("\n📂 Loading dataset...")
    
    # Use RGB + NIR bands for training (B02, B03, B04, B08)
    s2_bands = [S2Bands.B02, S2Bands.B03, S2Bands.B04, S2Bands.B08]
    
    dataset = CloudRemovalDataset(
        base_dir=data_dir,
        season=Seasons.WINTER,
        use_s1=use_s1,
        s2_bands=s2_bands,
        normalize=True
    )
    
    # Split train/val
    train_size = int(0.8 * len(dataset))
    val_size = len(dataset) - train_size
    train_dataset, val_dataset = torch.utils.data.random_split(
        dataset, [train_size, val_size]
    )
    
    print(f"Train samples: {len(train_dataset)}")
    print(f"Val samples: {len(val_dataset)}")
    
    # Create dataloaders
    train_loader = DataLoader(
        train_dataset, 
        batch_size=batch_size, 
        shuffle=True,
        num_workers=4,
        pin_memory=True if device == "cuda" else False
    )
    
    val_loader = DataLoader(
        val_dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=4,
        pin_memory=True if device == "cuda" else False
    )
    
    # Create model
    print("\n🏗️  Creating U-Net model...")
    in_channels = len(s2_bands) + (2 if use_s1 else 0)  # S2 + S1 (VV, VH)
    out_channels = len(s2_bands)
    
    model = UNet(in_channels=in_channels, out_channels=out_channels)
    model = model.to(device)
    
    print(f"Input channels: {in_channels}")
    print(f"Output channels: {out_channels}")
    print(f"Model parameters: {sum(p.numel() for p in model.parameters()):,}")
    
    # Loss and optimizer
    criterion = nn.L1Loss()  # MAE loss
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)
    scheduler = optim.lr_scheduler.ReduceLROnPlateau(
        optimizer, mode='min', factor=0.5, patience=5
    )
    
    # Tính Baseline PSNR trước khi bắt đầu train
    print("\n" + "="*70)
    print("📏 TÍNH CẬN DƯỚI (BASELINE) - CHỈ TÍNH 1 LẦN DUY NHẤT")
    print("="*70)
    baseline_psnr = calculate_baseline_psnr(val_loader, device)
    print("="*70)
    
    # Setup Tensorboard
    writer = SummaryWriter(log_dir=str(Path(save_dir) / "runs" / "cloud_removal"))
    print(f"\n📊 Tensorboard logging: {Path(save_dir) / 'runs' / 'cloud_removal'}")
    print(f"   Chạy lệnh: tensorboard --logdir={Path(save_dir) / 'runs'}")
    
    # Training loop
    print("\n🚀 Starting training...")
    best_val_loss = float('inf')
    best_val_psnr = 0
    train_losses = []
    val_losses = []
    val_psnrs = []
    
    for epoch in range(num_epochs):
        print(f"\n{'='*70}")
        print(f"Epoch {epoch + 1}/{num_epochs}")
        print(f"{'='*70}")
        
        # Train
        train_loss = train_epoch(model, train_loader, criterion, optimizer, device)
        train_losses.append(train_loss)
        
        # Validate
        val_loss, val_psnr = validate(model, val_loader, criterion, device)
        val_losses.append(val_loss)
        val_psnrs.append(val_psnr)
        
        # Update learning rate
        scheduler.step(val_loss)
        
        # Tensorboard logging
        writer.add_scalar('Loss/train', train_loss, epoch)
        writer.add_scalar('Loss/val', val_loss, epoch)
        writer.add_scalar('PSNR/model', val_psnr, epoch)
        writer.add_scalar('PSNR/baseline', baseline_psnr, epoch)
        writer.add_scalar('PSNR/improvement', val_psnr - baseline_psnr, epoch)
        writer.add_scalar('Learning_Rate', optimizer.param_groups[0]['lr'], epoch)
        
        # Tính độ cải thiện so với baseline
        improvement = val_psnr - baseline_psnr
        improvement_percent = (improvement / baseline_psnr) * 100 if baseline_psnr > 0 else 0
        
        print(f"\nEpoch {epoch + 1} Summary:")
        print(f"  Train Loss:     {train_loss:.6f}")
        print(f"  Val Loss:       {val_loss:.6f}")
        print(f"  ")
        print(f"  📊 PSNR Comparison:")
        print(f"  ├─ Model PSNR:      {val_psnr:.2f} dB")
        print(f"  ├─ Baseline PSNR:   {baseline_psnr:.2f} dB (cận dưới)")
        print(f"  └─ Improvement:     {improvement:+.2f} dB ({improvement_percent:+.1f}%)")
        if improvement > 0:
            print(f"     ✅ Model đang TỐT HƠN baseline!")
        else:
            print(f"     ⚠️  Model chưa vượt qua baseline")
        
        # Call status callback if provided
        if status_callback:
            try:
                status_callback(epoch + 1, train_loss, val_loss, val_psnr, baseline_psnr)
            except Exception as cb_err:
                print(f"[WARNING] Status callback error: {cb_err}")
        
        # Save best model (based on PSNR)
        if val_psnr > best_val_psnr:
            best_val_psnr = val_psnr
            best_val_loss = val_loss
            save_path = Path(save_dir) / "cloud_removal_unet_best.pth"
            save_path.parent.mkdir(exist_ok=True)
            
            torch.save({
                'epoch': epoch,
                'model_state_dict': model.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
                'train_loss': train_loss,
                'val_loss': val_loss,
                'val_psnr': val_psnr,
                'baseline_psnr': baseline_psnr,
                'use_s1': use_s1,
                'in_channels': in_channels,
                'out_channels': out_channels,
                # Hyperparameters for reproducibility
                'hyperparameters': {
                    'data_dir': data_dir,
                    'use_s1': use_s1,
                    'batch_size': batch_size,
                    'num_epochs': num_epochs,
                    'learning_rate': learning_rate,
                    'device': device,
                    'model_type': 'unet',
                    'architecture': 'U-Net',
                    'optimizer': 'Adam',
                    'criterion': 'L1Loss',
                    'scheduler': 'ReduceLROnPlateau',
                    'train_size': len(train_dataset),
                    'val_size': len(val_dataset),
                    'test_size': len(val_dataset) / (len(train_dataset) + len(val_dataset)),
                    's2_bands': len(s2_bands),
                    'feature_mode': 'image',
                    'random_seed': 42  # Default seed
                }
            }, save_path)
            
            print(f"  💾 Saved best model (highest PSNR): {save_path}")
        
        # Visualize every 10 epochs
        if (epoch + 1) % 10 == 0:
            print("\n📊 Generating visualizations...")
            fig = visualize_results(model, val_dataset, device, num_samples=3)
            
            viz_path = Path(save_dir) / f"cloud_removal_epoch_{epoch+1}.png"
            fig.savefig(viz_path, dpi=150, bbox_inches='tight')
            plt.close(fig)
            
            print(f"  💾 Saved visualization: {viz_path}")
    
    # Close Tensorboard writer
    writer.close()
    
    # Plot training curves
    print("\n📈 Plotting training curves...")
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 6))
    
    # Loss curves
    ax1.plot(train_losses, label='Train Loss', linewidth=2)
    ax1.plot(val_losses, label='Val Loss', linewidth=2)
    ax1.set_xlabel('Epoch', fontsize=12)
    ax1.set_ylabel('Loss (MAE)', fontsize=12)
    ax1.set_title('Loss Curves', fontsize=14, fontweight='bold')
    ax1.legend(fontsize=10)
    ax1.grid(True, alpha=0.3)
    
    # PSNR curves với baseline
    epochs = range(len(val_psnrs))
    ax2.plot(epochs, val_psnrs, label='Model PSNR', linewidth=2, color='green')
    ax2.axhline(y=baseline_psnr, color='red', linestyle='--', linewidth=2, 
                label=f'Baseline PSNR ({baseline_psnr:.2f} dB)')
    ax2.fill_between(epochs, baseline_psnr, val_psnrs, 
                      where=[p > baseline_psnr for p in val_psnrs],
                      alpha=0.3, color='green', label='Model > Baseline')
    ax2.fill_between(epochs, baseline_psnr, val_psnrs,
                      where=[p <= baseline_psnr for p in val_psnrs],
                      alpha=0.3, color='red', label='Model < Baseline')
    ax2.set_xlabel('Epoch', fontsize=12)
    ax2.set_ylabel('PSNR (dB)', fontsize=12)
    ax2.set_title('PSNR vs Baseline (Cận dưới)', fontsize=14, fontweight='bold')
    ax2.legend(fontsize=10)
    ax2.grid(True, alpha=0.3)
    
    plt.tight_layout()
    curve_path = Path(save_dir) / "training_curves.png"
    fig.savefig(curve_path, dpi=150, bbox_inches='tight')
    plt.close(fig)
    
    print(f"  💾 Saved training curves: {curve_path}")
    
    # Final summary
    print("\n" + "=" * 70)
    print("✅ TRAINING COMPLETED!")
    print("=" * 70)
    print(f"Best validation loss: {best_val_loss:.6f}")
    print(f"Best validation PSNR: {best_val_psnr:.2f} dB")
    print(f"Baseline PSNR:        {baseline_psnr:.2f} dB")
    print(f"Final improvement:    {best_val_psnr - baseline_psnr:+.2f} dB")
    print(f"Model saved to: {Path(save_dir) / 'cloud_removal_unet_best.pth'}")
    print(f"Tensorboard logs: {Path(save_dir) / 'runs' / 'cloud_removal'}")
    print("=" * 70)
    
    return model, train_losses, val_losses, val_psnrs, baseline_psnr


# ============ INFERENCE FUNCTION ============

def apply_cloud_removal(model_path, cloudy_image, s1_data=None, device="cuda"):
    """
    Áp dụng model để khử mây cho một ảnh
    
    Args:
        model_path: Đường dẫn đến model đã train
        cloudy_image: Ảnh S2 bị mây [C, H, W]
        s1_data: Dữ liệu S1 (optional) [2, H, W]
        device: 'cuda' hoặc 'cpu'
    
    Returns:
        cleaned_image: Ảnh đã khử mây [C, H, W]
    """
    # Load model
    checkpoint = torch.load(model_path, map_location=device)
    
    model = UNet(
        in_channels=checkpoint['in_channels'],
        out_channels=checkpoint['out_channels']
    )
    model.load_state_dict(checkpoint['model_state_dict'])
    model = model.to(device)
    model.eval()
    
    # Prepare input
    input_tensor = torch.from_numpy(cloudy_image).float().unsqueeze(0).to(device)
    
    if checkpoint['use_s1'] and s1_data is not None:
        s1_tensor = torch.from_numpy(s1_data).float().unsqueeze(0).to(device)
        input_tensor = torch.cat([input_tensor, s1_tensor], dim=1)
    
    # Inference
    with torch.no_grad():
        output = model(input_tensor)
    
    cleaned_image = output[0].cpu().numpy()
    
    return cleaned_image


if __name__ == "__main__":
    # Train model
    model, train_losses, val_losses = train_cloud_removal_model(
        data_dir="winter_dataset",
        use_s1=True,
        batch_size=8,
        num_epochs=50,
        learning_rate=1e-4
    )