#!/usr/bin/env python
# ==============================================================================
# SCENIC+ Step 3: SCENIC+ Pipeline 核心推断
# ==============================================================================
# 功能：整合 Step 0~2 的所有产出，初始化 Snakemake pipeline、生成 config.yaml、
#       并执行 snakemake 完成 TF→调控元件→靶基因的完整调控网络推断。
#
# 输入（来自前序步骤）：
#   - scRNA.h5ad                          ：Step 0 输出
#   - cistopic_obj_with_models.pkl         ：Step 1 输出
#   - binarized_topics.pkl                ：Step 1 输出
#   - <species>_custom.*.feather          ：Step 2 输出 (rankings + scores)
#
# 输出：
#   - scplusmdata.h5mu      ：最终 MuData 对象（含 direct/extended eRegulon）
#   - ctx_results.hdf5      ：cisTarget motif 富集结果
#   - dem_results.hdf5      ：差异可及性分析结果
#   - ACC_GEX.h5mu          ：ATAC+GEX 联合 MuData
#   - cistromes_*.h5ad      ：direct/extended cistrome
#   - AUCell_*.h5mu         ：direct/extended AUCell 矩阵
#   - genome_annotation.tsv ：基因组注释
#   - 其他中间文件           ：search_space、adjacencies、eRegulons 等
#
# 用法：
#   python scripts/3_run_scenicplus_pipeline.py \
#     --project_dir /path/to/project \
#     --scRNA_h5ad /path/to/scRNA.h5ad \
#     --cistopic_obj /path/to/cistopic_obj_with_models.pkl \
#     --binarized_topics /path/to/binarized_topics.pkl \
#     --cistarget_dir /path/to/cisTarget \
#     --motif_annotation_dir /path/to/motif_snapshots \
#     --species mouse \
#     --n_cpu 16
#
# 注意：脚本执行完毕后，需要用户手动运行 snakemake（或直接使用 --run_snakemake 自动执行）。
#
# 所有路径均通过命令行参数指定，无服务器本地路径依赖。
# ==============================================================================

import os
import sys
import yaml
import pickle
import pandas as pd
import scanpy as sc
import shutil
import subprocess

import argparse

# ==============================================================================
# 0. 命令行参数
# ==============================================================================
parser = argparse.ArgumentParser(description="Run SCENIC+ pipeline automation")
parser.add_argument("--project_dir", type=str, required=True,
                    help="Base project directory")
parser.add_argument("--scRNA_h5ad", type=str, required=True,
                    help="Path to scRNA.h5ad file (from step 0)")
parser.add_argument("--cistopic_obj", type=str, required=True,
                    help="Path to cistopic_obj_with_models.pkl (from step 1)")
parser.add_argument("--binarized_topics", type=str, required=True,
                    help="Path to binarized_topics.pkl (from step 1)")
parser.add_argument("--cistarget_dir", type=str, required=True,
                    help="Directory containing cisTarget databases (from step 2)")
parser.add_argument("--motif_annotation_dir", type=str, required=True,
                    help="Directory containing motif annotation .tbl files (snapshots)")
parser.add_argument("--species", type=str, required=True,
                    help="Species name: mouse, human, fly, chicken, rat")
parser.add_argument("--n_cpu", type=int, default=16,
                    help="Number of CPUs (default: 16)")
parser.add_argument("--java_home", type=str, default="",
                    help="Path to JAVA_HOME (default: use system JAVA_HOME)")
parser.add_argument("--fix_genome_script", type=str, default="",
                    help="Path to fix_genome_annotation.py (optional, skip if not provided)")
parser.add_argument("--chrom_sizes", type=str, default="",
                    help="Path to species.chrom.sizes file (optional, for chromsizes.tsv generation)")
parser.add_argument("--run_snakemake", action="store_true",
                    help="Also run snakemake after config generation")

args = parser.parse_args()

# 设置 JAVA 环境（如果提供了）
if args.java_home:
    os.environ["JAVA_HOME"] = args.java_home
    conda_lib_path = os.path.join(os.environ["JAVA_HOME"], "lib")
    if "LD_LIBRARY_PATH" in os.environ:
        os.environ["LD_LIBRARY_PATH"] = conda_lib_path + ":" + os.environ["LD_LIBRARY_PATH"]
    else:
        os.environ["LD_LIBRARY_PATH"] = conda_lib_path
    os.environ["PATH"] = os.environ["JAVA_HOME"] + "/bin:" + os.environ["PATH"]

# ==============================================================================
# 1. 路径配置
# ==============================================================================
PROJECT_DIR    = args.project_dir
DATA_DIR       = os.path.join(PROJECT_DIR, "data")
OUTPUT_DIR     = os.path.join(PROJECT_DIR, "scenicplus_pipeline")
PYCISTOPIC_DIR = os.path.join(PROJECT_DIR, "pycisTopic")
CISTARGET_DIR  = args.cistarget_dir
TEMP_DIR       = os.path.join(PROJECT_DIR, "tmp")
PIPELINE_DIR   = os.path.join(PROJECT_DIR, "scplus_pipeline")

# 输入文件（全部通过参数传入）
SCRNA_H5AD_PATH     = args.scRNA_h5ad
CISTOPIC_OBJ_PATH   = args.cistopic_obj
BINARIZED_TOPICS_PATH = args.binarized_topics

# 物种
SPECIES_NAME = args.species
N_CPU        = args.n_cpu

# Motif annotation
MOTIF_ANNOTATION_DIR = args.motif_annotation_dir

# ==============================================================================
# 2. 辅助函数
# ==============================================================================
def get_species_params(species):
    """将通用物种名映射到 SCENIC+ config 所需参数。"""
    params = {}
    if species == "mouse":
        params['biomart_species'] = "mmusculus"
        params['motif_species']   = "mus_musculus"
        params['annotation_file'] = "motifs-v10-nr.mgi-m0.00001-o0.0.tbl"
        params['assembly']        = "mm10"
    elif species == "human":
        params['biomart_species'] = "hsapiens"
        params['motif_species']   = "homo_sapiens"
        params['annotation_file'] = "motifs-v10-nr.hgnc-m0.00001-o0.0.tbl"
        params['assembly']        = "hg38"
    elif species == "fly":
        params['biomart_species'] = "dmelanogaster"
        params['motif_species']   = "drosophila_melanogaster"
        params['annotation_file'] = "motifs-v10-nr.flybase-m0.00001-o0.0.tbl"
        params['assembly']        = "dm6"
    elif species == "chicken":
        params['biomart_species'] = "ggallus"
        params['motif_species']   = "gallus_gallus"
        params['annotation_file'] = "motifs-v10-nr.chicken-m0.00001-o0.0.tbl"
        params['assembly']        = "GRCg7b"
    elif species == "rat":
        params['biomart_species'] = "rnorvegicus"
        params['motif_species']   = "rattus_norvegicus"
        params['annotation_file'] = "motifs-v10-nr.rat-m0.00001-o0.0.tbl"
        params['assembly']        = "rn7"
    else:
        raise ValueError(f"Species {species} not supported.")
    return params

def export_region_sets(cistopic_obj_path, output_dir):
    """从 binarized_topics.pkl 导出每个 topic 的 BED 文件（Snakemake 要求的格式）。"""
    print("[Step 3] Loading binarized topics to extract region sets...")
    try:
        with open(BINARIZED_TOPICS_PATH, 'rb') as f:
            binarized_topics = pickle.load(f)
        print("[Step 3] Loaded binarized topics from file.")
    except Exception as e:
        print(f"Error: Cannot load binarized topics: {e}")
        print("Please ensure step 1 completed successfully.")
        sys.exit(1)

    topics_dir = os.path.join(output_dir, "Topics")
    if not os.path.exists(topics_dir):
        os.makedirs(topics_dir)

    print(f"[Step 3] Exporting topic regions to {topics_dir}...")
    for topic, regions in binarized_topics.items():
        if isinstance(regions, pd.DataFrame):
            region_list = regions.index.tolist()
        elif isinstance(regions, dict):
            region_list = list(regions.keys())
        else:
            try:
                region_list = list(regions)
            except Exception:
                print(f"Warning: Cannot parse regions for topic {topic}. Skipping.")
                continue

        bed_data = []
        for r in region_list:
            if ':' in r:
                chrom, rest = r.split(':')
                start, end = rest.split('-')
            elif '-' in r:
                parts = r.split('-')
                chrom = parts[0]
                start = parts[1]
                end = parts[2]
            else:
                print(f"Warning: Skipping malformed region: {r}")
                continue
            bed_data.append([chrom, start, end])

        bed_df = pd.DataFrame(bed_data)
        safe_topic = topic.replace(' ', '_').replace(':', '_')
        bed_df.to_csv(os.path.join(topics_dir, f"{safe_topic}.bed"),
                      sep='\t', header=False, index=False)

    print(f"[Step 3] Exported {len(binarized_topics)} topic BED files.")
    return output_dir

# ==============================================================================
# 3. 主流程
# ==============================================================================
def main():
    print("[Step 3] Starting SCENIC+ Pipeline Automation...")

    # 3.1 创建目录
    for d in [OUTPUT_DIR, TEMP_DIR]:
        if not os.path.exists(d):
            os.makedirs(d)

    sp_params = get_species_params(SPECIES_NAME)

    # 3.2 导出 region sets
    region_sets_dir = os.path.join(OUTPUT_DIR, "region_sets")
    export_region_sets(CISTOPIC_OBJ_PATH, region_sets_dir)

    # 3.3 准备 chromsizes（3 列格式，可选）
    if args.chrom_sizes and os.path.exists(args.chrom_sizes):
        chromsizes_src = args.chrom_sizes
        chromsizes_dst = os.path.join(OUTPUT_DIR, "chromsizes.tsv")
        print(f"[Step 3] Converting chromsizes from {chromsizes_src}")
        try:
            df_chrom = pd.read_csv(chromsizes_src, sep='\t', header=None, names=['Chromosome', 'Size'])
            df_fixed = pd.DataFrame()
            df_fixed['Chromosome'] = df_chrom['Chromosome']
            df_fixed['Start'] = 0
            df_fixed['End'] = df_chrom['Size']
            df_fixed.to_csv(chromsizes_dst, sep='\t', index=False)
            print("[Step 3] chromsizes.tsv created (3 columns).")
        except Exception as e:
            print(f"Error preparing chromsizes: {e}")
    else:
        print("[Step 3] chromsizes file not provided or not found, skipping chromsizes.tsv generation.")

    # 3.4 准备 genome_annotation（可选，用户提供 fix 脚本时执行）
    genome_annot_dst = os.path.join(OUTPUT_DIR, "genome_annotation.tsv")
    if args.fix_genome_script and os.path.exists(args.fix_genome_script):
        print(f"[Step 3] Running GTF parser for {SPECIES_NAME}...")
        cmd = [sys.executable, args.fix_genome_script, "--species", SPECIES_NAME, "--output_dir", OUTPUT_DIR]
        ret = subprocess.run(cmd, check=False)
        if ret.returncode == 0 and os.path.exists(genome_annot_dst):
            print("[Step 3] genome_annotation.tsv generated successfully.")
        else:
            print("Warning: Failed to generate genome_annotation.tsv from GTF.")
    else:
        print("[Step 3] fix_genome_annotation script not provided, skipping genome_annotation generation.")

    # 3.5 初始化 Snakemake pipeline
    print("[Step 3] Initializing Snakemake pipeline...")
    if not os.path.exists(os.path.join(PIPELINE_DIR, "Snakemake")):
        if not os.path.exists(PIPELINE_DIR):
            os.makedirs(PIPELINE_DIR)
        cmd = f"scenicplus init_snakemake --out_dir {PIPELINE_DIR}"
        print(f"[Step 3] Running: {cmd}")
        ret = os.system(cmd)
        if ret != 0:
            print("Error: Failed to initialize Snakemake pipeline.")
            sys.exit(1)
    else:
        print("[Step 3] Snakemake pipeline already initialized, skipping init...")

    # 3.6 生成 config.yaml
    print("[Step 3] Generating config.yaml...")

    db_prefix = f"{SPECIES_NAME}_custom"
    ctx_db    = os.path.join(CISTARGET_DIR, f"{db_prefix}.regions_vs_motifs.rankings.feather")
    dem_db    = os.path.join(CISTARGET_DIR, f"{db_prefix}.regions_vs_motifs.scores.feather")
    motif_annot = os.path.join(MOTIF_ANNOTATION_DIR, sp_params['annotation_file'])

    required_files = [SCRNA_H5AD_PATH, CISTOPIC_OBJ_PATH, ctx_db, dem_db, motif_annot]
    for f in required_files:
        if not os.path.exists(f):
            print(f"Warning: Required file not found: {f}")

    config = {
        'input_data': {
            'cisTopic_obj_fname': CISTOPIC_OBJ_PATH,
            'GEX_anndata_fname': SCRNA_H5AD_PATH,
            'region_set_folder': region_sets_dir,
            'ctx_db_fname': ctx_db,
            'dem_db_fname': dem_db,
            'path_to_motif_annotations': motif_annot,
        },
        'output_data': {
            'combined_GEX_ACC_mudata': os.path.join(OUTPUT_DIR, "ACC_GEX.h5mu"),
            'dem_result_fname': os.path.join(OUTPUT_DIR, "dem_results.hdf5"),
            'ctx_result_fname': os.path.join(OUTPUT_DIR, "ctx_results.hdf5"),
            'output_fname_dem_html': os.path.join(OUTPUT_DIR, "dem_results.html"),
            'output_fname_ctx_html': os.path.join(OUTPUT_DIR, "ctx_results.html"),
            'cistromes_direct': os.path.join(OUTPUT_DIR, "cistromes_direct.h5ad"),
            'cistromes_extended': os.path.join(OUTPUT_DIR, "cistromes_extended.h5ad"),
            'tf_names': os.path.join(OUTPUT_DIR, "tf_names.txt"),
            'genome_annotation': os.path.join(OUTPUT_DIR, "genome_annotation.tsv"),
            'chromsizes': os.path.join(OUTPUT_DIR, "chromsizes.tsv"),
            'search_space': os.path.join(OUTPUT_DIR, "search_space.tsv"),
            'tf_to_gene_adjacencies': os.path.join(OUTPUT_DIR, "tf_to_gene_adj.tsv"),
            'region_to_gene_adjacencies': os.path.join(OUTPUT_DIR, "region_to_gene_adj.tsv"),
            'eRegulons_direct': os.path.join(OUTPUT_DIR, "eRegulon_direct.tsv"),
            'eRegulons_extended': os.path.join(OUTPUT_DIR, "eRegulons_extended.tsv"),
            'AUCell_direct': os.path.join(OUTPUT_DIR, "AUCell_direct.h5mu"),
            'AUCell_extended': os.path.join(OUTPUT_DIR, "AUCell_extended.h5mu"),
            'scplus_mdata': os.path.join(OUTPUT_DIR, "scplusmdata.h5mu"),
        },
        'params_general': {
            'temp_dir': TEMP_DIR,
            'n_cpu': N_CPU,
            'seed': 666,
        },
        'params_data_preparation': {
            'use_raw_for_GEX_anndata': False,
            'bc_transform_func': '"lambda x: x"',
            'is_multiome': True,
            'key_to_group_by': "",
            'nr_cells_per_metacells': 10,
            'direct_annotation': "Direct_annot",
            'extended_annotation': "Orthology_annot",
            'species': sp_params['biomart_species'],
            'biomart_host': "http://www.ensembl.org",
            'search_space_upstream': "1000 150000",
            'search_space_downstream': "1000 150000",
            'search_space_extend_tss': "10 10",
        },
        'params_motif_enrichment': {
            'species': sp_params['motif_species'],
            'annotation_version': "v10nr_clust",
            'motif_similarity_fdr': 0.001,
            'orthologous_identity_threshold': 0.0,
            'annotations_to_use': "Direct_annot Orthology_annot",
            'fraction_overlap_w_dem_database': 0.4,
            'dem_max_bg_regions': 500,
            'dem_balance_number_of_promoters': True,
            'dem_promoter_space': 1000,
            'dem_adj_pval_thr': 0.05,
            'dem_log2fc_thr': 1.0,
            'dem_mean_fg_thr': 0.0,
            'dem_motif_hit_thr': 3.0,
            'fraction_overlap_w_ctx_database': 0.4,
            'ctx_auc_threshold': 0.005,
            'ctx_nes_threshold': 3.0,
            'ctx_rank_threshold': 0.05,
        },
        'params_inference': {
            'tf_to_gene_importance_method': "GBM",
            'region_to_gene_importance_method': "GBM",
            'region_to_gene_correlation_method': "SR",
            'order_regions_to_genes_by': "importance",
            'order_TFs_to_genes_by': "importance",
            'gsea_n_perm': 1000,
            'quantile_thresholds_region_to_gene': "0.85 0.90 0.95",
            'top_n_regionTogenes_per_gene': "5 10 15",
            'top_n_regionTogenes_per_region': "",
            'min_regions_per_gene': 0,
            'rho_threshold': 0.05,
            'min_target_genes': 10,
        },
    }

    config_path = os.path.join(PIPELINE_DIR, "Snakemake", "config", "config.yaml")
    print(f"[Step 3] Writing config to {config_path}...")
    with open(config_path, 'w') as f:
        yaml.dump(config, f, default_flow_style=False, sort_keys=False)
    print("[Step 3] Config file generated successfully.")

    # 3.7 Patch Snakefile（跳过 download_genome_annotations，使用本地生成的文件）
    snakefile_path = os.path.join(PIPELINE_DIR, "Snakemake", "workflow", "Snakefile")
    print(f"[Step 3] Patching Snakefile at {snakefile_path}...")
    try:
        with open(snakefile_path, 'r') as f:
            snakefile_content = f.read()

        if "scenicplus prepare_data download_genome_annotations" in snakefile_content:
            old_cmd = """scenicplus prepare_data download_genome_annotations \\
                    --biomart_host {params.biomart_host} \\
                    --species {params.species} \\
                    --out_dir {params.out_dir}"""
            new_cmd = """echo 'Skipping download_genome_annotations: using locally generated genome_annotation.tsv and chromsizes.tsv'"""
            snakefile_content = snakefile_content.replace(old_cmd, new_cmd)
            with open(snakefile_path, 'w') as f:
                f.write(snakefile_content)
            print("[Step 3] Snakefile patched successfully.")
        else:
            print("[Step 3] Snakefile already patched or rule not found.")
    except Exception as e:
        print(f"Warning: Could not patch Snakefile: {e}")

    # 3.8 执行 Snakemake（可选）
    if args.run_snakemake:
        print("[Step 3] Running Snakemake...")
        snakemake_dir = os.path.join(PIPELINE_DIR, "Snakemake")
        cmd = (
            f"cd {snakemake_dir} && "
            f"snakemake --cores {N_CPU} --use-conda "
            f"--envvars PATH LD_LIBRARY_PATH "
            f"OPENBLAS_NUM_THREADS=1 OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 "
            f"NUMEXPR_NUM_THREADS=1 VECLIB_MAXIMUM_THREADS=1 BLAS_NUM_THREADS=1"
        )
        print(f"[Step 3] Executing: {cmd}")
        ret = os.system(cmd)
        if ret != 0:
            print("Error: Snakemake execution failed.")
            sys.exit(1)
        print("[Step 3] Snakemake completed successfully.")
    else:
        print("\n[Step 3] Config 已生成。如需自动运行 Snakemake，请添加 --run_snakemake 参数。")
        print(f"手动运行命令:")
        print(f"  cd {os.path.join(PIPELINE_DIR, 'Snakemake')}")
        print(f"  snakemake --cores {N_CPU} --use-conda --envvars PATH LD_LIBRARY_PATH "
              f"OPENBLAS_NUM_THREADS=1 OMP_NUM_THREADS=1")

    print("\n[Step 3] SCENIC+ pipeline automation completed (config ready).")

if __name__ == "__main__":
    main()