#!/usr/bin/env python
#
# file: transfer_learning_demo.py
#
# description:
#  This script provides an educational demonstration of Transfer Learning 
#  using PyTorch. It loads a pre-trained ResNet18 model, freezes its early 
#  layers, and trains a newly added final classification layer on a synthetic 
#  image dataset. It prints train/test accuracies and outputs plots 
#  for both the training metrics and sample visual predictions.
#
# revision history:
#  20260416 (AM): initial version
#  20260416 (AM): added train/test splits, test evaluation, and prediction plots
#------------------------------------------------------------------------------

# import system modules
#
import os
import sys
import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
import torchvision.models as models

#------------------------------------------------------------------------------
#
# global variables are listed here
#
#------------------------------------------------------------------------------

# set the filename using basename
#
__FILE__ = os.path.basename(__file__)

# define default values for the dataset, model, and plotting
#
DEF_NUM_TRAIN     = 100
DEF_NUM_TEST      = 20
DEF_BATCH_SIZE    = 10
DEF_NUM_EPOCHS    = 5
DEF_LEARNING_RATE = 0.01
DEF_RANDOM_SEED   = 42
DEF_METRICS_FILE  = "transfer_learning_metrics.png"
DEF_PRED_FILE     = "transfer_learning_predictions.png"
DEF_PLOT_TITLE    = "Transfer Learning: Training Progression"

#------------------------------------------------------------------------------
#
# functions are listed here
#
#------------------------------------------------------------------------------

# set the random seeds for reproducibility
#
def set_seed(seed=DEF_RANDOM_SEED):
    """
    method: set_seed
    
    arguments:
     seed: integer value for the random seed
     
    return:
     none
     
    description:
     Locks the random seed for NumPy and PyTorch to ensure the training 
     results and synthetic data generation are fully reproducible.
    """
    
    np.random.seed(seed)
    torch.manual_seed(seed)
    
    # exit gracefully
    #
    return

# generate synthetic image data for training and testing
#
def generate_synthetic_data(n_train=DEF_NUM_TRAIN, n_test=DEF_NUM_TEST):
    """
    method: generate_synthetic_data
    
    arguments:
     n_train: number of training images to generate per class
     n_test: number of testing images to generate per class
     
    return:
     train_loader: PyTorch DataLoader for the training set
     test_loader: PyTorch DataLoader for the testing set
     
    description:
     Generates fake 3-channel 224x224 "images" for train and test sets.
     Class 0: dark background (~0.1) with a bright horizontal stripe
              spanning the middle third of the image height.
     Class 1: bright background (~0.9) with a dark vertical stripe
              spanning the middle third of the image width.
     Both classes are visually distinct even after normalization to [0,1].
    """

    def _create_dataset(n_samples):
        # class 0: dark background with bright horizontal band
        #
        X_class0 = torch.ones(n_samples, 3, 224, 224) * 0.1 \
                   + torch.randn(n_samples, 3, 224, 224) * 0.05
        X_class0[:, :, 74:150, :] = 0.9 \
                   + torch.randn(n_samples, 3, 76, 224) * 0.05
        y_class0 = torch.zeros(n_samples, dtype=torch.long)

        # class 1: bright background with dark vertical band
        #
        X_class1 = torch.ones(n_samples, 3, 224, 224) * 0.9 \
                   + torch.randn(n_samples, 3, 224, 224) * 0.05
        X_class1[:, :, :, 74:150] = 0.1 \
                   + torch.randn(n_samples, 3, 224, 76) * 0.05
        y_class1 = torch.ones(n_samples, dtype=torch.long)

        X = torch.cat([X_class0, X_class1], dim=0)
        y = torch.cat([y_class0, y_class1], dim=0)

        return TensorDataset(X, y)

    train_dataset = _create_dataset(n_train)
    test_dataset  = _create_dataset(n_test)

    train_loader = DataLoader(train_dataset, batch_size=DEF_BATCH_SIZE, shuffle=True)
    test_loader  = DataLoader(test_dataset,  batch_size=DEF_BATCH_SIZE, shuffle=False)

    # exit gracefully
    #
    return train_loader, test_loader

# configure the pre-trained model for transfer learning
#
def setup_transfer_learning_model():
    """
    method: setup_transfer_learning_model
    
    arguments:
     none
     
    return:
     model: modified PyTorch model ready for transfer learning
     
    description:
     Loads a pre-trained ResNet18, freezes all internal convolutional 
     layers so their weights do not update, and replaces the final 
     fully connected (fc) layer to predict exactly 2 classes.
    """
    
    print("Loading pre-trained ResNet18 model...")
    
    weights = models.ResNet18_Weights.DEFAULT
    model = models.resnet18(weights=weights)
    
    # freeze all early layers in the network
    #
    for param in model.parameters():
        param.requires_grad = False
        
    # replace the final layer (which automatically has requires_grad=True)
    #
    num_ftrs = model.fc.in_features
    model.fc = nn.Linear(num_ftrs, 2)
    
    # exit gracefully
    #
    return model

# train and evaluate the modified model
#
def train_model(model, train_loader, test_loader, epochs=DEF_NUM_EPOCHS, lr=DEF_LEARNING_RATE):
    """
    method: train_model
    
    arguments:
     model: the PyTorch model to train
     train_loader: dataset loader for training data
     test_loader: dataset loader for testing data
     epochs: number of times to iterate over the dataset
     lr: learning rate for the optimizer
     
    return:
     metrics: dictionary containing lists of train/test losses and accuracies
     
    description:
     Executes the training loop and evaluates on the test set every epoch.
     Prints accuracy for both sets and tracks metrics over time.
    """
    
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=lr)
    
    metrics = {
        'train_loss': [], 'train_acc': [],
        'test_loss': [], 'test_acc': []
    }
    
    print("Beginning training loop for %d epochs..." % epochs)
    
    for epoch in range(epochs):
        
        # --- TRAINING PHASE ---
        #
        model.train()
        running_loss = 0.0
        correct_preds = 0
        total_preds = 0
        
        for inputs, labels in train_loader:
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            
            running_loss += loss.item() * inputs.size(0)
            _, predictions = torch.max(outputs, 1)
            correct_preds += torch.sum(predictions == labels.data).item()
            total_preds += labels.size(0)
            
        train_loss = running_loss / total_preds
        train_acc = (correct_preds / total_preds) * 100.0
        
        # --- TESTING PHASE ---
        #
        model.eval()
        test_loss_val = 0.0
        test_correct = 0
        test_total = 0
        
        with torch.no_grad():
            for inputs, labels in test_loader:
                outputs = model(inputs)
                loss = criterion(outputs, labels)
                
                test_loss_val += loss.item() * inputs.size(0)
                _, predictions = torch.max(outputs, 1)
                test_correct += torch.sum(predictions == labels.data).item()
                test_total += labels.size(0)
                
        test_loss = test_loss_val / test_total
        test_acc = (test_correct / test_total) * 100.0
        
        # store metrics
        #
        metrics['train_loss'].append(train_loss)
        metrics['train_acc'].append(train_acc)
        metrics['test_loss'].append(test_loss)
        metrics['test_acc'].append(test_acc)
        
        print(" Epoch %d/%d | Train Acc: %5.2f%% (Loss: %.4f) | Test Acc: %5.2f%% (Loss: %.4f)" % 
              (epoch + 1, epochs, train_acc, train_loss, test_acc, test_loss))
              
    # exit gracefully
    #
    return metrics

# plot the training metrics
#
def plot_training_results(metrics, outfile=DEF_METRICS_FILE):
    """
    method: plot_training_results
    
    arguments:
     metrics: dict containing tracked loss and accuracy values
     outfile: path to save the resulting image
     
    return:
     status: boolean indicating success
     
    description:
     Generates a 2-panel side-by-side plot comparing train and test 
     loss and accuracy over the epochs.
    """
    
    epochs_range = range(1, len(metrics['train_loss']) + 1)
    
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))
    
    # --- PANEL 1: LOSS ---
    #
    ax1.plot(epochs_range, metrics['train_loss'], marker='o', color='crimson', label='Train Loss')
    ax1.plot(epochs_range, metrics['test_loss'], marker='x', color='darkred', linestyle='--', label='Test Loss')
    ax1.set_title("Cross-Entropy Loss", fontsize=14)
    ax1.set_xlabel("Epoch")
    ax1.set_ylabel("Loss")
    ax1.legend()
    ax1.grid(True, linestyle='--', alpha=0.6)
    
    # --- PANEL 2: ACCURACY ---
    #
    ax2.plot(epochs_range, metrics['train_acc'], marker='s', color='seagreen', label='Train Acc')
    ax2.plot(epochs_range, metrics['test_acc'], marker='^', color='darkgreen', linestyle='--', label='Test Acc')
    ax2.set_title("Classification Accuracy", fontsize=14)
    ax2.set_xlabel("Epoch")
    ax2.set_ylabel("Accuracy (%)")
    ax2.legend()
    ax2.grid(True, linestyle='--', alpha=0.6)
    
    fig.suptitle(DEF_PLOT_TITLE, fontsize=16, y=1.02)
    plt.tight_layout()
    
    try:
        plt.savefig(outfile, dpi=150, bbox_inches='tight')
        print("\nSaved metrics visualization to: %s" % outfile)
    except Exception as e:
        print("**> Error saving metrics plot: %s" % str(e))
        return False
        
    plt.show()
    
    # exit gracefully
    #
    return True

# plot sample predictions from the test set
#
# plot balanced sample predictions from the test set
#
def plot_predictions(model, test_loader, num_per_class=2, outfile=DEF_PRED_FILE):
    """
    method: plot_predictions
    
    arguments:
     model: the trained PyTorch model
     test_loader: dataset loader for testing data
     num_per_class: number of images to plot per class (default is 2)
     outfile: path to save the resulting image
     
    return:
     status: boolean indicating success
     
    description:
     Searches the test set to extract a specific number of images from 
     each class (ensuring a balanced display), runs them through the model, 
     and visualizes the synthetic images alongside their true and predicted labels.
    """
    
    model.eval()
    
    # initialize lists to hold our selected samples
    #
    selected_inputs = []
    selected_labels = []
    
    # track how many we have found for each class
    # (assuming classes are 0 and 1 based on previous code)
    #
    class_counts = {0: 0, 1: 0}
    
    # iterate through the test loader until we have enough of each class
    #
    with torch.no_grad():
        for inputs, labels in test_loader:
            for i in range(inputs.size(0)):
                label_val = labels[i].item()
                
                # if we need more of this specific class, grab it
                #
                if label_val in class_counts and class_counts[label_val] < num_per_class:
                    selected_inputs.append(inputs[i])
                    selected_labels.append(labels[i])
                    class_counts[label_val] += 1
                    
            # check if we have reached our target for both classes
            #
            if class_counts[0] == num_per_class and class_counts[1] == num_per_class:
                break
                
    # fallback check in case the dataset didn't have enough images
    #
    total_images = len(selected_inputs)
    if total_images == 0:
        print("**> Error: No images found to plot.")
        return False
        
    # stack the lists into proper tensors for the model
    #
    inputs_tensor = torch.stack(selected_inputs)
    labels_tensor = torch.stack(selected_labels)
    
    # get predictions
    #
    with torch.no_grad():
        outputs = model(inputs_tensor)
        _, preds = torch.max(outputs, 1)
        
    # setup plot
    #
    fig, axes = plt.subplots(1, total_images, figsize=(15, 4))
    if total_images == 1:
        axes = [axes]
        
    for i in range(total_images):
        # convert tensor to numpy and reorder channels for matplotlib (C,H,W) -> (H,W,C)
        #
        img = inputs_tensor[i].numpy().transpose((1, 2, 0))
        
        # normalize to [0, 1] for visualization purposes
        #
        img = (img - img.min()) / (img.max() - img.min())
        
        true_label = labels_tensor[i].item()
        pred_label = preds[i].item()
        
        color = 'green' if true_label == pred_label else 'red'
        
        axes[i].imshow(img)
        axes[i].axis('off')
        axes[i].set_title(f"GT: Class {true_label} | Pred: Class {pred_label}", color=color, fontweight='bold')
        
    fig.suptitle("Sample Predictions on Test Set (Balanced)", fontsize=16, y=1.05)
    plt.tight_layout()
    
    try:
        plt.savefig(outfile, dpi=150, bbox_inches='tight')
        print("Saved predictions visualization to: %s" % outfile)
    except Exception as e:
        print("**> Error saving predictions plot: %s" % str(e))
        return False
        
    plt.show()
    
    # exit gracefully
    #
    return True

# function: main
#
def main(argv):
    
    print("--- Starting Transfer Learning Demonstration ---")
    
    # 1. initialize setup
    #
    set_seed()
    
    # 2. prepare data
    #
    print("Generating synthetic train and test data (224x224 tensors)...")
    train_loader, test_loader = generate_synthetic_data()
    
    # 3. prepare transfer learning model
    #
    model = setup_transfer_learning_model()
    
    # 4. train only the final layer and evaluate
    #
    metrics = train_model(model, train_loader, test_loader)
    
    # 5. visualize the learning process and predictions
    #
    print("\nPlotting results...")
    status_metrics = plot_training_results(metrics)
    status_preds = plot_predictions(model, test_loader)
    
    if not (status_metrics and status_preds):
        print("**> Process failed during plotting.")
        return False
        
    print("--- Demonstration Complete ---")
    
    # exit gracefully
    #
    return True

# begin gracefully
#
if __name__ == '__main__':
    main(sys.argv[0:])

#
# end of file