import xgboost as xgb
import joblib
from pathlib import Path
import scipy.sparse as sp
import numpy as np


def load_celltype_model(model_path,label_encoder_path):
    """
    Load the XGBoost model from the specified path.
    """
    model = xgb.XGBClassifier()
    model.load_model(model_path)

    # Access the booster and retrieve the feature names
    booster = model.get_booster()

    # Get the feature names (these should match the features used for training)
    feature_names = model.feature_names_in_

    label_encoder = joblib.load(label_encoder_path)

    return model, label_encoder, feature_names




def prepare_matrix_celltype(adata, feature_names,active_assay="sc"):
    """
    Prepare the AnnData object for prediction by selecting the relevant features.
    """

    adata_to_handle = adata.copy()

    # Define all potential assay features based on feature_names_sorted1
    assay_features = [f for f in feature_names if f.startswith('assay_')]

    #number of cells
    n_obs = adata_to_handle.n_obs

    #initialise assay data with zeros
    assay_data1 = np.zeros((n_obs, len(assay_features)), dtype=np.float32)

    # Create a mapping from assay feature name to its column index
    assay_feature_indices1 = {name: i for i, name in enumerate(assay_features)}

    #active assay must be one of sc, sn, multi_rna, if neither, set it as sc
    if active_assay not in ['sc', 'sn', 'multi_rna']:
        print(f"Warning: Active assay '{active_assay}' not recognized. Defaulting to 'sc'.")
        active_assay = 'sc'

    if active_assay in assay_feature_indices1:
        assay_data1[:, assay_feature_indices1[active_assay]] = 1.0

    # Convert assay data to sparse if original data is sparse
    assay_data_matrix1 = sp.csr_matrix(assay_data1) if sp.issparse(adata_to_handle.X) else assay_data1

    # --- 2. Combine Gene Expression and Assay Features for Model 1 ---
    # We need to combine these temporarily to easily subset later

    # Ensure original data is CSR for efficient column slicing if sparse
    if sp.issparse(adata_to_handle.X) and not isinstance(adata_to_handle.X, sp.csr_matrix):
        adata_X = adata_to_handle.X.tocsr().copy()
        print("Converted adata_orig.X to CSR format.")
    else:
        adata_X = adata_to_handle.X.copy()

    # Combine the matrices horizontally
    combined_X1 = sp.hstack([adata_X, assay_data_matrix1], format='csr') if sp.issparse(adata_X) else np.hstack([adata_X, assay_data1])

    # Create combined feature names list
    combined_feature_names1 = adata_to_handle.var_names.tolist() + assay_features

    # Create a mapping from the combined feature names to their column index
    combined_feature_indices1 = {name: i for i, name in enumerate(combined_feature_names1)}
    print(f"Combined matrix shape for model 1: {combined_X1.shape}")

    # Initialize the final matrix with NaNs (XGBoost can handle NaNs)
    X_final1 = np.full((n_obs, len(feature_names)), np.nan, dtype=np.float32)

    # Create a mapping for the target feature order
    target_feature_indices1 = {name: i for i, name in enumerate(feature_names)}

    # Find which features required by the model are present in our combined data
    available_features1 = [f for f in feature_names if f in combined_feature_indices1]
    missing_features1 = [f for f in feature_names if f not in combined_feature_indices1]

    if missing_features1:
        print(f"Warning: {len(missing_features1)} features required by model 1 are missing from the data: {missing_features1[:5]}...") # Print first 5

    print(f"Found {len(available_features1)} available features out of {len(feature_names)} required for model 1.")

    # Get the column indices in the *combined* data for the available features - indices of where it is in the combined matrix
    source_indices1 = [combined_feature_indices1[f] for f in available_features1]

    # Get the column indices in the *final* matrix for these available features - indices of where the model expects it
    target_indices1 = [target_feature_indices1[f] for f in available_features1]

    # Fill the final matrix with data from the available features
    # Ensure data is dense for assignment; handle potential memory issues for large datasets
    if sp.issparse(combined_X1):
        # Slice sparse matrix efficiently and convert to dense for assignment
        X_final1[:, target_indices1] = combined_X1[:, source_indices1].toarray()
        print("Filled final matrix for model 1 from sparse data.")
    else:
        X_final1[:, target_indices1] = combined_X1[:, source_indices1]
        print("Filled final matrix for model 1 from dense data.")
    
    
    return X_final1


def perform_celltype_prediction(matrix, model, label_encoder, return_probas=True):
    """
    Perform cell type prediction using the provided model and label encoder.
    """
    probas = model.predict_proba(matrix)
    predicted_labels = model.predict(matrix)
    predicted_cell_types = label_encoder.inverse_transform(predicted_labels)

    if return_probas:
        return predicted_cell_types, probas
    else:
        return predicted_cell_types