"""

Basic script for pltoting pca of a dataset(s)

"""

import os

import sys

import argparse

import logging

import itertools

import numpy as np

import matplotlib.pyplot as plt

SRC_DIR = os.path.join(

    os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "babel",

)

assert os.path.isdir(SRC_DIR)

sys.path.append(SRC_DIR)

import interpretation

import adata_utils

import plot_utils

import utils

from evaluate_bulk_rna_concordance import load_file_flex_format

def var_barplot(var: np.ndarray, fname: str = ""):

    """Basic barplot of explained variance"""

    fig, ax = plt.subplots(dpi=300)

    ax.bar(np.arange(len(var)), var)

    ax.set(

        xlabel=f"Principal component",

        ylabel="Explained variance",

        title=f"Top {len(var)} PCs ({np.sum(var):.4f} explained variance)",

    )

    if fname:

        fig.savefig(fname, bbox_inches="tight")

    return fig

def build_parser():

    """Build a basic CLI parser"""

    parser = argparse.ArgumentParser()

    parser.add_argument("adata_fname", type=str, help="Adata object to plot PCA for")

    parser.add_argument("plot_prefix", type=str, help="Prefix to save plots to")

    parser.add_argument(

        "--numdims", "-n", type=int, default=16, help="Number of top PCs to consider"

    )

    return parser

def main():

    """Run script"""

    parser = build_parser()

    args = parser.parse_args()

    adata = load_file_flex_format(args.adata_fname)

    var = adata.uns["pca"]["variance_ratio"][: args.numdims]

    var_barplot(var, fname=args.plot_prefix + "_explained_var.pdf")

if __name__ == "__main__":

    main()