"""
General utilities shared across all courseworks.

Provides:
    - set_seed(): Set random seeds for reproducibility
    - Timer: Context manager for timing code blocks
    - EarlyStopping: Stop training when validation loss stops improving
    - save_results() / load_results(): Save/load training history as JSON
    - print_config(): Pretty-print configuration
"""

import os
import json
import time
import numpy as np


def set_seed(seed):
    """
    Set random seed for reproducibility (numpy + torch if available).

    Args:
        seed: Integer seed value
    """
    np.random.seed(seed)
    try:
        import torch
        torch.manual_seed(seed)
        if torch.cuda.is_available():
            torch.cuda.manual_seed_all(seed)
    except ImportError:
        pass  # PyTorch not required for CW2/CW3


class Timer:
    """
    Context manager for timing code blocks.

    Usage:
        with Timer("Training"):
            train()
        # Prints: "Training completed in 12.34s"
    """

    def __init__(self, name="Block"):
        self.name = name

    def __enter__(self):
        self.start = time.time()
        return self

    def __exit__(self, *args):
        elapsed = time.time() - self.start
        if elapsed >= 60:
            minutes = int(elapsed // 60)
            seconds = elapsed % 60
            print(f"{self.name} completed in {minutes}m {seconds:.1f}s")
        else:
            print(f"{self.name} completed in {elapsed:.2f}s")


class EarlyStopping:
    """
    Stop training when validation loss stops improving.

    Usage:
        early_stop = EarlyStopping(patience=5)
        for epoch in range(num_epochs):
            val_loss = validate()
            if early_stop(val_loss):
                print("Early stopping triggered!")
                break
    """

    def __init__(self, patience=5, min_delta=1e-4):
        """
        Args:
            patience: Number of epochs to wait after last improvement
            min_delta: Minimum change to qualify as an improvement
        """
        self.patience = patience
        self.min_delta = min_delta
        self.best_loss = float('inf')
        self.counter = 0

    def __call__(self, val_loss):
        """
        Check if training should stop.

        Args:
            val_loss: Current epoch's validation loss

        Returns:
            True if training should stop, False otherwise
        """
        if val_loss < self.best_loss - self.min_delta:
            self.best_loss = val_loss
            self.counter = 0
            return False
        else:
            self.counter += 1
            return self.counter >= self.patience


def save_results(results_dict, path):
    """
    Save training results to a JSON file.

    Args:
        results_dict: Dictionary of results (lists, scalars, strings)
        path: File path to save to (e.g., "outputs/results.json")
    """
    os.makedirs(os.path.dirname(path), exist_ok=True) if os.path.dirname(path) else None

    # Convert numpy types to Python types for JSON serialization
    def convert(obj):
        if isinstance(obj, np.integer):
            return int(obj)
        elif isinstance(obj, np.floating):
            return float(obj)
        elif isinstance(obj, np.ndarray):
            return obj.tolist()
        return obj

    serializable = {}
    for key, val in results_dict.items():
        if isinstance(val, list):
            serializable[key] = [convert(v) for v in val]
        else:
            serializable[key] = convert(val)

    with open(path, 'w') as f:
        json.dump(serializable, f, indent=2)
    print(f"  Results saved to: {path}")


def load_results(path):
    """
    Load training results from a JSON file.

    Args:
        path: File path to load from

    Returns:
        Dictionary of results
    """
    with open(path, 'r') as f:
        return json.load(f)


def print_config(config):
    """
    Pretty-print a configuration (dataclass or dict).

    Args:
        config: A dataclass instance or dictionary
    """
    print("\n=== Configuration ===")
    if hasattr(config, '__dataclass_fields__'):
        for field_name in config.__dataclass_fields__:
            print(f"  {field_name}: {getattr(config, field_name)}")
    elif isinstance(config, dict):
        for key, val in config.items():
            print(f"  {key}: {val}")
    print("=" * 24)
