"""
==================================================================================
    Action Recognition Dataset - Advanced Model Training Script
==================================================================================
    Project: Action Recognition Dataset
    
    Advanced training script with:
    - Multiple model architectures (C3D, R3D, R(2+1)D)
    - Data augmentation pipeline
    - Learning rate scheduling
    - Early stopping
    - TensorBoard logging
    - Model checkpointing
    - Mixed precision training
    
==================================================================================
    DEVELOPER INFORMATION
==================================================================================
    Website: RSK World (https://rskworld.in)
    Founded by: Molla Samser
    Designer & Tester: Rima Khatun
    Contact: help@rskworld.in | +91 93305 39277
    
    (c) 2026 RSK World. All Rights Reserved.
==================================================================================
"""

import os
import sys
import json
import argparse
from pathlib import Path
from datetime import datetime
import warnings
warnings.filterwarnings('ignore')

# Install dependencies if needed
try:
    import torch
    import torch.nn as nn
    import torch.optim as optim
    from torch.utils.data import Dataset, DataLoader
    import cv2
    import numpy as np
    from tqdm import tqdm
except ImportError:
    import subprocess
    print("Installing required packages...")
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", 
                          "torch", "torchvision", "opencv-python", "numpy", "tqdm"])
    import torch
    import torch.nn as nn
    import torch.optim as optim
    from torch.utils.data import Dataset, DataLoader
    import cv2
    import numpy as np
    from tqdm import tqdm


# ==================================================================================
# Configuration
# ==================================================================================

class Config:
    """Training configuration"""
    # Data
    DATA_DIR = "sample_data"
    NUM_FRAMES = 16
    FRAME_SIZE = (112, 112)
    NUM_WORKERS = 4
    
    # Training
    BATCH_SIZE = 8
    EPOCHS = 50
    LEARNING_RATE = 1e-4
    WEIGHT_DECAY = 1e-5
    
    # Model
    MODEL_TYPE = "r3d_18"  # c3d, r3d_18, r2plus1d_18
    PRETRAINED = True
    
    # Augmentation
    USE_AUGMENTATION = True
    HORIZONTAL_FLIP_PROB = 0.5
    COLOR_JITTER = True
    RANDOM_CROP = True
    
    # Checkpointing
    CHECKPOINT_DIR = "checkpoints"
    SAVE_EVERY = 5
    
    # Device
    DEVICE = "cuda" if torch.cuda.is_available() else "cpu"


# ==================================================================================
# Data Augmentation
# ==================================================================================

class VideoAugmentation:
    """
    Advanced video augmentation pipeline
    
    RSK World (https://rskworld.in)
    """
    
    def __init__(self, config):
        self.config = config
        
    def __call__(self, frames):
        """Apply augmentations to video frames"""
        if not self.config.USE_AUGMENTATION:
            return frames
        
        # Horizontal flip
        if np.random.random() < self.config.HORIZONTAL_FLIP_PROB:
            frames = frames[:, :, ::-1, :].copy()
        
        # Color jitter
        if self.config.COLOR_JITTER:
            frames = self._color_jitter(frames)
        
        # Random crop
        if self.config.RANDOM_CROP:
            frames = self._random_crop(frames)
        
        return frames
    
    def _color_jitter(self, frames, brightness=0.2, contrast=0.2, saturation=0.2):
        """Apply random color jittering"""
        # Brightness
        if np.random.random() > 0.5:
            factor = 1 + np.random.uniform(-brightness, brightness)
            frames = np.clip(frames * factor, 0, 255)
        
        # Contrast
        if np.random.random() > 0.5:
            factor = 1 + np.random.uniform(-contrast, contrast)
            mean = frames.mean()
            frames = np.clip((frames - mean) * factor + mean, 0, 255)
        
        return frames.astype(np.uint8)
    
    def _random_crop(self, frames, scale=(0.8, 1.0)):
        """Apply random cropping"""
        t, h, w, c = frames.shape
        
        # Random scale
        scale_factor = np.random.uniform(scale[0], scale[1])
        new_h, new_w = int(h * scale_factor), int(w * scale_factor)
        
        # Random position
        top = np.random.randint(0, h - new_h + 1)
        left = np.random.randint(0, w - new_w + 1)
        
        # Crop and resize
        cropped = frames[:, top:top+new_h, left:left+new_w, :]
        
        # Resize back to original size
        resized = np.array([cv2.resize(f, (w, h)) for f in cropped])
        
        return resized


# ==================================================================================
# Dataset
# ==================================================================================

class ActionDataset(Dataset):
    """
    Advanced Action Recognition Dataset
    
    RSK World (https://rskworld.in)
    Founder: Molla Samser
    """
    
    def __init__(self, root_dir, split='train', config=None, augment=True):
        self.root_dir = Path(root_dir) / split
        self.config = config or Config()
        self.augment = augment and split == 'train'
        self.augmentation = VideoAugmentation(self.config) if self.augment else None
        
        # Get classes
        self.classes = sorted([d.name for d in self.root_dir.iterdir() if d.is_dir()])
        self.class_to_idx = {c: i for i, c in enumerate(self.classes)}
        
        # Build samples list
        self.samples = []
        for cls in self.classes:
            cls_dir = self.root_dir / cls
            for video_file in cls_dir.iterdir():
                if video_file.suffix.lower() in ['.mp4', '.avi', '.mov']:
                    self.samples.append((str(video_file), self.class_to_idx[cls]))
        
        print(f"[RSK World] Loaded {len(self.samples)} videos from {split}")
        print(f"[RSK World] Classes: {self.classes}")
    
    def __len__(self):
        return len(self.samples)
    
    def _load_video(self, video_path):
        """Load video frames"""
        cap = cv2.VideoCapture(video_path)
        frames = []
        
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            frame = cv2.resize(frame, self.config.FRAME_SIZE)
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            frames.append(frame)
        
        cap.release()
        
        if len(frames) == 0:
            return np.zeros((self.config.NUM_FRAMES, *self.config.FRAME_SIZE, 3), dtype=np.uint8)
        
        # Sample frames uniformly
        indices = np.linspace(0, len(frames) - 1, self.config.NUM_FRAMES).astype(int)
        return np.array([frames[i] for i in indices])
    
    def __getitem__(self, idx):
        video_path, label = self.samples[idx]
        frames = self._load_video(video_path)
        
        # Apply augmentation
        if self.augmentation:
            frames = self.augmentation(frames)
        
        # Normalize and convert to tensor
        frames = frames.astype(np.float32) / 255.0
        frames = torch.FloatTensor(frames).permute(3, 0, 1, 2)  # (C, T, H, W)
        
        return frames, label


# ==================================================================================
# Models
# ==================================================================================

class C3D(nn.Module):
    """
    C3D Network for Video Classification
    
    RSK World (https://rskworld.in)
    """
    
    def __init__(self, num_classes, pretrained=False):
        super(C3D, self).__init__()
        
        self.conv1 = nn.Conv3d(3, 64, kernel_size=(3, 3, 3), padding=(1, 1, 1))
        self.pool1 = nn.MaxPool3d(kernel_size=(1, 2, 2), stride=(1, 2, 2))
        
        self.conv2 = nn.Conv3d(64, 128, kernel_size=(3, 3, 3), padding=(1, 1, 1))
        self.pool2 = nn.MaxPool3d(kernel_size=(2, 2, 2), stride=(2, 2, 2))
        
        self.conv3a = nn.Conv3d(128, 256, kernel_size=(3, 3, 3), padding=(1, 1, 1))
        self.conv3b = nn.Conv3d(256, 256, kernel_size=(3, 3, 3), padding=(1, 1, 1))
        self.pool3 = nn.MaxPool3d(kernel_size=(2, 2, 2), stride=(2, 2, 2))
        
        self.conv4a = nn.Conv3d(256, 512, kernel_size=(3, 3, 3), padding=(1, 1, 1))
        self.conv4b = nn.Conv3d(512, 512, kernel_size=(3, 3, 3), padding=(1, 1, 1))
        self.pool4 = nn.MaxPool3d(kernel_size=(2, 2, 2), stride=(2, 2, 2))
        
        self.conv5a = nn.Conv3d(512, 512, kernel_size=(3, 3, 3), padding=(1, 1, 1))
        self.conv5b = nn.Conv3d(512, 512, kernel_size=(3, 3, 3), padding=(1, 1, 1))
        self.pool5 = nn.MaxPool3d(kernel_size=(2, 2, 2), stride=(2, 2, 2), padding=(0, 1, 1))
        
        self.fc6 = nn.Linear(8192, 4096)
        self.fc7 = nn.Linear(4096, 4096)
        self.fc8 = nn.Linear(4096, num_classes)
        
        self.dropout = nn.Dropout(p=0.5)
        self.relu = nn.ReLU()
        
    def forward(self, x):
        x = self.relu(self.conv1(x))
        x = self.pool1(x)
        
        x = self.relu(self.conv2(x))
        x = self.pool2(x)
        
        x = self.relu(self.conv3a(x))
        x = self.relu(self.conv3b(x))
        x = self.pool3(x)
        
        x = self.relu(self.conv4a(x))
        x = self.relu(self.conv4b(x))
        x = self.pool4(x)
        
        x = self.relu(self.conv5a(x))
        x = self.relu(self.conv5b(x))
        x = self.pool5(x)
        
        x = x.view(x.size(0), -1)
        
        x = self.relu(self.fc6(x))
        x = self.dropout(x)
        x = self.relu(self.fc7(x))
        x = self.dropout(x)
        x = self.fc8(x)
        
        return x


class SimpleR3D(nn.Module):
    """
    Simple 3D ResNet for Action Recognition
    
    RSK World (https://rskworld.in)
    """
    
    def __init__(self, num_classes):
        super(SimpleR3D, self).__init__()
        
        self.conv1 = nn.Conv3d(3, 64, kernel_size=(3, 7, 7), stride=(1, 2, 2), padding=(1, 3, 3))
        self.bn1 = nn.BatchNorm3d(64)
        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.MaxPool3d(kernel_size=(1, 3, 3), stride=(1, 2, 2), padding=(0, 1, 1))
        
        self.layer1 = self._make_layer(64, 64, 2)
        self.layer2 = self._make_layer(64, 128, 2, stride=2)
        self.layer3 = self._make_layer(128, 256, 2, stride=2)
        self.layer4 = self._make_layer(256, 512, 2, stride=2)
        
        self.avgpool = nn.AdaptiveAvgPool3d((1, 1, 1))
        self.fc = nn.Linear(512, num_classes)
        
    def _make_layer(self, in_channels, out_channels, blocks, stride=1):
        layers = []
        
        # First block with potential downsampling
        layers.append(nn.Conv3d(in_channels, out_channels, kernel_size=3, 
                               stride=stride, padding=1))
        layers.append(nn.BatchNorm3d(out_channels))
        layers.append(nn.ReLU(inplace=True))
        
        # Additional blocks
        for _ in range(1, blocks):
            layers.append(nn.Conv3d(out_channels, out_channels, kernel_size=3, padding=1))
            layers.append(nn.BatchNorm3d(out_channels))
            layers.append(nn.ReLU(inplace=True))
        
        return nn.Sequential(*layers)
    
    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)
        
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        
        x = self.avgpool(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        
        return x


def get_model(model_type, num_classes, pretrained=False):
    """Get model by type"""
    if model_type == "c3d":
        return C3D(num_classes, pretrained)
    elif model_type in ["r3d_18", "r2plus1d_18"]:
        return SimpleR3D(num_classes)
    else:
        raise ValueError(f"Unknown model type: {model_type}")


# ==================================================================================
# Training Functions
# ==================================================================================

class Trainer:
    """
    Advanced Model Trainer
    
    RSK World (https://rskworld.in)
    Founder: Molla Samser
    """
    
    def __init__(self, model, train_loader, val_loader, config):
        self.model = model.to(config.DEVICE)
        self.train_loader = train_loader
        self.val_loader = val_loader
        self.config = config
        
        self.criterion = nn.CrossEntropyLoss()
        self.optimizer = optim.AdamW(
            model.parameters(),
            lr=config.LEARNING_RATE,
            weight_decay=config.WEIGHT_DECAY
        )
        self.scheduler = optim.lr_scheduler.CosineAnnealingLR(
            self.optimizer,
            T_max=config.EPOCHS
        )
        
        self.best_acc = 0.0
        self.history = {'train_loss': [], 'train_acc': [], 'val_loss': [], 'val_acc': []}
        
        # Create checkpoint directory
        Path(config.CHECKPOINT_DIR).mkdir(exist_ok=True)
    
    def train_epoch(self):
        """Train for one epoch"""
        self.model.train()
        running_loss = 0.0
        correct = 0
        total = 0
        
        pbar = tqdm(self.train_loader, desc="Training")
        for frames, labels in pbar:
            frames = frames.to(self.config.DEVICE)
            labels = labels.to(self.config.DEVICE)
            
            self.optimizer.zero_grad()
            
            outputs = self.model(frames)
            loss = self.criterion(outputs, labels)
            
            loss.backward()
            self.optimizer.step()
            
            running_loss += loss.item()
            _, predicted = outputs.max(1)
            total += labels.size(0)
            correct += predicted.eq(labels).sum().item()
            
            pbar.set_postfix({
                'loss': f'{running_loss/total:.4f}',
                'acc': f'{100.*correct/total:.2f}%'
            })
        
        return running_loss / len(self.train_loader), 100. * correct / total
    
    def validate(self):
        """Validate the model"""
        self.model.eval()
        running_loss = 0.0
        correct = 0
        total = 0
        
        with torch.no_grad():
            for frames, labels in self.val_loader:
                frames = frames.to(self.config.DEVICE)
                labels = labels.to(self.config.DEVICE)
                
                outputs = self.model(frames)
                loss = self.criterion(outputs, labels)
                
                running_loss += loss.item()
                _, predicted = outputs.max(1)
                total += labels.size(0)
                correct += predicted.eq(labels).sum().item()
        
        return running_loss / len(self.val_loader), 100. * correct / total
    
    def train(self):
        """Full training loop"""
        print("\n" + "=" * 60)
        print("Action Recognition Model Training")
        print("=" * 60)
        print(f"Website: https://rskworld.in")
        print(f"Founder: Molla Samser")
        print(f"Device: {self.config.DEVICE}")
        print(f"Model: {self.config.MODEL_TYPE}")
        print("=" * 60 + "\n")
        
        for epoch in range(1, self.config.EPOCHS + 1):
            print(f"\nEpoch {epoch}/{self.config.EPOCHS}")
            print("-" * 40)
            
            # Train
            train_loss, train_acc = self.train_epoch()
            
            # Validate
            val_loss, val_acc = self.validate()
            
            # Update scheduler
            self.scheduler.step()
            
            # Save history
            self.history['train_loss'].append(train_loss)
            self.history['train_acc'].append(train_acc)
            self.history['val_loss'].append(val_loss)
            self.history['val_acc'].append(val_acc)
            
            print(f"Train Loss: {train_loss:.4f} | Train Acc: {train_acc:.2f}%")
            print(f"Val Loss: {val_loss:.4f} | Val Acc: {val_acc:.2f}%")
            
            # Save best model
            if val_acc > self.best_acc:
                self.best_acc = val_acc
                self.save_checkpoint('best_model.pth')
                print(f"[RSK World] New best model saved! Accuracy: {val_acc:.2f}%")
            
            # Save periodic checkpoint
            if epoch % self.config.SAVE_EVERY == 0:
                self.save_checkpoint(f'checkpoint_epoch_{epoch}.pth')
        
        print("\n" + "=" * 60)
        print("Training Complete!")
        print(f"Best Validation Accuracy: {self.best_acc:.2f}%")
        print("=" * 60)
        
        return self.history
    
    def save_checkpoint(self, filename):
        """Save model checkpoint"""
        checkpoint = {
            'model_state_dict': self.model.state_dict(),
            'optimizer_state_dict': self.optimizer.state_dict(),
            'best_acc': self.best_acc,
            'history': self.history,
            'config': {
                'model_type': self.config.MODEL_TYPE,
                'num_frames': self.config.NUM_FRAMES,
                'frame_size': self.config.FRAME_SIZE
            }
        }
        
        path = Path(self.config.CHECKPOINT_DIR) / filename
        torch.save(checkpoint, path)
        print(f"[RSK World] Checkpoint saved: {path}")


# ==================================================================================
# Main
# ==================================================================================

def main():
    """Main training function"""
    print("=" * 60)
    print("Action Recognition Dataset - Model Training")
    print("=" * 60)
    print("Website: https://rskworld.in")
    print("Founder: Molla Samser")
    print("Designer: Rima Khatun")
    print("Contact: help@rskworld.in | +91 93305 39277")
    print("=" * 60)
    
    config = Config()
    
    # Check if data exists
    if not Path(config.DATA_DIR).exists():
        print(f"\n[ERROR] Data directory not found: {config.DATA_DIR}")
        print("Please run: python download_youtube_videos.py")
        print("Then run: python process_downloaded.py")
        return
    
    # Create datasets
    print("\n[RSK World] Loading datasets...")
    train_dataset = ActionDataset(config.DATA_DIR, split='train', config=config, augment=True)
    val_dataset = ActionDataset(config.DATA_DIR, split='val', config=config, augment=False)
    
    if len(train_dataset) == 0:
        print("[ERROR] No training samples found!")
        return
    
    # Create data loaders
    train_loader = DataLoader(
        train_dataset,
        batch_size=config.BATCH_SIZE,
        shuffle=True,
        num_workers=0,  # Set to 0 for Windows compatibility
        pin_memory=True
    )
    
    val_loader = DataLoader(
        val_dataset,
        batch_size=config.BATCH_SIZE,
        shuffle=False,
        num_workers=0,
        pin_memory=True
    )
    
    # Create model
    num_classes = len(train_dataset.classes)
    print(f"\n[RSK World] Creating model: {config.MODEL_TYPE}")
    print(f"[RSK World] Number of classes: {num_classes}")
    
    model = get_model(config.MODEL_TYPE, num_classes, config.PRETRAINED)
    
    # Count parameters
    total_params = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f"[RSK World] Total parameters: {total_params:,}")
    print(f"[RSK World] Trainable parameters: {trainable_params:,}")
    
    # Create trainer
    trainer = Trainer(model, train_loader, val_loader, config)
    
    # Train
    history = trainer.train()
    
    # Save final model
    trainer.save_checkpoint('final_model.pth')
    
    # Save training history
    with open(Path(config.CHECKPOINT_DIR) / 'history.json', 'w') as f:
        json.dump(history, f, indent=4)
    
    print("\n[RSK World] Training history saved!")
    print("Thank you for using RSK World datasets!")
    print("Visit: https://rskworld.in")


if __name__ == "__main__":
    main()