main.py

"""
Multi-Source Transfer Learning for Time Series Classification

This script implements multi-source transfer learning using shapelet-based similarity
metrics for source dataset selection. It supports pre-training on multiple source datasets
followed by fine-tuning on target datasets.

Author: Time Series Transferability Research
"""

import os
import argparse
from typing import Tuple, List, Optional

import numpy as np
import pandas as pd
import tensorflow as tf
from tensorflow import keras

from utils.constants import UCR_list
from utils.transfer_learning import pre_train_multi_source, fine_tuning


# ============================================================================
# Source Dataset Selection
# ============================================================================

def select_source_datasets(target: str, 
                          metric: str, 
                          num_sources: int) -> np.ndarray:
    """
    Select source datasets based on similarity metric.
    
    Reads the precomputed similarity metric scores and selects the top-k
    most similar datasets to the target dataset.
    
    Args:
        target: Name of the target dataset
        metric: Name of the similarity metric (e.g., 'Minimum_Shapelet')
        num_sources: Number of source datasets to select
        
    Returns:
        Array of selected source dataset names
        
    Raises:
        FileNotFoundError: If metric score file is not found
        ValueError: If target dataset not in UCR_list
    """
    if target not in UCR_list:
        raise ValueError(f"Target dataset '{target}' not found in UCR_list")
    
    # Load similarity scores
    score_file = f"score/{metric}.npy"
    if not os.path.exists(score_file):
        raise FileNotFoundError(f"Metric score file not found: {score_file}")
    
    score_matrix = np.load(score_file)
    score_matrix = score_matrix.astype(np.float64).T
    
    # Get scores for target dataset
    target_index = UCR_list.index(target)
    target_scores = score_matrix[target_index]

    # Sort datasets by similarity (ascending order - lower distance = more similar)
    sorted_indices = np.argsort(target_scores)
    ranked_indices = sorted_indices.astype(int)[:num_sources + 1]
    
    source_list = np.array(UCR_list)[ranked_indices]
    
    # Remove target dataset if it appears in the source list
    if target in source_list:
        source_list = np.delete(source_list, np.where(source_list == target))
    else:
        # Remove the least similar dataset
        source_list = np.delete(source_list, -1)
    
    return source_list


# ============================================================================
# Transfer Learning
# ============================================================================

def run_transfer_learning(target: str,
                         metric: str,
                         num_sources: int,
                         model_arch: str = 'vgg',
                         pre_train_iterations: int = 10000,
                         fine_tune_iterations: int = 5000,
                         save_pre_trained: bool = False,
                         save_fine_tuned: bool = False,
                         fine_tuning_only: bool = False) -> Tuple[float, float, Optional[np.ndarray]]:
    """
    Run complete transfer learning pipeline for a target dataset.
    
    Args:
        target: Target dataset name
        metric: Similarity metric for source selection
        num_sources: Number of source datasets for pre-training
        model_arch: Model architecture (default: 'vgg')
        pre_train_iterations: Number of pre-training iterations
        fine_tune_iterations: Number of fine-tuning iterations
        save_pre_trained: Whether to save pre-trained model
        save_fine_tuned: Whether to save fine-tuned model
        fine_tuning_only: If True, load pre-trained model and only do fine-tuning
        
    Returns:
        Tuple of (loss, accuracy, source_list)
    """
    print(f"\n{'='*80}")
    print(f"Target Dataset: {target}")
    print(f"{'='*80}")
    
    # Fine-tuning only mode
    if fine_tuning_only:
        model_path = f"model_save/{metric}/{num_sources}/{target}.h5"
        
        if not os.path.isfile(model_path):
            print(f"Error: No pre-trained model found at {model_path}")
            return None, None, None
        
        print(f"Loading pre-trained model from: {model_path}")
        print(f"Model size: {os.stat(model_path).st_size / (1024*1024):.2f} MB")
        
        pre_trained_model = keras.models.load_model(model_path)
        loss, acc = fine_tuning(
            pre_trained_model=pre_trained_model,
            target=target,
            nb_iterations=fine_tune_iterations
        )
        
        return loss, acc, None
    
    # Full pipeline: source selection + pre-training + fine-tuning
    else:
        # Select source datasets
        print(f"\nSelecting source datasets using {metric} metric...")
        source_list = select_source_datasets(target, metric, num_sources)
        print(f"Selected sources ({len(source_list)}): {', '.join(source_list)}")
        
        # Create directory for saving models
        model_dir = f"model_save/{metric}/{num_sources}"
        os.makedirs(model_dir, exist_ok=True)
        
        model_path = f"{model_dir}/{target}.h5"
        
        # Pre-training on multi-source datasets
        print(f"\n--- Phase 1: Multi-Source Pre-Training ---")
        print(f"Sources: {len(source_list)} datasets")
        print(f"Iterations: {pre_train_iterations}")
        
        pre_trained_model = pre_train_multi_source(
            source_list=source_list,
            target=target,
            model_architecture=model_arch,
            dataset_balancing=True,
            save_model=save_pre_trained,
            metric=metric,
            nb_iterations=pre_train_iterations,
            save_path=model_path
        )
        
        # Fine-tuning on target dataset
        print(f"\n--- Phase 2: Fine-Tuning on Target ---")
        print(f"Target: {target}")
        print(f"Iterations: {fine_tune_iterations}")
        
        loss, acc = fine_tuning(
            pre_trained_model=pre_trained_model,
            target=target,
            nb_iterations=fine_tune_iterations
        )
        
        print(f"\n{'='*80}")
        print(f"Results for {target}:")
        print(f"  Loss: {loss:.4f}")
        print(f"  Accuracy: {acc:.4f}")
        print(f"{'='*80}")
        
        return loss, acc, source_list


def run_experiment(metric: str,
                  num_sources: int,
                  model_arch: str = 'vgg',
                  pre_train_iterations: int = 10000,
                  fine_tune_iterations: int = 5000,
                  save_pre_trained: bool = False,
                  save_fine_tuned: bool = False) -> pd.DataFrame:
    """
    Run transfer learning experiment on all UCR datasets.
    
    Args:
        metric: Similarity metric for source selection
        num_sources: Number of source datasets for pre-training
        model_arch: Model architecture
        pre_train_iterations: Number of pre-training iterations
        fine_tune_iterations: Number of fine-tuning iterations
        save_pre_trained: Whether to save pre-trained models
        save_fine_tuned: Whether to save fine-tuned models
        
    Returns:
        DataFrame with results for all datasets
    """
    print(f"\n{'#'*80}")
    print(f"# Running Full Experiment on {len(UCR_list)} UCR Datasets")
    print(f"# Metric: {metric}")
    print(f"# Number of Sources: {num_sources}")
    print(f"{'#'*80}\n")
    
    results = {
        'dataset': [],
        'source_list': [],
        'accuracy': [],
        'loss': []
    }
    
    for idx, target in enumerate(UCR_list, 1):
        print(f"\nProgress: [{idx}/{len(UCR_list)}]")
        
        try:
            loss, acc, source_list = run_transfer_learning(
                target=target,
                metric=metric,
                num_sources=num_sources,
                model_arch=model_arch,
                pre_train_iterations=pre_train_iterations,
                fine_tune_iterations=fine_tune_iterations,
                save_pre_trained=save_pre_trained,
                save_fine_tuned=save_fine_tuned
            )
            
            results['dataset'].append(target)
            results['source_list'].append(source_list.tolist() if source_list is not None else [])
            results['accuracy'].append(acc)
            results['loss'].append(loss)
            
        except Exception as e:
            print(f"Error processing {target}: {e}")
            results['dataset'].append(target)
            results['source_list'].append([])
            results['accuracy'].append(None)
            results['loss'].append(None)
    
    # Create results DataFrame
    df_results = pd.DataFrame(results)
    
    # Save results
    os.makedirs("result", exist_ok=True)
    output_file = f"result/{metric}_{num_sources}_result.csv"
    df_results.to_csv(output_file, index=False)
    
    print(f"\n{'#'*80}")
    print(f"# Experiment Complete!")
    print(f"# Results saved to: {output_file}")
    print(f"{'#'*80}\n")
    
    # Print summary statistics
    valid_accuracies = df_results['accuracy'].dropna()
    if len(valid_accuracies) > 0:
        print(f"Summary Statistics:")
        print(f"  Mean Accuracy: {valid_accuracies.mean():.4f}")
        print(f"  Std Accuracy: {valid_accuracies.std():.4f}")
        print(f"  Min Accuracy: {valid_accuracies.min():.4f}")
        print(f"  Max Accuracy: {valid_accuracies.max():.4f}")
    
    return df_results


# ============================================================================
# Main Execution
# ============================================================================

def main():
    """Main execution function."""
    parser = argparse.ArgumentParser(
        description='Multi-Source Transfer Learning for Time Series Classification',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Run transfer learning on a single dataset
  python main.py --target Coffee --metric Minimum_Shapelet --dataset-number 10
  
  # Run full experiment on all UCR datasets
  python main.py --target experiment --metric Minimum_Shapelet --dataset-number 10
  
  # Fine-tuning only from saved pre-trained model
  python main.py --target Coffee --metric Minimum_Shapelet --dataset-number 10 --fine-tuning-only
        """
    )
    
    # GPU settings
    parser.add_argument(
        '--gpus',
        type=str,
        default="",
        help="GPU device IDs to use (e.g., '0,1' for GPUs 0 and 1)"
    )
    
    # Target dataset
    parser.add_argument(
        '--target', '-t',
        type=str,
        required=True,
        help="Target dataset name (use 'experiment' to run on all UCR datasets)"
    )
    
    # Model settings
    parser.add_argument(
        '--model',
        type=str,
        default='vgg',
        help="Model architecture (default: vgg)"
    )
    
    # Training iterations
    parser.add_argument(
        '--pre-iteration', '-pi',
        type=int,
        default=10000,
        help='Number of iterations for pre-training (default: 10000)'
    )
    
    parser.add_argument(
        '--transfer-iteration', '-ti',
        type=int,
        default=5000,
        help='Number of iterations for fine-tuning (default: 5000)'
    )
    
    # Source selection
    parser.add_argument(
        '--dataset-number', '-dn',
        type=int,
        required=True,
        help='Number of source datasets for multi-source pre-training'
    )
    
    parser.add_argument(
        '--metric',
        type=str,
        default='Minimum_Shapelet',
        help='Similarity metric for source dataset selection (default: Minimum_Shapelet)'
    )
    
    # Model saving options
    parser.add_argument(
        '--save-pre-trained-model', '-sv-pm',
        action='store_true',
        help='Save pre-trained model to disk'
    )
    
    parser.add_argument(
        '--save-fine-tuned-model', '-sv-fm',
        action='store_true',
        help='Save fine-tuned model to disk'
    )
    
    # Training mode
    parser.add_argument(
        '--fine-tuning-only',
        action='store_true',
        help='Only perform fine-tuning using saved pre-trained model'
    )
    
    args = parser.parse_args()
    
    # Set GPU availability
    if args.gpus:
        os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus
        print(f"Using GPU(s): {args.gpus}")
    
    # Run experiment on all datasets or single target
    if args.target == "experiment":
        run_experiment(
            metric=args.metric,
            num_sources=args.dataset_number,
            model_arch=args.model,
            pre_train_iterations=args.pre_iteration,
            fine_tune_iterations=args.transfer_iteration,
            save_pre_trained=args.save_pre_trained_model,
            save_fine_tuned=args.save_fine_tuned_model
        )
    else:
        # Run on single target dataset
        loss, acc, source_list = run_transfer_learning(
            target=args.target,
            metric=args.metric,
            num_sources=args.dataset_number,
            model_arch=args.model,
            pre_train_iterations=args.pre_iteration,
            fine_tune_iterations=args.transfer_iteration,
            save_pre_trained=args.save_pre_trained_model,
            save_fine_tuned=args.save_fine_tuned_model,
            fine_tuning_only=args.fine_tuning_only
        )
        
        if loss is not None and acc is not None:
            print("\n" + "="*80)
            print("Transfer Learning Complete!")
            print(f"  Loss: {loss:.4f}")
            print(f"  Accuracy: {acc:.4f}")
            if source_list is not None:
                print(f"  Source Datasets: {', '.join(source_list)}")
            print("="*80)
        else:
            print("\nTransfer learning failed. Please check the error messages above.")


if __name__ == "__main__":
    main()