import sys
import scanpy as sc


def run_dimensionality_reduction(h5ad_file, sample_id, normalization="log1p_norm"):
    """PCA, t-SNE, UMAP. 

       https://www.sc-best-practices.org/preprocessing_visualization/dimensionality_reduction.html
    """
    lateral_root_markers = [
                            "Chr5g0400511", "Chr8g0389721", "Chr8g0389731",
                            "Chr2g0284281", "Chr3g0113941", "Chr8g0378811",
                            "Chr7g0265311"
                           ]
    sc.settings.verbosity = 0
    sc.settings.set_figure_params(
        dpi=80,
        facecolor="white",
        frameon=False,
    )

    adata = sc.read(
        filename=h5ad_file,
    )

    ## for now use log1p_norm as defailt we also have scran and pearson residuals
    adata.X = adata.layers[normalization]

    ## setting highly variable as highly deviant to use scanpy 'use_highly_variable' argument in sc.pp.pca
    adata.var["highly_variable"] = adata.var["highly_deviant"]
#    for o in adata.obs:
#        print(o)
#    sys.exit(1)
    sc.pp.pca(adata, svd_solver="arpack", use_highly_variable=True)  # PCA
    sc.pl.pca_scatter(adata, color="total_counts", save=f".{sample_id}.{normalization}.pca_scatter.png")
#    sc.pl.pca_scatter(adata, color="total_counts_dr5", save=f".{sample_id}.{normalization}.pca_scatter.dr5.png")
#    sc.pl.pca_scatter(adata, color="total_counts_tcsn", save=f".{sample_id}.{normalization}.pca_scatter.tcsn.png")

    ## t-SNE
    sc.tl.tsne(adata, use_rep="X_pca")
    sc.pl.tsne(adata, color="total_counts", save=f".{sample_id}.{normalization}.tsne.png")
#    sc.pl.tsne(adata, color="total_counts_dr5", save=f".{sample_id}.{normalization}.tsne.dr5.png")
#    sc.pl.tsne(adata, color="total_counts_tcsn", save=f".{sample_id}.{normalization}.tsne.tcsn.png")

    ## UMAP uses Neighbor graph
    sc.pp.neighbors(adata)
    sc.tl.umap(adata)
    sc.pl.umap(adata, color="total_counts", save=f".{sample_id}.{normalization}.umap.png")
#    sc.pl.umap(adata, color="total_counts_dr5", save=f".{sample_id}.{normalization}.umap.dr5.png")
#    sc.pl.umap(adata, color="total_counts_tcsn", save=f".{sample_id}.{normalization}.umap.tcsn.png")

    ## plot total counts, percentage mt counts, doublet detection score, is doublet in 4 UMAPs
    sc.pl.umap(
        adata,
#        color=[
#               "total_counts", "pct_counts_mt", "scDblFinder_score", 
#               "scDblFinder_class", "total_counts_dr5", "total_counts_tcsn"
#              ],
        color=[
               "total_counts", "pct_counts_mt", "scDblFinder_score", 
               "scDblFinder_class"
              ],
        save=f".{sample_id}.total_mt_dblscore_dbl.umap.png",
    )
    sc.pl.umap(adata, color=[ f"total_counts_{m}" for m in lateral_root_markers ],
               save=f".{sample_id}.all_marker_genes.umap.png")
    adata.write(f"{sample_id}.dimensionality_reduction.h5ad")


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}.feature_selection.h5ad"
    run_dimensionality_reduction(h5_input, sample_id, "log1p_norm")