# Project: Surveillance Video Dataset
# Author: Molla Samser
# Website: https://rskworld.in/
# Contact: help@rskworld.in
# Phone: +91 93305 39277
# Address: Nutanhat, Mongolkote, Purba Burdwan, West Bengal, India, 713147

"""
Machine Learning model training script for surveillance video analysis.
Supports training custom models for person detection, activity recognition, and anomaly detection.
"""

import os
import json
import cv2
import numpy as np
from pathlib import Path
import argparse

try:
    from ultralytics import YOLO
    YOLO_AVAILABLE = True
except ImportError:
    YOLO_AVAILABLE = False
    print("Warning: ultralytics not available. Install with: pip install ultralytics")

try:
    import torch
    import torch.nn as nn
    TORCH_AVAILABLE = True
except ImportError:
    TORCH_AVAILABLE = False
    print("Warning: PyTorch not available. Install with: pip install torch")


class ActivityRecognitionModel:
    """Simple CNN model for activity recognition."""
    
    def __init__(self, num_classes=4):
        self.num_classes = num_classes
        if TORCH_AVAILABLE:
            self.model = self._build_model()
        else:
            self.model = None
    
    def _build_model(self):
        """Build CNN model for activity recognition."""
        model = nn.Sequential(
            nn.Conv2d(3, 32, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Flatten(),
            nn.Linear(128 * 28 * 28, 512),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(512, self.num_classes)
        )
        return model
    
    def train(self, train_data, epochs=10, batch_size=32):
        """Train the activity recognition model."""
        if not TORCH_AVAILABLE:
            print("PyTorch not available. Cannot train model.")
            return
        
        # Training logic here
        print(f"Training activity recognition model for {epochs} epochs...")
        # Implementation would go here
        print("Training completed!")


class AnomalyDetectionModel:
    """Autoencoder-based anomaly detection model."""
    
    def __init__(self, input_dim=224):
        self.input_dim = input_dim
        if TORCH_AVAILABLE:
            self.model = self._build_autoencoder()
        else:
            self.model = None
    
    def _build_autoencoder(self):
        """Build autoencoder for anomaly detection."""
        encoder = nn.Sequential(
            nn.Conv2d(3, 32, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(64, 128, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )
        
        decoder = nn.Sequential(
            nn.ConvTranspose2d(128, 64, 3, stride=2, padding=1, output_padding=1),
            nn.ReLU(),
            nn.ConvTranspose2d(64, 32, 3, stride=2, padding=1, output_padding=1),
            nn.ReLU(),
            nn.ConvTranspose2d(32, 3, 3, stride=2, padding=1, output_padding=1),
            nn.Sigmoid()
        )
        
        return nn.ModuleDict({'encoder': encoder, 'decoder': decoder})
    
    def train(self, normal_frames, epochs=20):
        """Train on normal frames only."""
        if not TORCH_AVAILABLE:
            print("PyTorch not available. Cannot train model.")
            return
        
        print(f"Training anomaly detection model on {len(normal_frames)} normal frames...")
        # Training logic here
        print("Training completed!")


def prepare_training_data(videos_dir='videos', annotations_dir='annotations', frames_dir='frames/extracted_frames'):
    """Prepare training data from videos and annotations."""
    print("Preparing training data...")
    
    # Load annotations
    annotations_file = os.path.join(annotations_dir, 'annotations.json')
    activities = []
    
    if os.path.exists(annotations_file):
        with open(annotations_file, 'r') as f:
            data = json.load(f)
            activities = data.get('activities', [])
    
    # Extract frames with labels
    training_data = []
    
    for video_file in os.listdir(videos_dir):
        if video_file.endswith(('.mp4', '.avi')):
            video_path = os.path.join(videos_dir, video_file)
            cap = cv2.VideoCapture(video_path)
            fps = cap.get(cv2.CAP_PROP_FPS)
            frame_count = 0
            
            while cap.isOpened():
                ret, frame = cap.read()
                if not ret:
                    break
                
                timestamp = frame_count / fps
                
                # Find activity label for this timestamp
                activity_label = None
                for activity in activities:
                    if activity['start_time'] <= timestamp <= activity['end_time']:
                        activity_label = activity['type']
                        break
                
                if activity_label:
                    training_data.append({
                        'frame': frame,
                        'label': activity_label,
                        'timestamp': timestamp,
                        'video': video_file
                    })
                
                frame_count += 1
            
            cap.release()
    
    print(f"Prepared {len(training_data)} training samples")
    return training_data


def train_yolo_custom(video_dir='videos', epochs=50):
    """Train custom YOLO model on surveillance videos."""
    if not YOLO_AVAILABLE:
        print("YOLO not available. Install ultralytics.")
        return
    
    print("Training custom YOLO model for person detection...")
    
    # Load base model
    model = YOLO('yolov8n.pt')
    
    # Prepare dataset (would need YOLO format annotations)
    # For now, fine-tune on existing model
    print(f"Fine-tuning YOLO model for {epochs} epochs...")
    
    # This would require proper YOLO dataset format
    # model.train(data='dataset.yaml', epochs=epochs, imgsz=640)
    
    # Save model
    model_path = 'models/custom_yolo_surveillance.pt'
    os.makedirs('models', exist_ok=True)
    # model.save(model_path)
    
    print(f"Model training completed. Saved to {model_path}")


def train_activity_model(training_data, epochs=20):
    """Train activity recognition model."""
    print("Training activity recognition model...")
    
    model = ActivityRecognitionModel(num_classes=4)  # walking, standing, running, other
    
    if model.model:
        model.train(training_data, epochs=epochs)
        
        # Save model
        model_path = 'models/activity_recognition.pt'
        os.makedirs('models', exist_ok=True)
        # torch.save(model.model.state_dict(), model_path)
        print(f"Activity recognition model saved to {model_path}")
    else:
        print("Cannot train: PyTorch not available")


def train_anomaly_model(normal_frames, epochs=30):
    """Train anomaly detection model."""
    print("Training anomaly detection model...")
    
    model = AnomalyDetectionModel()
    
    if model.model:
        model.train(normal_frames, epochs=epochs)
        
        # Save model
        model_path = 'models/anomaly_detection.pt'
        os.makedirs('models', exist_ok=True)
        # torch.save(model.model.state_dict(), model_path)
        print(f"Anomaly detection model saved to {model_path}")
    else:
        print("Cannot train: PyTorch not available")


def main():
    parser = argparse.ArgumentParser(description='Train ML models for surveillance video analysis')
    parser.add_argument('--model', '-m', choices=['yolo', 'activity', 'anomaly', 'all'],
                       default='all', help='Model to train')
    parser.add_argument('--epochs', '-e', type=int, default=20, help='Number of epochs')
    parser.add_argument('--data-dir', '-d', default='videos', help='Video data directory')
    
    args = parser.parse_args()
    
    print("=" * 60)
    print("ML Model Training for Surveillance Video Dataset")
    print("=" * 60)
    
    if args.model in ['yolo', 'all']:
        train_yolo_custom(args.data_dir, epochs=args.epochs)
    
    if args.model in ['activity', 'all']:
        training_data = prepare_training_data()
        if training_data:
            train_activity_model(training_data, epochs=args.epochs)
    
    if args.model in ['anomaly', 'all']:
        # Extract normal frames (non-anomaly)
        normal_frames = []  # Would extract from videos excluding anomaly timestamps
        if normal_frames:
            train_anomaly_model(normal_frames, epochs=args.epochs)
    
    print("\n" + "=" * 60)
    print("Training completed!")
    print("=" * 60)


if __name__ == '__main__':
    main()