"""

Script to convert files to adata objects for feeding into BABEL

Formatting for all files is autodetected

"""

import os, sys

import argparse

import logging

from typing import *

import pandas as pd

import anndata as ad

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 utils

logging.basicConfig(level=logging.INFO)

def auto_read_matrix_values(fname: str) -> pd.DataFrame:

    """Read the given counts file"""

    assert os.path.isfile(fname)

    ext = utils.get_file_extension_no_gz(fname)

    if ext == "csv":

        df = pd.read_csv(fname, sep=",")

    elif ext == "txt" or ext == "tsv":

        df = pd.read_csv(fname, sep="\t")

    else:

        raise ValueError(f"Cannot recognize file extension for {fname}")

    return df

def auto_read_metadata(fname: str, index_col: Optional[int] = None) -> pd.DataFrame:

    """Read the given metadata file"""

    assert os.path.isfile(fname)

    ext = utils.get_file_extension_no_gz(fname)

    if ext == "csv":

        df = pd.read_csv(fname, sep=",")

    elif ext == "tsv" or ext == "txt":

        df = pd.read_csv(fname, sep="\t")

    else:

        raise ValueError(f"Cannot recognize file extension for {fname}")

    if index_col:

        df.set_index(df.columns[index_col], inplace=True)

    return df

def build_parser():

    """Build CLI parser"""

    parser = argparse.ArgumentParser(

        usage=__doc__, formatter_class=argparse.ArgumentDefaultsHelpFormatter

    )

    parser.add_argument(

        "matfile",

        type=str,

        help="File containing matrix of vlaues. Expects (cell x feature), use transpose otherwise.",

    )

    parser.add_argument("out_h5ad", type=str, help="Output file (should end in .h5ad)")

    parser.add_argument(

        "-t",

        "--transpose",

        action="store_true",

        help="Apply transpose to matrix values to match expected (cell x feature) format",

    )

    parser.add_argument(

        "--obsinfo", type=str, help="Optional file for cell (observation) annotations"

    )

    parser.add_argument(

        "--obscol",

        type=int,

        help="Column of obs table to use as obs names",

    )

    parser.add_argument(

        "--varinfo",

        type=str,

        help="Optional file for feature (gene or peaks) annotations",

    )

    parser.add_argument(

        "--varcol",

        type=int,

        help="Column of var table to use as var names",

    )

    return parser

def main():

    """Run script"""

    parser = build_parser()

    args = parser.parse_args()

    # Read in the counts file

    counts_df = auto_read_matrix_values(args.matfile)

    if args.transpose:

        counts_df = counts_df.T

    logging.info(f"Read input matrix of (cell x feature) {counts_df.shape}")

    # Read in metadata if given

    obs_df = None

    if args.obsinfo:

        obs_df = auto_read_metadata(args.obsinfo, index_col=args.obscol)

    var_df = None

    if args.varinfo:

        var_df = auto_read_metadata(args.varinfo, index_col=args.varcol)

    adata = ad.AnnData(counts_df, obs=obs_df, var=var_df)

    logging.info(f"Created anndata object: {adata}")

    logging.info(f"Writing adata object to {args.out_h5ad}")

    adata.write(args.out_h5ad)

if __name__ == "__main__":

    main()