import sys
import pandas as pd
import scanpy as sc


def run_clustering(h5ad_file, sample_id):
    """Clustering based on

       https://www.sc-best-practices.org/cellular_structure/clustering.html
    """
    sc.settings.verbosity = 0
    sc.settings.set_figure_params(
        dpi=80,
        facecolor="white",
        frameon=False,
    )

    adata = sc.read(
        filename=h5ad_file,
    )

    ## Neighbors for UMAP
    sc.pp.neighbors(adata, n_pcs=30)
    sc.tl.umap(adata)

    ## Run leiden clustering
    sc.tl.leiden(adata)

    ## Leiden clustering at 3 resolutions controlling the KNN cluster density
    sc.tl.leiden(adata, key_added="leiden_res0_25", resolution=0.25)
    sc.tl.leiden(adata, key_added="leiden_res0_5", resolution=0.5)
    sc.tl.leiden(adata, key_added="leiden_res1", resolution=1.0)
    pd.DataFrame(adata.obsm["X_umap"], index=adata.obs_names).to_csv(f"{sample_id}.umap.csv")
    leiden_res025 = open(f"./{sample_id}.ledien_resolution0_25.csv", "w")
    leiden_res05 = open(f"./{sample_id}.ledien_resolution0_5.csv", "w")
    leiden_res1 = open(f"./{sample_id}.ledien_resolution1.csv", "w")
    for barcode, cluster in adata.obs.leiden_res0_25.items():
        leiden_res025.write(f"{barcode},{cluster}\n")
    for barcode, cluster in adata.obs.leiden_res0_5.items():
        leiden_res05.write(f"{barcode},{cluster}\n")
    for barcode, cluster in adata.obs.leiden_res1.items():
        leiden_res1.write(f"{barcode},{cluster}\n")
    sc.pl.umap(
        adata,
        color=["leiden_res0_25", "leiden_res0_5", "leiden_res1"],
        legend_loc="on data",
        save=f".{sample_id}.leiden_clustering.umap.png",
    )
    leiden_res1.close()
    leiden_res05.close()
    leiden_res025.close()


if __name__ == "__main__":
    sample_id = sys.argv[1]  # get sample_id from command line, REMOVE LATER
    if not sample_id:
        sample_id = "2Cinoc_wholeroot"  # default for testing
    h5_input = f"./{sample_id}.dimensionality_reduction.h5ad"
    run_clustering(h5_input, sample_id)