"""
CLI Module - Command-Line Interface for DeepParameters
======================================================

Provides terminal access to the core DeepParameters API via the
``deepparameters`` (or ``dp``) command installed as a console script.

Commands
--------
  learn           Learn a CPD for a single node from a CSV dataset.
  learn-network   Learn CPDs for every node in a network, in parallel.
  compare         Compare a saved learned CPD against a saved true CPD.
  info            Print version, available architectures, and sampling methods.

Usage examples
--------------
  deepparameters info

  deepparameters learn \\
      --node heart_disease \\
      --data data.csv \\
      --edges edges.csv \\
      --num-parameters 20 \\
      --network-type vae \\
      --sampling-method 4 \\
      --output cpd_heart_disease.json

  deepparameters learn-network \\
      --data data.csv \\
      --edges edges.csv \\
      --parallel-style topological \\
      --max-workers 4 \\
      --output-dir ./cpds/

  deepparameters compare \\
      --learned cpd_heart_disease.json \\
      --true    cpd_true.json
"""

import argparse
import json
import os
import sys


# ---------------------------------------------------------------------------
# CPD serialisation helpers (stdlib only — no pgmpy at import time)
# ---------------------------------------------------------------------------

def _to_json_safe(obj):
    """Recursively convert numpy scalars / arrays to Python-native types."""
    try:
        import numpy as np
        if isinstance(obj, np.ndarray):
            return obj.tolist()
        if isinstance(obj, (np.integer,)):
            return int(obj)
        if isinstance(obj, (np.floating,)):
            return float(obj)
        if isinstance(obj, (np.bool_,)):
            return bool(obj)
    except ImportError:
        pass
    if isinstance(obj, dict):
        return {str(k): _to_json_safe(v) for k, v in obj.items()}
    if isinstance(obj, (list, tuple)):
        return [_to_json_safe(x) for x in obj]
    return obj


def _cpd_to_dict(cpd):
    """
    Serialise a pgmpy ``TabularCPD`` to a plain dict that is safe for JSON.

    The ``values`` array is stored in 2-D form: shape ``[variable_card, -1]``.
    ``state_names`` values are stored as lists so they round-trip cleanly.
    """
    import numpy as np

    values_2d = cpd.values.reshape(cpd.variable_card, -1)

    # evidence / evidence_card — handle both old and new pgmpy APIs
    if hasattr(cpd, "evidence") and not callable(getattr(type(cpd), "evidence", None)):
        evidence = list(cpd.evidence) if cpd.evidence else []
    else:
        evidence = list(cpd.variables[1:]) if len(cpd.variables) > 1 else []

    if hasattr(cpd, "evidence_card"):
        evidence_card = [int(x) for x in cpd.evidence_card] if cpd.evidence_card is not None else []
    else:
        evidence_card = [int(x) for x in cpd.cardinality[1:]] if len(cpd.cardinality) > 1 else []

    state_names = {}
    if hasattr(cpd, "state_names") and cpd.state_names:
        for var, states in cpd.state_names.items():
            state_names[str(var)] = _to_json_safe(list(states))

    return {
        "variable": str(cpd.variable),
        "variable_card": int(cpd.variable_card),
        "values": _to_json_safe(values_2d.tolist()),
        "evidence": [str(e) for e in evidence],
        "evidence_card": evidence_card,
        "state_names": state_names,
    }


def _dict_to_cpd(d):
    """Reconstruct a ``TabularCPD`` from a JSON-loaded dict."""
    from pgmpy.factors.discrete import TabularCPD
    import numpy as np

    values = np.array(d["values"], dtype=float)
    evidence = d.get("evidence") or []
    evidence_card = d.get("evidence_card") or []

    # Reconstruct state_names — JSON keys are always str; variable names stay str
    raw_sn = d.get("state_names") or {}
    state_names = {}
    for var, states in raw_sn.items():
        # Attempt to restore integer states that were stringified by JSON
        restored = []
        for s in states:
            if isinstance(s, str):
                try:
                    restored.append(int(s))
                except ValueError:
                    try:
                        restored.append(float(s))
                    except ValueError:
                        restored.append(s)
            else:
                restored.append(s)
        state_names[var] = restored

    kwargs = dict(
        variable=d["variable"],
        variable_card=int(d["variable_card"]),
        values=values,
    )
    if evidence:
        kwargs["evidence"] = evidence
        kwargs["evidence_card"] = [int(x) for x in evidence_card]
    if state_names:
        kwargs["state_names"] = state_names

    return TabularCPD(**kwargs)


def _save_cpd(cpd, path):
    """Write a CPD to *path* as JSON."""
    os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
    with open(path, "w") as fh:
        json.dump(_cpd_to_dict(cpd), fh, indent=2)


def _load_cpd(path):
    """Load a CPD from a JSON file created by ``_save_cpd``."""
    with open(path) as fh:
        d = json.load(fh)
    return _dict_to_cpd(d)


# ---------------------------------------------------------------------------
# Data / network loading helpers
# ---------------------------------------------------------------------------

def _load_data(path):
    """Load a CSV file and return a pandas DataFrame."""
    import pandas as pd

    if not os.path.isfile(path):
        _die(f"Data file not found: {path}")
    try:
        df = pd.read_csv(path)
    except Exception as exc:
        _die(f"Could not read data CSV '{path}': {exc}")

    if df.empty:
        _die(f"Data file '{path}' is empty.")
    return df


def _load_network(edges_path):
    """
    Load a two-column CSV (``parent,child``) and return a ``BayesianNetwork``.

    The CSV may optionally have a header row; if the first row looks like
    data (not a literal "parent,child" header), it is treated as an edge.
    """
    import pandas as pd

    try:
        from pgmpy.models import DiscreteBayesianNetwork as BayesianNetwork
    except ImportError:
        from pgmpy.models import BayesianNetwork

    if not os.path.isfile(edges_path):
        _die(f"Edges file not found: {edges_path}")

    try:
        df = pd.read_csv(edges_path, header=None, names=["parent", "child"])
    except Exception as exc:
        _die(f"Could not read edges CSV '{edges_path}': {exc}")

    # Drop header-like rows if user included a header
    df = df[~((df["parent"].str.lower() == "parent") & (df["child"].str.lower() == "child"))]
    df = df.dropna()

    if df.empty:
        _die(f"Edges file '{edges_path}' contains no edges.")

    edges = list(zip(df["parent"].str.strip(), df["child"].str.strip()))
    try:
        bn = BayesianNetwork(edges)
    except Exception as exc:
        _die(f"Could not build BayesianNetwork from edges: {exc}")
    return bn


# ---------------------------------------------------------------------------
# Utility
# ---------------------------------------------------------------------------

def _die(msg, code=1):
    """Print an error message to stderr and exit."""
    print(f"ERROR: {msg}", file=sys.stderr)
    sys.exit(code)


def _print_cpd_summary(cpd):
    """Print a compact human-readable summary of a CPD."""
    import numpy as np

    evidence = []
    if hasattr(cpd, "evidence") and not callable(getattr(type(cpd), "evidence", None)):
        evidence = list(cpd.evidence) if cpd.evidence else []
    else:
        evidence = list(cpd.variables[1:]) if len(cpd.variables) > 1 else []

    print(f"\n  Variable : {cpd.variable}  (card={cpd.variable_card})")
    print(f"  Parents  : {evidence if evidence else '(root node)'}")
    print(f"  Shape    : {cpd.values.shape}")
    col_sums = cpd.values.reshape(cpd.variable_card, -1).sum(axis=0)
    ok = bool(all(abs(s - 1.0) < 1e-4 for s in col_sums))
    print(f"  Valid    : {'YES (columns sum to 1)' if ok else 'NO - check values!'}")
    if hasattr(cpd, "state_names") and cpd.state_names:
        node_states = cpd.state_names.get(cpd.variable, [])
        print(f"  States   : {node_states}")


# ---------------------------------------------------------------------------
# Command handlers
# ---------------------------------------------------------------------------

def cmd_info(_args):
    """``deepparameters info`` — print package information."""
    import deepparameters

    print(f"\nDeepParameters v{deepparameters.__version__}")
    print(f"Python {sys.version.split()[0]}  |  {sys.executable}")
    print()

    print("AVAILABLE ARCHITECTURES (--network-type)")
    print("-" * 44)
    archs = {
        "simple":          "MLP (64→32) — fast baseline, always works",
        "advanced":        "MLP (128→64→32) — general purpose",
        "ultra":           "MLP (256→128→64→32) — higher accuracy",
        "mega":            "MLP (512→256→128→64→32) — maximum MLP",
        "lstm":            "LSTM — temporal / sequential patterns",
        "bnn":             "Bayesian NN + MC Dropout — uncertainty",
        "vae":             "Variational Autoencoder — generative",
        "autoencoder":     "Autoencoder — feature compression",
        "normalizing_flow":"Affine coupling layers — exact density",
    }
    for name, desc in archs.items():
        print(f"  {name:<20} {desc}")

    print()
    print("AVAILABLE SAMPLING METHODS (--sampling-method)")
    print("-" * 50)
    methods = {
        "1": "Gibbs Sampling",
        "2": "Metropolis-Hastings",
        "3": "Importance Sampling",
        "4": "Bayesian Parameter Estimation (BPE) — recommended fast default",
        "5": "Variational Inference",
        "6": "Hamiltonian Monte Carlo (HMC)",
        "7": "Sequential Monte Carlo",
        "8": "Adaptive KDE — good for sparse data",
    }
    for key, desc in methods.items():
        print(f"  {key}   {desc}")
    print()


def cmd_learn(args):
    """``deepparameters learn`` — learn a CPD for a single node."""
    from deepparameters import learn_cpd_for_node

    data = _load_data(args.data)
    bn = _load_network(args.edges)

    print(f"\nData   : {args.data}  ({len(data)} rows × {len(data.columns)} cols)")
    print(f"Network: {len(bn.nodes())} nodes, {len(bn.edges())} edges")
    print(f"Node   : {args.node}")
    print()

    cpd = learn_cpd_for_node(
        node=args.node,
        data=data,
        true_model=bn,
        learnt_bn_structure=bn,
        num_parameters=args.num_parameters,
        network_type=args.network_type,
        sampling_method=str(args.sampling_method),
        epochs=args.epochs,
        batch_size=args.batch_size,
        learning_rate=args.learning_rate,
        validation_split=args.validation_split,
        early_stopping=args.early_stopping,
        optimizer=args.optimizer,
        early_stopping_patience=args.early_stopping_patience,
        verbose=args.verbose,
        random_state=args.random_state,
    )

    print("\nLearned CPD:")
    _print_cpd_summary(cpd)

    if args.output:
        _save_cpd(cpd, args.output)
        print(f"\nSaved to: {args.output}")
    else:
        # Pretty-print the raw CPD dict to stdout
        print("\nCPD (JSON):")
        print(json.dumps(_cpd_to_dict(cpd), indent=2))


def cmd_learn_network(args):
    """``deepparameters learn-network`` — learn CPDs for all nodes in parallel."""
    from deepparameters import DeepParametersLearner

    data = _load_data(args.data)
    bn = _load_network(args.edges)

    print(f"\nData    : {args.data}  ({len(data)} rows × {len(data.columns)} cols)")
    print(f"Network : {len(bn.nodes())} nodes, {len(bn.edges())} edges")
    print(f"Style   : {args.parallel_style}  (max-workers={args.max_workers or 'auto'})")
    print()

    learner = DeepParametersLearner()
    result = learner.learn_network_parallel(
        data=data,
        network_structure=bn,
        parallel_style=args.parallel_style,
        max_workers=args.max_workers,
        network_type=args.network_type,
        sampling_method=str(args.sampling_method),
        epochs=args.epochs,
        verbose=args.verbose,
    )

    # result may be a dict {node: CPD} or a list [CPDs]
    if isinstance(result, dict):
        cpd_items = result.items()
    else:
        # list — derive node names from CPD objects
        cpd_items = [(cpd.variable, cpd) for cpd in result]

    output_dir = args.output_dir or "."
    os.makedirs(output_dir, exist_ok=True)

    print(f"\nResults: {len(list(cpd_items))} CPDs")
    # Re-evaluate since we may have consumed the iterator
    if isinstance(result, dict):
        cpd_items = result.items()
    else:
        cpd_items = [(cpd.variable, cpd) for cpd in result]

    saved = []
    for node, cpd in cpd_items:
        _print_cpd_summary(cpd)
        fpath = os.path.join(output_dir, f"{node}.json")
        _save_cpd(cpd, fpath)
        saved.append(fpath)

    print(f"\nSaved {len(saved)} CPD file(s) to: {output_dir}")
    for p in saved:
        print(f"  {p}")


def cmd_compare(args):
    """``deepparameters compare`` — compare two saved CPDs and print metrics."""
    from deepparameters import evaluate_cpd_performance, print_comparison_metrics

    if not os.path.isfile(args.learned):
        _die(f"Learned CPD file not found: {args.learned}")
    if not os.path.isfile(args.true):
        _die(f"True CPD file not found: {args.true}")

    learned_cpd = _load_cpd(args.learned)
    true_cpd = _load_cpd(args.true)

    print(f"\nLearned CPD : {args.learned}")
    print(f"True CPD    : {args.true}")
    print()

    try:
        metrics = evaluate_cpd_performance(learned_cpd, true_cpd)
    except Exception as exc:
        _die(f"Comparison failed: {exc}")

    print_comparison_metrics(metrics)

    if args.output:
        with open(args.output, "w") as fh:
            json.dump({k: float(v) for k, v in metrics.items()}, fh, indent=2)
        print(f"\nMetrics saved to: {args.output}")


# ---------------------------------------------------------------------------
# Argument parser
# ---------------------------------------------------------------------------

def _build_parser():
    parser = argparse.ArgumentParser(
        prog="deepparameters",
        description=(
            "DeepParameters CLI — learn Bayesian network CPDs from CSV data.\n"
            "Run 'deepparameters <command> --help' for command-specific help."
        ),
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=(
            "Examples:\n"
            "  deepparameters info\n"
            "  deepparameters learn --node B --data data.csv --edges edges.csv "
            "--num-parameters 20 --output cpd_B.json\n"
            "  deepparameters learn-network --data data.csv --edges edges.csv "
            "--output-dir ./cpds/\n"
            "  deepparameters compare --learned cpd_B.json --true cpd_true.json"
        ),
    )
    subparsers = parser.add_subparsers(dest="command", metavar="<command>")
    subparsers.required = True

    # ---- info ---------------------------------------------------------------
    subparsers.add_parser(
        "info",
        help="Print version and available architectures / sampling methods.",
    )

    # ---- learn --------------------------------------------------------------
    p_learn = subparsers.add_parser(
        "learn",
        help="Learn a CPD for a single node.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    # required
    p_learn.add_argument("--node", required=True,
                         help="Name of the node to learn (must be a column in --data).")
    p_learn.add_argument("--data", required=True, metavar="CSV",
                         help="Path to the CSV data file.")
    p_learn.add_argument("--edges", required=True, metavar="CSV",
                         help="Path to edges CSV (columns: parent, child).")
    p_learn.add_argument("--num-parameters", type=int, required=True,
                         help="Neural network capacity (typical range 5–100).")
    # optional — model
    p_learn.add_argument("--network-type", default="simple",
                         choices=["simple","advanced","ultra","mega","lstm",
                                  "bnn","vae","autoencoder","normalizing_flow"],
                         help="Neural network architecture.")
    p_learn.add_argument("--sampling-method", default="4", metavar="1-8",
                         help="CPD refinement method (1=Gibbs … 8=Adaptive KDE).")
    # optional — training
    p_learn.add_argument("--epochs", type=int, default=100)
    p_learn.add_argument("--batch-size", type=int, default=32)
    p_learn.add_argument("--learning-rate", type=float, default=0.001)
    p_learn.add_argument("--validation-split", type=float, default=0.2)
    p_learn.add_argument("--optimizer", default="adam",
                         choices=["adam","adamw","sgd","rmsprop","nadam"])
    p_learn.add_argument("--early-stopping-patience", type=int, default=10)
    p_learn.add_argument("--no-early-stopping", dest="early_stopping",
                         action="store_false", default=True,
                         help="Disable early stopping.")
    # optional — misc
    p_learn.add_argument("--verbose", action="store_true", default=False)
    p_learn.add_argument("--random-state", type=int, default=42)
    p_learn.add_argument("--output", metavar="PATH",
                         help="Save learned CPD to this JSON file.")

    # ---- learn-network ------------------------------------------------------
    p_net = subparsers.add_parser(
        "learn-network",
        help="Learn CPDs for all nodes in a network (parallel).",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    p_net.add_argument("--data", required=True, metavar="CSV")
    p_net.add_argument("--edges", required=True, metavar="CSV")
    p_net.add_argument("--network-type", default="simple",
                       choices=["simple","advanced","ultra","mega","lstm",
                                "bnn","vae","autoencoder","normalizing_flow"])
    p_net.add_argument("--sampling-method", default="4", metavar="1-8")
    p_net.add_argument("--parallel-style", default="topological",
                       choices=["topological","parent_child"])
    p_net.add_argument("--max-workers", type=int, default=None,
                       metavar="N", help="Worker pool size (default: CPU count).")
    p_net.add_argument("--epochs", type=int, default=100)
    p_net.add_argument("--verbose", action="store_true", default=False)
    p_net.add_argument("--output-dir", metavar="DIR", default=None,
                       help="Directory to save per-node CPD JSON files (default: cwd).")

    # ---- compare ------------------------------------------------------------
    p_cmp = subparsers.add_parser(
        "compare",
        help="Compare a learned CPD JSON against a true CPD JSON.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    p_cmp.add_argument("--learned", required=True, metavar="JSON",
                       help="Path to the learned CPD JSON (from 'learn --output').")
    p_cmp.add_argument("--true", required=True, dest="true", metavar="JSON",
                       help="Path to the true/reference CPD JSON.")
    p_cmp.add_argument("--output", metavar="JSON",
                       help="Optionally save metrics dict to this JSON file.")

    return parser


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------

def main():
    parser = _build_parser()
    args = parser.parse_args()

    dispatch = {
        "info":            cmd_info,
        "learn":           cmd_learn,
        "learn-network":   cmd_learn_network,
        "compare":         cmd_compare,
    }

    handler = dispatch.get(args.command)
    if handler is None:
        parser.print_help()
        sys.exit(1)

    try:
        handler(args)
    except SystemExit:
        raise
    except KeyboardInterrupt:
        print("\nInterrupted.", file=sys.stderr)
        sys.exit(130)
    except Exception as exc:  # pragma: no cover
        print(f"\nERROR: {exc}", file=sys.stderr)
        import traceback
        traceback.print_exc(file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()
