import logging
import os
import random
from collections import OrderedDict
from typing import Any

import numpy as np
import torch
import yaml
from torch.utils.data import (
    DataLoader,
    SubsetRandomSampler,
    TensorDataset,
    random_split,
)

from gsMap.config import FindLatentRepresentationsConfig

from .gnn.st_model import StEmbeding
from .gnn.train_step import ModelTrain
from .st_process import (
    InferenceData,
    TrainingData,
    apply_module_score_qc,
    calculate_module_score,
    calculate_module_scores_from_degs,
    convert_to_human_genes,
    create_subsampled_adata,
    find_common_hvg,
)

logger = logging.getLogger(__name__)



def set_seed(seed_value):
    """
    Set seed for reproducibility in PyTorch and other libraries.
    """
    torch.manual_seed(seed_value)
    np.random.seed(seed_value)
    random.seed(seed_value)
    if torch.cuda.is_available():
        logger.info("Using GPU for computations.")
        torch.cuda.manual_seed(seed_value)
        torch.cuda.manual_seed_all(seed_value)
    else:
        logger.info("Using CPU for computations.")


def index_splitter(n, splits):
    idx = torch.arange(n)
    splits_tensor = torch.as_tensor(splits)
    multiplier = n / splits_tensor.sum()
    splits_tensor = (multiplier * splits_tensor).long()
    diff = n - splits_tensor.sum()
    splits_tensor[0] += diff
    return random_split(idx, splits_tensor)


def run_find_latent_representation(config: FindLatentRepresentationsConfig) -> dict[str, Any]:
    """
    Run the find latent representation pipeline.

    Args:
        config: FindLatentRepresentationsConfig object with all necessary parameters

    Returns:
        Dictionary containing metadata about the run including config, model info,
        training info, outputs, and annotation info
    """
    logger.info(f'Project dir: {config.project_dir}')
    set_seed(2024)

    # Find the hvg
    hvg, n_cell_used, gene_homolog_dict = find_common_hvg(config.sample_h5ad_dict, config)
    common_genes = np.array(list(gene_homolog_dict.keys()))

    # Create subsampled concatenated adata with sample-specific stratified sampling
    training_adata = create_subsampled_adata(config.sample_h5ad_dict, n_cell_used, config)

    # Prepare the trainning data
    get_trainning_data = TrainingData(config)
    get_trainning_data.prepare(training_adata, hvg)

    # Configure the distribution
    if config.data_layer in ["count", "counts"]:
        distribution = config.distribution
        variational = True
        use_tf = config.use_tf
    else:
        distribution = "gaussian"
        variational = False
        use_tf = False

    # Instantiation the LGCN VAE
    input_size = [
        get_trainning_data.expression_merge.size(1),
        get_trainning_data.expression_gcn_merge.size(1),
    ]
    class_size = len(torch.unique(get_trainning_data.label_merge))
    batch_size = get_trainning_data.batch_size
    cell_size, out_size = get_trainning_data.expression_merge.shape
    label_name = get_trainning_data.label_name

    # Configure the batch embedding dim
    batch_embedding_size = 64

    # Configure the model
    gsmap_lgcn_model = StEmbeding(
        # parameter of VAE
        input_size=input_size,
        hidden_size=config.hidden_size,
        embedding_size=config.embedding_size,
        batch_embedding_size=batch_embedding_size,
        out_put_size=out_size,
        batch_size=batch_size,
        class_size=class_size,
        # parameter of transformer
        module_dim=config.module_dim,
        hidden_gmf=config.hidden_gmf,
        n_modules=config.n_modules,
        nhead=config.nhead,
        n_enc_layer=config.n_enc_layer,
        # parameter of model structure
        distribution=distribution,
        use_tf=use_tf,
        variational=variational,
    )

    # Configure the optimizer
    optimizer = torch.optim.Adam(gsmap_lgcn_model.parameters(), lr=1e-3)
    logger.info(
        f"gsMap-LGCN parameters: {sum(p.numel() for p in gsmap_lgcn_model.parameters())}."
    )
    logger.info(f"Number of cells used in trainning: {cell_size}.")

    # Split the data to trainning (80%) and validation (20%).
    train_idx, val_idx = index_splitter(
        get_trainning_data.expression_gcn_merge.size(0), [80, 20]
    )
    train_sampler = SubsetRandomSampler(train_idx)
    val_sampler = SubsetRandomSampler(val_idx)

    # Configure the data loader
    dataset = TensorDataset(
        get_trainning_data.expression_gcn_merge,
        get_trainning_data.batch_merge,
        get_trainning_data.expression_merge,
        get_trainning_data.label_merge,
    )
    train_loader = DataLoader(
        dataset=dataset, batch_size=config.batch_size, sampler=train_sampler
    )
    val_loader = DataLoader(
        dataset=dataset, batch_size=config.batch_size, sampler=val_sampler
    )

    # Model trainning
    gsMap_embedding_finder = ModelTrain(
        gsmap_lgcn_model,
        optimizer,
        distribution,
        mode="reconstruction",
        lr=1e-3,
        model_path=config.model_path,
    )
    gsMap_embedding_finder.set_loaders(train_loader, val_loader)
    print(gsMap_embedding_finder.model)

    if not os.path.exists(config.model_path):
        # reconstruction
        gsMap_embedding_finder.train(config.itermax, patience=config.patience)

        # classification
        if config.two_stage and config.annotation is not None:
            gsMap_embedding_finder.model.load_state_dict(torch.load(config.model_path))
            gsMap_embedding_finder.mode = "classification"
            gsMap_embedding_finder.train(config.itermax, patience=config.patience)
    else:
        logger.info(f"Model found at {config.model_path}. Skipping training.")

    # Load the best model
    gsMap_embedding_finder.model.load_state_dict(torch.load(config.model_path))
    gsmap_embedding_model = gsMap_embedding_finder.model

    # Configure the inference
    infer = InferenceData(hvg, batch_size, gsmap_embedding_model, label_name, config)


    output_h5ad_path_dict = OrderedDict(
        {sample_name: config.latent_dir / f"{sample_name}_add_latent.h5ad"
            for sample_name in config.sample_h5ad_dict.keys()}
    )
    # Do the DEG in the training adata if annotation is provided and high quality cell QC is enabled
    module_score_threshold_dict = {}
    if config.high_quality_cell_qc and config.annotation is not None:
        # Calculate module scores for training data
        training_adata = calculate_module_score(training_adata, config.annotation)

        # Calculate thresholds for each annotation
        for label in training_adata.obs[config.annotation].cat.categories:
            scores = training_adata.obs.loc[training_adata.obs[config.annotation] == label, f"{label}_module_score"]
            Q1 = np.percentile(scores, 25)
            Q2 = np.median(scores)
            Q3 = np.percentile(scores, 75)
            IQR = Q3 - Q1
            threshold = max(0, Q2 - 1 * IQR)  # Ensure threshold is not negative
            module_score_threshold_dict[label] = threshold
            logger.info(f"High quality module score threshold for {label}: {threshold:.3f}")

        training_adata.uns['module_score_thresholds'] = module_score_threshold_dict

        # Apply QC to training data as well
        training_adata = apply_module_score_qc(training_adata, config.annotation, module_score_threshold_dict)

    # save the training adata
    training_adata_path = config.find_latent_metadata_path.parent / "training_adata.h5ad"
    training_adata.write_h5ad(training_adata_path)
    logger.info(f"Saved training adata to {training_adata_path}")

    for st_id, (sample_name, st_file) in enumerate(config.sample_h5ad_dict.items()):

        output_path = output_h5ad_path_dict[sample_name]

        # Infer the embedding
        adata = infer.infer_embedding_single(st_id, st_file)

        # Calculate module scores and apply QC for each annotation
        if config.high_quality_cell_qc and config.annotation is not None:
            # Calculate module scores for this sample using the same DEGs from training
            logger.info(f"Calculating module scores for {sample_name}...")

            # Get DEG results from training data
            deg_results = training_adata.uns['rank_genes_groups']

            # Keep the same gene list as training data, because the module score is based on the training data.
            # This is critical for proper normalization: without this filter, normalize_total()
            # would compute different scaling factors compared to training data.
            adata = adata[:, training_adata.var_names].copy()

            # Calculate module scores using existing DEG results
            adata = calculate_module_scores_from_degs(adata, deg_results, config.annotation)

            # Apply QC based on module score thresholds
            if config.annotation in adata.obs.columns:
                adata = apply_module_score_qc(adata, config.annotation, module_score_threshold_dict)
            else:
                logger.warning(f"Annotation '{config.annotation}' not found in {sample_name}, skipping QC")

        # Transfer to human gene names
        adata = convert_to_human_genes(adata, gene_homolog_dict, species=config.species)


        # Compute the depth
        if config.data_layer in ["count", "counts"]:
            adata.obs['depth'] = np.array(adata.layers[config.data_layer].sum(axis=1)).flatten()

        # Save the ST data with embeddings
        adata.write_h5ad(output_path)

        logger.info(f"Saved latent representation to {output_path}")

    # Convert config to dict with all Path objects as strings
    config_dict = config.to_dict_with_paths_as_strings()

    # Convert output_h5ad_path_dict to strings
    output_h5ad_path_dict_str = {k: str(v) for k, v in output_h5ad_path_dict.items()}

    # Save metadata
    metadata = {
        "config": config_dict,
        "model_info": {
            "model_path": str(config.model_path),
            "n_parameters": int(sum(p.numel() for p in gsmap_embedding_model.parameters())),
            "input_size": [int(x) for x in input_size],
            "hidden_size": int(config.hidden_size),
            "embedding_size": int(config.embedding_size),
            "batch_embedding_size": int(batch_embedding_size),
            "class_size": int(class_size),
            "distribution": distribution,
            "variational": variational,
            "use_tf": use_tf
        },
        "training_info": {
            "n_cells_used": int(cell_size),
            "n_genes_used": int(len(hvg)),
            "n_common_genes": int(len(common_genes)),
            "batch_size": int(config.batch_size),
            "n_epochs": int(config.itermax),
            "patience": int(config.patience),
            "two_stage": config.two_stage
        },
        "outputs": {
            "latent_files": output_h5ad_path_dict_str,
            "n_sections": len(config.sample_h5ad_dict)
        },
        "annotation_info": {
            "annotation_key": config.annotation,
            "n_classes": int(class_size),
            "label_names": label_name if isinstance(label_name, list) else label_name.tolist() if hasattr(label_name, 'tolist') else list(label_name),
        }
    }

    # Save metadata to YAML file
    metadata_path = config.find_latent_metadata_path
    with open(metadata_path, 'w') as f:
        yaml.dump(metadata, f, default_flow_style=False, sort_keys=False)

    logger.info(f"Saved metadata to {metadata_path}")

    return metadata
