#!/usr/bin/env python3
# utils.py - Common utilities for XSimGCL project
"""
Shared helper functions for data conversion, hardware setup,
configuration management, checkpoint handling.

FIXED: resolve_config_path now checks script directory FIRST before going up
ADDED: load_active_config() — czyta train_dataset.yaml i zwraca (config_path, dataset_type)
"""

import torch
import numpy as np
import scipy.sparse as sp
import yaml
import os
import sys
from typing import Dict, Optional, List, Tuple


# ============================================================================
# SECTION 1: Data Conversion
# ============================================================================

def scipy_to_torch_sparse(mat):
    """
    Convert scipy sparse matrix to torch sparse tensor.

    Używa formatu CSR (Compressed Sparse Row) jeśli PyTorch >= 1.11.
    CSR jest cache-friendly dla torch.sparse.mm na CUDA — szybszy forward pass GCN.
    Fallback do COO dla starszych wersji PyTorch.

    Args:
        mat: Scipy sparse matrix or dense array

    Returns:
        torch.sparse.Tensor (CSR preferred) or original input if not sparse
    """
    if not sp.issparse(mat):
        return mat

    csr = mat.tocsr()
    crow_indices = torch.from_numpy(csr.indptr.astype(np.int64))
    col_indices  = torch.from_numpy(csr.indices.astype(np.int64))
    values       = torch.from_numpy(csr.data.astype(np.float32))
    try:
        return torch.sparse_csr_tensor(
            crow_indices, col_indices, values, torch.Size(csr.shape)
        )
    except Exception:
        # Fallback: COO dla PyTorch < 1.11
        coo = mat.tocoo()
        indices = torch.from_numpy(np.vstack((coo.row, coo.col)).astype(np.int64))
        values_coo = torch.from_numpy(coo.data.astype(np.float32))
        return torch.sparse_coo_tensor(indices, values_coo, torch.Size(coo.shape)).coalesce()


# ============================================================================
# SECTION 2: Hardware Setup
# ============================================================================

def setup_device(device_arg: str) -> str:
    """
    Auto-detect or select compute device.

    Args:
        device_arg: 'auto', 'cuda', or 'cpu'

    Returns:
        Selected device string ('cuda' or 'cpu')
    """
    if device_arg == 'auto':
        if torch.cuda.is_available():
            print(f"[Device] CUDA: {torch.cuda.get_device_name(0)}")
            return 'cuda'
        print("[Device] Using CPU")
        return 'cpu'

    if device_arg == 'cuda' and not torch.cuda.is_available():
        print("[Device] CUDA requested but not available → falling back to CPU")
        return 'cpu'

    return device_arg


# ============================================================================
# SECTION 3: Config Management
# ============================================================================

def load_yaml_config(path: str, required: bool = True) -> Optional[Dict]:
    """
    Load YAML configuration file.

    Args:
        path: Path to YAML file
        required: If True, exit on missing file

    Returns:
        Configuration dictionary or None
    """
    if not os.path.exists(path):
        if required:
            print(f"[ERROR] Missing configuration file: {path}")
            sys.exit(1)
        return None

    try:
        with open(path, 'r') as f:
            return yaml.safe_load(f)
    except Exception as e:
        print(f"[ERROR] Failed to load YAML: {e}")
        if required:
            sys.exit(1)
        return None


def load_active_config(script_dir: Optional[str] = None) -> Tuple[str, str]:
    """
    Czyta train_dataset.yaml i zwraca (config_path, dataset_type).

    Pozwala wszystkim skryptom (train_model, png, test_recommend, grid_train)
    automatycznie wiedzieć którego datasetu i której konfiguracji używać
    bez ręcznego podawania argumentów.

    train_dataset.yaml:
      dataset: "f95"    # lub "steam"

    Returns:
        (config_path, dataset_type) gdzie dataset_type in {'f95', 'steam'}
    """
    if script_dir is None:
        script_dir = get_script_directory()

    td_path = os.path.join(script_dir, 'train_dataset.yaml')
    if not os.path.exists(td_path):
        # Fallback: szukaj w katalogu roboczym
        td_path = os.path.join(os.getcwd(), 'train_dataset.yaml')

    td = load_yaml_config(td_path, required=False)
    if not td:
        # Brak pliku → domyślnie f95
        print(f"   [load_active_config] train_dataset.yaml not found, defaulting to f95")
        config_path = os.path.join(script_dir, 'config.yaml')
        return config_path, 'f95'

    dataset_type = str(td.get('dataset', 'f95')).lower().strip()

    if dataset_type == 'steam':
        config_path = os.path.join(script_dir, 'config_steam.yaml')
    else:
        dataset_type = 'f95'
        config_path = os.path.join(script_dir, 'config.yaml')

    print(f"   [Dataset] train_dataset.yaml → dataset={dataset_type}, config={os.path.basename(config_path)}")
    return config_path, dataset_type


def load_dataset_class(dataset_type: str):
    """
    Zwraca klasę datasetu na podstawie dataset_type.

    Args:
        dataset_type: 'f95' lub 'steam'

    Returns:
        Klasa datasetu (F95Dataset lub SteamDataset)
    """
    if dataset_type == 'steam':
        from game_dataset_steam import SteamDataset
        return SteamDataset
    else:
        from game_dataset import F95Dataset
        return F95Dataset


# ============================================================================
# SECTION 4: Checkpoint Management
# ============================================================================

def find_checkpoint(base_dir: str, checkpoint_patterns: List[str]) -> Optional[str]:
    """
    Search for checkpoint directory matching patterns.

    Args:
        base_dir: Base directory to search in
        checkpoint_patterns: List of subdirectory names to try

    Returns:
        Path to checkpoint directory or None
    """
    for pattern in checkpoint_patterns:
        checkpoint_dir = os.path.join(base_dir, pattern)
        if os.path.exists(os.path.join(checkpoint_dir, 'best_model.pt')):
            return checkpoint_dir
    return None


# ============================================================================
# SECTION 5: Path Utilities (FIXED)
# ============================================================================

def get_script_directory() -> str:
    """
    Get directory where the calling script is located.

    Returns:
        Absolute path to script directory
    """
    if hasattr(sys.modules['__main__'], '__file__'):
        script_path = os.path.abspath(sys.modules['__main__'].__file__)
        return os.path.dirname(script_path)
    return os.path.abspath(os.getcwd())


def get_project_root() -> str:
    """
    Get project root directory (where config.yaml is expected).

    FIXED: Now searches from SCRIPT directory first, then goes up.

    Returns:
        Absolute path to project root
    """
    current = get_script_directory()

    if os.path.exists(os.path.join(current, 'config.yaml')):
        return current

    search_path = current
    while search_path != '/':
        if os.path.exists(os.path.join(search_path, 'config.yaml')):
            return search_path
        search_path = os.path.dirname(search_path)

    return current


def resolve_config_path(custom_path: Optional[str] = None) -> str:
    """
    Resolve configuration file path.

    UWAGA: Ta funkcja zwraca config.yaml (F95) bez uwzględnienia train_dataset.yaml.
    Używaj load_active_config() jeśli chcesz automatyczną selekcję datasetu.

    Args:
        custom_path: Custom path to config.yaml (optional)

    Returns:
        Resolved absolute path to config.yaml
    """
    if custom_path:
        return os.path.abspath(custom_path)

    project_root = get_project_root()
    return os.path.join(project_root, 'config.yaml')