"""Command-line interface: hmoe-da predict input.h5ad -o output.h5ad"""
from __future__ import annotations

import argparse
import sys
from pathlib import Path


def _cmd_predict(args) -> int:
    from .io import read_adata
    from .model import predict

    ad = read_adata(args.input)
    print(f"loaded {args.input}: {ad.shape}")

    if args.layer:
        if args.layer not in ad.layers:
            print(f"ERROR: layer {args.layer!r} not in adata.layers ({list(ad.layers)})",
                  file=sys.stderr)
            return 2
        ad.X = ad.layers[args.layer]
        print(f"using layer {args.layer!r} as counts")
    elif args.use_raw:
        if ad.raw is None:
            print("ERROR: --use-raw given but adata.raw is None", file=sys.stderr)
            return 2
        ad = ad.raw.to_adata()
        print("using adata.raw as counts")

    ad = predict(
        ad,
        include_gad2=not args.exclude_gad2,
        assume_raw=args.assume_raw,
        key_added=args.key,
    )

    obs_cols = [c for c in ad.obs.columns if c.startswith(args.key)]
    print("\nper-cell results written to adata.obs:")
    for c in obs_cols:
        print(f"  {c}")
    print("\nsubtype counts:")
    print(ad.obs[f"{args.key}_subtype"].value_counts().to_string())

    out = args.output
    ad.write_h5ad(out)
    print(f"\nwrote {out}")

    if args.csv:
        df = ad.obs[obs_cols].copy()
        df.index.name = "cell"
        df.to_csv(args.csv)
        print(f"wrote {args.csv}")
    return 0


def main(argv=None) -> int:
    p = argparse.ArgumentParser(
        prog="hmoe-da",
        description="Assign dopaminergic-neuron subtypes to cells in an sc/snRNA-seq dataset.",
    )
    sub = p.add_subparsers(dest="command", required=True)

    pp = sub.add_parser("predict", help="predict subtypes for an .h5ad file")
    pp.add_argument("input", help="input .h5ad or .rds (Seurat) with RAW COUNTS (.h5ad: in .X, or use --layer/--use-raw)")
    pp.add_argument("-o", "--output", required=True, help="output .h5ad path")
    pp.add_argument("--csv", help="also write per-cell labels to this CSV path")
    pp.add_argument("--exclude-gad2", action="store_true",
                    help="exclude the Gad2 family: reassign Gad2-called cells to their best non-Gad2 subtype")
    pp.add_argument("--layer", help="use adata.layers[LAYER] as the count matrix")
    pp.add_argument("--use-raw", action="store_true", help="use adata.raw as the count matrix")
    pp.add_argument("--assume-raw", action="store_true",
                    help="skip the raw-counts sanity check (use if you are sure .X is counts)")
    pp.add_argument("--key", default="hmoe", help="prefix for output columns (default: hmoe)")
    pp.set_defaults(func=_cmd_predict)

    args = p.parse_args(argv)
    return args.func(args)


if __name__ == "__main__":
    raise SystemExit(main())