from __future__ import annotations

import argparse

from .pipeline import run_proxai


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="proxai-ms",
        description="ProXAI-MS: train models and generate pseudo-MS1 from saliency.",
    )

    subparsers = parser.add_subparsers(dest="command", required=True)

    run_parser = subparsers.add_parser(
        "run",
        help="Run the full ProXAI pipeline on one or more bins.",
    )
    run_parser.add_argument("--csv", required=True, help="Path to input CSV file.")
    run_parser.add_argument("--label-column", default="target", help="Label column name.")
    run_parser.add_argument("--bin-column", default="bin", help="Bin column name.")
    run_parser.add_argument(
        "--bin-values",
        required=True,
        nargs="+",
        type=int,
        help="One or more bin values to process.",
    )
    run_parser.add_argument(
        "--control-labels",
        required=True,
        nargs="+",
        help="One or more label values from the label column to treat as controls.",
    )
    run_parser.add_argument(
        "--experiment-labels",
        required=True,
        nargs="+",
        help="One or more label values from the label column to treat as experiments.",
    )
    run_parser.add_argument("--mz-start", default=200.0, type=float, help="Starting m/z.")
    run_parser.add_argument("--bin-size", default=0.1, type=float, help="m/z bin size.")
    run_parser.add_argument(
        "--out-prefix",
        default="proxai_output",
        help="Output prefix for CSV and PNG files.",
    )

    return parser


def main() -> None:
    parser = build_parser()
    args = parser.parse_args()

    if args.command == "run":
        out = run_proxai(
            csv_path=args.csv,
            label_column=args.label_column,
            bin_column=args.bin_column,
            bin_values=args.bin_values,
            control_labels=args.control_labels,
            experiment_labels=args.experiment_labels,
            mz_start=args.mz_start,
            bin_size=args.bin_size,
            out_prefix=args.out_prefix,
        )
        print("ProXAI run complete.")
        print(f"Controls: {out['control_labels']}")
        print(f"Experiments: {out['experiment_labels']}")
        for bin_value, bin_out in out["bins"].items():
            print(f"Bin {bin_value} outputs:")
            print(f"  positive csv : {bin_out['csv']['positive']}")
            print(f"  negative csv : {bin_out['csv']['negative']}")
            print(f"  positive png : {bin_out['png']['positive']}")
            print(f"  negative png : {bin_out['png']['negative']}")


if __name__ == "__main__":
    main()
