"""
LINE Model JSON Save/Load

Provides save_model() and load_model() functions for serializing LINE
Network and LayeredNetwork models to/from JSON format conforming to
the line-model.schema.json specification.

Copyright (c) 2012-2026, Imperial College London
All rights reserved.
"""

import json
import numpy as np
from typing import Any, Dict, Optional, Union


# ---------------------------------------------------------------------------
# Distribution serialization
# ---------------------------------------------------------------------------

def _dist_to_json(dist) -> Optional[Dict[str, Any]]:
    """Convert a LINE distribution object to JSON-compatible dict."""
    if dist is None:
        return None

    from ..distributions.continuous import (
        Exp, Det, Erlang, HyperExp, Gamma, Lognormal, Uniform, Immediate, Disabled, Pareto
    )
    from ..distributions.discrete import Zipf, DiscreteSampler
    from ..distributions.markovian import PH, APH, Coxian, MAP

    if isinstance(dist, Disabled):
        return {"type": "Disabled"}
    if isinstance(dist, Immediate):
        return {"type": "Immediate"}
    if isinstance(dist, Exp):
        return {"type": "Exp", "params": {"lambda": dist._rate}}
    if isinstance(dist, Det):
        return {"type": "Det", "params": {"value": dist._value}}
    if isinstance(dist, Erlang):
        return {"type": "Erlang", "params": {"lambda": dist._phase_rate, "k": int(dist._phases)}}
    if isinstance(dist, HyperExp):
        return {"type": "HyperExp", "params": {
            "p": _to_list(dist._probs), "lambda": _to_list(dist._rates)
        }}
    if isinstance(dist, Gamma):
        return {"type": "Gamma", "params": {"alpha": dist._shape, "beta": dist._rate}}
    if isinstance(dist, Lognormal):
        return {"type": "Lognormal", "params": {"mu": dist._mu, "sigma": dist._sigma}}
    if isinstance(dist, Uniform):
        return {"type": "Uniform", "params": {"a": dist._a, "b": dist._b}}
    if isinstance(dist, Zipf):
        return {"type": "Zipf", "params": {"s": dist._skew, "n": int(dist._n)}}
    if isinstance(dist, Pareto):
        return {"type": "Pareto", "params": {"alpha": dist._alpha, "scale": dist._scale}}
    if isinstance(dist, DiscreteSampler):
        return {"type": "DiscreteSampler", "params": {
            "p": _to_list(dist._probs), "x": _to_list(dist._values)
        }}
    if isinstance(dist, MAP):
        return {"type": "MAP", "map": {
            "D0": _matrix_to_list(dist._D0),
            "D1": _matrix_to_list(dist._D1)
        }}
    # APH/Coxian before PH (subclass check)
    if isinstance(dist, (APH, Coxian)):
        return {"type": "PH", "ph": {
            "alpha": _to_list(dist._alpha),
            "T": _matrix_to_list(dist._T)
        }}
    if isinstance(dist, PH):
        return {"type": "PH", "ph": {
            "alpha": _to_list(dist._alpha),
            "T": _matrix_to_list(dist._T)
        }}

    # Fallback: use fitMean with mean
    try:
        mean = dist.getMean()
        return {"type": "Exp", "fit": {"method": "fitMean", "mean": mean}}
    except Exception:
        return None


def _json_to_dist(d: Dict[str, Any]):
    """Convert a JSON dist dict back to a LINE distribution object."""
    from ..distributions.continuous import (
        Exp, Det, Erlang, HyperExp, Gamma, Lognormal, Uniform, Immediate, Disabled, Pareto
    )
    from ..distributions.discrete import Zipf, DiscreteSampler
    from ..distributions.markovian import PH, APH, MAP

    dtype = d.get("type")
    params = d.get("params", {})
    fit = d.get("fit")
    ph = d.get("ph")
    mapspec = d.get("map")

    if dtype == "Disabled":
        return Disabled.getInstance()
    if dtype == "Immediate":
        return Immediate.getInstance()

    # Direct params
    if params:
        if dtype == "Exp":
            rate = params.get("lambda", params.get("rate"))
            return Exp(rate)
        if dtype == "Det":
            return Det(params["value"])
        if dtype == "Erlang":
            return Erlang(params["lambda"], int(params["k"]))
        if dtype == "HyperExp":
            p = params["p"]
            lam = params["lambda"]
            if len(p) == 2:
                return HyperExp(p[0], lam[0], lam[1])
            return HyperExp(np.array(p), np.array(lam))
        if dtype == "Gamma":
            return Gamma(params["alpha"], params["beta"])
        if dtype == "Lognormal":
            return Lognormal(params["mu"], params["sigma"])
        if dtype == "Uniform":
            return Uniform(params["a"], params["b"])
        if dtype == "Zipf":
            return Zipf(params["s"], int(params["n"]))
        if dtype == "Pareto":
            return Pareto(params["alpha"], params["scale"])
        if dtype == "DiscreteSampler":
            pv = np.array(params["p"])
            xv = np.array(params["x"])
            return DiscreteSampler(xv, pv)

    # PH representation
    if ph:
        alpha = np.array(ph["alpha"])
        T = np.array(ph["T"])
        return PH(alpha, T)

    # MAP representation
    if mapspec:
        D0 = np.array(mapspec["D0"])
        D1 = np.array(mapspec["D1"])
        return MAP(D0, D1)

    # Fit specification
    if fit:
        method = fit["method"]
        if method == "fitMean":
            mean = fit["mean"]
            if dtype == "Exp":
                return Exp.fit_mean(mean)
            if dtype == "Erlang":
                order = fit.get("order", 1)
                return Erlang.fit_mean_and_order(mean, order)
            if dtype == "Det":
                return Det(mean)
            # Default: Exp with given mean
            return Exp.fit_mean(mean)
        if method == "fitMeanAndSCV":
            mean = fit["mean"]
            scv = fit["scv"]
            if dtype == "Erlang":
                return Erlang.fit_mean_and_scv(mean, scv)
            if dtype == "HyperExp":
                return HyperExp.fit_mean_and_scv(mean, scv)
            if dtype == "Gamma":
                return Gamma.fit_mean_and_scv(mean, scv)
            if dtype == "Lognormal":
                return Lognormal.fit_mean_and_scv(mean, scv)
            return Exp.fit_mean(mean)
        if method == "fitMeanAndOrder":
            mean = fit["mean"]
            order = fit["order"]
            if dtype == "Erlang":
                return Erlang.fit_mean_and_order(mean, order)
            return Exp.fit_mean(mean)

    # Fallback
    return Exp(1.0)


# ---------------------------------------------------------------------------
# Network serialization
# ---------------------------------------------------------------------------

def _network_to_json(model) -> Dict[str, Any]:
    """Convert a Network model to JSON-compatible dict."""
    from ..lang.nodes import Queue, Delay, Source, Sink, Fork, Join, Router, ClassSwitch, Cache, Place, Transition
    from ..lang.classes import OpenClass, ClosedClass
    from ..lang.base import SchedStrategy, NodeType

    nodes = model.get_nodes()
    classes = model.get_classes()

    # Build nodes array
    nodes_json = []
    for node in nodes:
        # Skip implicit ClassSwitch nodes (auto-created by link())
        if isinstance(node, ClassSwitch) and getattr(node, '_auto_added', False):
            continue

        nj = {"name": node.name, "type": _node_type_str(node)}

        # Scheduling
        if isinstance(node, (Queue, Delay)):
            sched = node.get_sched_strategy()
            if sched is not None and sched != SchedStrategy.INF:
                nj["scheduling"] = sched.name
            elif isinstance(node, Delay):
                nj["scheduling"] = "INF"

        # Servers
        if isinstance(node, Queue) and not isinstance(node, Delay):
            ns = getattr(node, '_number_of_servers', 1)
            if ns is not None and np.isfinite(ns) and int(ns) != 1:
                nj["servers"] = int(ns)

        # Buffer
        if isinstance(node, Queue):
            cap = getattr(node, '_capacity', None)
            if cap is not None and np.isfinite(cap):
                nj["buffer"] = int(cap)

        # Service/arrival distributions
        if isinstance(node, Source):
            svc = {}
            for jc in classes:
                d = node._arrival_process.get(jc)
                if d is not None:
                    dj = _dist_to_json(d)
                    if dj is not None:
                        svc[jc.name] = dj
            if svc:
                nj["service"] = svc
        elif isinstance(node, (Queue, Delay)):
            svc = {}
            for jc in classes:
                d = node._service_process.get(jc)
                if d is not None:
                    dj = _dist_to_json(d)
                    if dj is not None:
                        svc[jc.name] = dj
            if svc:
                nj["service"] = svc

        # Scheduling params (DPS weights, etc.)
        if isinstance(node, Queue) and hasattr(node, '_sched_param') and node._sched_param:
            sp = {}
            for jc, val in node._sched_param.items():
                if val is not None:
                    sp[jc.name] = float(val)
            if sp:
                nj["schedParams"] = sp

        # ClassSwitch matrix
        if isinstance(node, ClassSwitch):
            csm = node.get_class_switching_matrix()
            if csm is not None:
                cs_dict = {}
                for ri, rc in enumerate(classes):
                    row = {}
                    for ci, cc in enumerate(classes):
                        if ri < csm.shape[0] and ci < csm.shape[1] and csm[ri, ci] != 0:
                            row[cc.name] = float(csm[ri, ci])
                    if row:
                        cs_dict[rc.name] = row
                if cs_dict:
                    nj["classSwitchMatrix"] = cs_dict

        # Join strategy
        if isinstance(node, Join):
            js = getattr(node, '_join_strategy', None)
            if js is not None:
                nj["joinStrategy"] = js.name if hasattr(js, 'name') else str(js)

        # Cache config
        if isinstance(node, Cache):
            cc = {}
            cc["items"] = int(node._num_items)
            cap = node._item_level_cap
            if isinstance(cap, np.ndarray):
                cc["capacity"] = [int(x) for x in cap]
            else:
                cc["capacity"] = int(cap)
            rs = node._replacement_strategy
            cc["replacement"] = rs.name if hasattr(rs, 'name') else str(rs)

            # Hit/miss class mappings
            if node._hit_class:
                hit_map = {}
                for in_cls, out_cls in node._hit_class.items():
                    hit_map[in_cls.name] = out_cls.name
                if hit_map:
                    cc["hitClass"] = hit_map
            if node._miss_class:
                miss_map = {}
                for in_cls, out_cls in node._miss_class.items():
                    miss_map[in_cls.name] = out_cls.name
                if miss_map:
                    cc["missClass"] = miss_map

            # Read popularity distributions (set_read)
            if node._read_process:
                pop_map = {}
                for jc, dist in node._read_process.items():
                    dj = _dist_to_json(dist)
                    if dj is not None:
                        pop_map[jc.name] = dj
                if pop_map:
                    cc["popularity"] = pop_map

            nj["cache"] = cc

        # Fork tasksPerLink
        if isinstance(node, Fork):
            tpl = getattr(node, '_tasks_per_link', None)
            if tpl is not None and tpl > 1:
                nj["tasksPerLink"] = int(tpl)

        # Join paired fork
        if isinstance(node, Join):
            fork_ref = getattr(node, '_fork', None)
            if fork_ref is not None and hasattr(fork_ref, 'name'):
                nj["forkNode"] = fork_ref.name

        # Transition modes
        if isinstance(node, Transition):
            modes_json = []
            n_modes = node.get_number_of_modes()
            all_nodes = model.get_nodes()
            for mi_idx in range(n_modes):
                mj = {}
                if mi_idx < len(node._mode_names):
                    mj["name"] = node._mode_names[mi_idx]
                else:
                    mj["name"] = f"Mode{mi_idx + 1}"
                # Distribution
                if mi_idx < len(node._distributions) and node._distributions[mi_idx] is not None:
                    dj = _dist_to_json(node._distributions[mi_idx])
                    if dj is not None:
                        mj["distribution"] = dj
                # Timing strategy
                if mi_idx < len(node._timing_strategies):
                    from ..lang.base import TimingStrategy
                    if node._timing_strategies[mi_idx] == TimingStrategy.IMMEDIATE:
                        mj["timingStrategy"] = "IMMEDIATE"
                    else:
                        mj["timingStrategy"] = "TIMED"
                # Number of servers
                if mi_idx < len(node._number_of_servers) and node._number_of_servers[mi_idx] > 1:
                    mj["numServers"] = int(node._number_of_servers[mi_idx])
                # Firing priority
                if mi_idx < len(node._firing_priorities) and node._firing_priorities[mi_idx] > 0:
                    mj["firingPriority"] = float(node._firing_priorities[mi_idx])
                # Firing weight
                if mi_idx < len(node._firing_weights) and node._firing_weights[mi_idx] != 1.0:
                    mj["firingWeight"] = float(node._firing_weights[mi_idx])
                # Enabling conditions
                if mi_idx < len(node._enabling_conditions):
                    ec_mat = node._enabling_conditions[mi_idx]
                    ec_list = []
                    for ni in range(ec_mat.shape[0]):
                        for ci in range(ec_mat.shape[1]):
                            if ec_mat[ni, ci] > 0:
                                ec_list.append({
                                    "node": all_nodes[ni].name,
                                    "class": classes[ci].name,
                                    "count": float(ec_mat[ni, ci])
                                })
                    if ec_list:
                        mj["enablingConditions"] = ec_list
                # Inhibiting conditions
                if mi_idx < len(node._inhibiting_conditions):
                    ic_mat = node._inhibiting_conditions[mi_idx]
                    ic_list = []
                    for ni in range(ic_mat.shape[0]):
                        for ci in range(ic_mat.shape[1]):
                            if np.isfinite(ic_mat[ni, ci]):
                                ic_list.append({
                                    "node": all_nodes[ni].name,
                                    "class": classes[ci].name,
                                    "count": float(ic_mat[ni, ci])
                                })
                    if ic_list:
                        mj["inhibitingConditions"] = ic_list
                # Firing outcomes
                if mi_idx < len(node._firing_outcomes):
                    fo_mat = node._firing_outcomes[mi_idx]
                    fo_list = []
                    for ni in range(fo_mat.shape[0]):
                        for ci in range(fo_mat.shape[1]):
                            if fo_mat[ni, ci] != 0:
                                fo_list.append({
                                    "node": all_nodes[ni].name,
                                    "class": classes[ci].name,
                                    "count": float(fo_mat[ni, ci])
                                })
                    if fo_list:
                        mj["firingOutcomes"] = fo_list
                modes_json.append(mj)
            if modes_json:
                nj["modes"] = modes_json

        nodes_json.append(nj)

    # Build classes array
    classes_json = []
    for jc in classes:
        cj = {"name": jc.name}
        if isinstance(jc, OpenClass):
            cj["type"] = "Open"
        elif isinstance(jc, ClosedClass):
            cj["type"] = "Closed"
            cj["population"] = int(jc._njobs)
            ref = jc._refstat
            if ref is not None:
                cj["refNode"] = ref.name
        else:
            cj["type"] = "Open"

        prio = getattr(jc, '_priority', 0)
        if prio != 0:
            cj["priority"] = int(prio)

        classes_json.append(cj)

    # Build routing matrix
    routing_json = _build_routing_json(model, nodes, classes)

    result = {
        "type": "Network",
        "name": model.name,
        "nodes": nodes_json,
        "classes": classes_json,
        "routing": routing_json
    }

    return result


def _build_routing_json(model, nodes, classes) -> Dict[str, Any]:
    """Build routing JSON from the model's routing matrix."""
    matrix = {}

    # Try to get the routing from the internal _routing_matrix
    # Prefer _original_routes (pre-ClassSwitch expansion) over _routes
    rm = getattr(model, '_routing_matrix', None)
    if rm is not None and hasattr(rm, '_routes'):
        routes = rm._routes
        original_routes = getattr(rm, '_original_routes', None)
        if original_routes is not None:
            routes = original_routes
        for (cs, cd), route_dict in routes.items():
            key = f"{cs.name},{cd.name}"
            from_to = {}
            for (ns, nd), prob in route_dict.items():
                if prob > 0:
                    if ns.name not in from_to:
                        from_to[ns.name] = {}
                    from_to[ns.name][nd.name] = prob
            if from_to:
                matrix[key] = from_to

    # If the _routes approach didn't work, try the struct approach
    if not matrix:
        try:
            sn = model.getStruct()
            if sn is not None and hasattr(sn, 'rtnodes') and sn.rtnodes is not None:
                rt = sn.rtnodes
                K = len(classes)
                N = len(nodes)
                for r in range(K):
                    for s in range(K):
                        from_to = {}
                        for i in range(N):
                            for j in range(N):
                                val = rt[i * K + r, j * K + s]
                                if val > 1e-14:
                                    ni = nodes[i].name
                                    nj = nodes[j].name
                                    if ni not in from_to:
                                        from_to[ni] = {}
                                    from_to[ni][nj] = float(val)
                        if from_to:
                            key = f"{classes[r].name},{classes[s].name}"
                            matrix[key] = from_to
        except Exception:
            pass

    return {"type": "matrix", "matrix": matrix}


def _json_to_network(data: Dict[str, Any]):
    """Reconstruct a Network from JSON data."""
    from ..lang.network import Network
    from ..lang.nodes import Queue, Delay, Source, Sink, Fork, Join, Router, ClassSwitch, Cache, Place, Transition
    from ..lang.classes import OpenClass, ClosedClass
    from ..lang.routing import RoutingMatrix
    from ..lang.base import SchedStrategy, ReplacementStrategy

    model = Network(data.get("name", "model"))

    # Create nodes
    node_map = {}
    for nd in data.get("nodes", []):
        name = nd["name"]
        ntype = nd["type"]
        node = _create_node(model, nd)
        node_map[name] = node

    # Deferred Join fork linking
    for nd in data.get("nodes", []):
        if nd.get("type") == "Join" and "forkNode" in nd:
            join_node = node_map.get(nd["name"])
            fork_node = node_map.get(nd["forkNode"])
            if join_node is not None and fork_node is not None:
                join_node._fork = fork_node

    # Create classes
    class_map = {}
    for cd in data.get("classes", []):
        cname = cd["name"]
        ctype = cd["type"]
        if ctype == "Open":
            prio = cd.get("priority", 0)
            jc = OpenClass(model, cname, prio)
        elif ctype == "Closed":
            pop = cd["population"]
            ref_name = cd.get("refNode")
            ref_node = node_map.get(ref_name) if ref_name else None
            prio = cd.get("priority", 0)
            jc = ClosedClass(model, cname, pop, ref_node, prio)
        else:
            jc = OpenClass(model, cname)
        class_map[cname] = jc

    # Set service/arrival distributions
    for nd in data.get("nodes", []):
        name = nd["name"]
        node = node_map[name]
        svc = nd.get("service", {})
        for cname, dist_json in svc.items():
            jc = class_map.get(cname)
            if jc is None:
                continue
            dist = _json_to_dist(dist_json)
            if dist is not None:
                if isinstance(node, Source):
                    node.set_arrival(jc, dist)
                elif isinstance(node, (Queue, Delay)):
                    node.set_service(jc, dist)

        # Set scheduling params
        sp = nd.get("schedParams", {})
        if sp and isinstance(node, Queue):
            for cname, val in sp.items():
                jc = class_map.get(cname)
                if jc is not None:
                    node._sched_param[jc] = val

        # Set class switch matrix (dict format: "classSwitchMatrix")
        csm = nd.get("classSwitchMatrix")
        if csm and isinstance(node, ClassSwitch):
            classes_list = model.get_classes()
            K = len(classes_list)
            mat = np.zeros((K, K))
            class_idx = {c.name: i for i, c in enumerate(classes_list)}
            for from_name, to_dict in csm.items():
                ri = class_idx.get(from_name, -1)
                if ri < 0:
                    continue
                for to_name, prob in to_dict.items():
                    ci = class_idx.get(to_name, -1)
                    if ci >= 0:
                        mat[ri, ci] = prob
            node.set_class_switching_matrix(mat)
        # Legacy 2D array format: "csMatrix" (from older JAR saves)
        elif not csm and isinstance(node, ClassSwitch):
            cs_arr = nd.get("csMatrix")
            if cs_arr and isinstance(cs_arr, list):
                mat = np.array(cs_arr, dtype=float)
                node.set_class_switching_matrix(mat)

        # Cache hit/miss class mappings and popularity distributions
        if isinstance(node, Cache) and "cache" in nd:
            cc = nd["cache"]
            # Hit class mapping
            hc = cc.get("hitClass", {})
            for in_name, out_name in hc.items():
                in_cls = class_map.get(in_name)
                out_cls = class_map.get(out_name)
                if in_cls is not None and out_cls is not None:
                    node.set_hit_class(in_cls, out_cls)
            # Miss class mapping
            mc = cc.get("missClass", {})
            for in_name, out_name in mc.items():
                in_cls = class_map.get(in_name)
                out_cls = class_map.get(out_name)
                if in_cls is not None and out_cls is not None:
                    node.set_miss_class(in_cls, out_cls)
            # Popularity distributions (set_read)
            pop = cc.get("popularity", {})
            for cname, dist_json in pop.items():
                jc = class_map.get(cname)
                if jc is not None:
                    pop_dist = _json_to_dist(dist_json)
                    if pop_dist is not None:
                        node.set_read(jc, pop_dist)

    # Configure Transition modes
    for nd in data.get("nodes", []):
        if nd.get("type") != "Transition" or "modes" not in nd:
            continue
        tnode = node_map[nd["name"]]
        for md in nd["modes"]:
            mode_name = md.get("name", "Mode")
            mode = tnode.add_mode(mode_name)
            # Distribution
            dist_data = md.get("distribution")
            if dist_data:
                dist = _json_to_dist(dist_data)
                if dist is not None:
                    tnode.set_distribution(mode, dist)
            # Timing strategy
            ts = md.get("timingStrategy")
            if ts:
                from ..lang.base import TimingStrategy
                if ts == "IMMEDIATE":
                    tnode.set_timing_strategy(mode, TimingStrategy.IMMEDIATE)
                else:
                    tnode.set_timing_strategy(mode, TimingStrategy.TIMED)
            # Number of servers
            ns = md.get("numServers")
            if ns is not None and ns > 1:
                tnode.set_number_of_servers(mode, ns)
            # Firing priority
            fp = md.get("firingPriority")
            if fp is not None:
                tnode.set_firing_priorities(mode, fp)
            # Firing weight
            fw = md.get("firingWeight")
            if fw is not None:
                tnode.set_firing_weights(mode, fw)
            # Enabling conditions
            for ec in md.get("enablingConditions", []):
                ec_node = node_map.get(ec["node"])
                ec_cls = class_map.get(ec["class"])
                if ec_node is not None and ec_cls is not None:
                    tnode.set_enabling_conditions(mode, ec_cls, ec_node, ec["count"])
            # Inhibiting conditions
            for ic in md.get("inhibitingConditions", []):
                ic_node = node_map.get(ic["node"])
                ic_cls = class_map.get(ic["class"])
                if ic_node is not None and ic_cls is not None:
                    tnode.set_inhibiting_conditions(mode, ic_cls, ic_node, ic["count"])
            # Firing outcomes
            for fo in md.get("firingOutcomes", []):
                fo_node = node_map.get(fo["node"])
                fo_cls = class_map.get(fo["class"])
                if fo_node is not None and fo_cls is not None:
                    tnode.set_firing_outcome(mode, fo_cls, fo_node, fo["count"])

    # Build routing
    routing_data = data.get("routing", {})
    if routing_data.get("type") == "matrix":
        P = RoutingMatrix(model)
        matrix = routing_data.get("matrix", {})
        for key, from_to in matrix.items():
            parts = key.split(",")
            if len(parts) != 2:
                continue
            cs_name, cd_name = parts[0].strip(), parts[1].strip()
            cs = class_map.get(cs_name)
            cd = class_map.get(cd_name)
            if cs is None or cd is None:
                continue
            for from_name, to_dict in from_to.items():
                ns = node_map.get(from_name)
                if ns is None:
                    continue
                for to_name, prob in to_dict.items():
                    nd = node_map.get(to_name)
                    if nd is not None:
                        P.set(cs, cd, ns, nd, prob)
        model.link(P)

    return model


def _create_node(model, nd: Dict[str, Any]):
    """Create a node from JSON data (before classes are set up)."""
    from ..lang.nodes import Queue, Delay, Source, Sink, Fork, Join, Router, ClassSwitch, Cache, Place, Transition
    from ..lang.base import SchedStrategy, ReplacementStrategy

    name = nd["name"]
    ntype = nd["type"]

    if ntype == "Source":
        return Source(model, name)
    elif ntype == "Sink":
        return Sink(model, name)
    elif ntype == "Delay":
        return Delay(model, name)
    elif ntype == "Queue":
        sched_str = nd.get("scheduling", "FCFS")
        sched = getattr(SchedStrategy, sched_str, SchedStrategy.FCFS)
        node = Queue(model, name, sched)
        servers = nd.get("servers", 1)
        if servers > 1:
            node.set_number_of_servers(servers)
        buf = nd.get("buffer")
        if buf is not None:
            node._capacity = buf
        return node
    elif ntype == "Fork":
        node = Fork(model, name)
        tpl = nd.get("tasksPerLink")
        if tpl is not None and tpl > 1:
            node.set_tasks_per_link(tpl)
        return node
    elif ntype == "Join":
        node = Join(model, name)
        js = nd.get("joinStrategy")
        if js is not None:
            from ..lang.base import JoinStrategy
            node._join_strategy = getattr(JoinStrategy, js, JoinStrategy.STD)
        return node
    elif ntype == "Router":
        return Router(model, name)
    elif ntype == "ClassSwitch":
        return ClassSwitch(model, name)
    elif ntype == "Cache":
        cc = nd.get("cache", {})
        nitems = cc.get("items", 10)
        cap = cc.get("capacity", 1)
        if isinstance(cap, list):
            cap = np.array(cap)
        repl_str = cc.get("replacement", "LRU")
        repl = getattr(ReplacementStrategy, repl_str, ReplacementStrategy.LRU)
        return Cache(model, name, nitems, cap, repl)
    elif ntype == "Place":
        return Place(model, name)
    elif ntype == "Transition":
        return Transition(model, name)
    else:
        # Default to Queue FCFS
        return Queue(model, name, SchedStrategy.FCFS)


def _node_type_str(node) -> str:
    """Get the schema node type string for a node object."""
    from ..lang.nodes import Queue, Delay, Source, Sink, Fork, Join, Router, ClassSwitch, Cache, Place, Transition
    if isinstance(node, Source):
        return "Source"
    if isinstance(node, Sink):
        return "Sink"
    if isinstance(node, Delay):
        return "Delay"
    if isinstance(node, Cache):
        return "Cache"
    if isinstance(node, Place):
        return "Place"
    if isinstance(node, Transition):
        return "Transition"
    if isinstance(node, Queue):
        return "Queue"
    if isinstance(node, Fork):
        return "Fork"
    if isinstance(node, Join):
        return "Join"
    if isinstance(node, Router):
        return "Router"
    if isinstance(node, ClassSwitch):
        return "ClassSwitch"
    return "Queue"


# ---------------------------------------------------------------------------
# LayeredNetwork serialization
# ---------------------------------------------------------------------------

def _layered_to_json(model) -> Dict[str, Any]:
    """Convert a LayeredNetwork to JSON-compatible dict."""
    from ..layered import (
        LayeredNetwork, Processor, Task, Entry, Activity,
        ActivityPrecedence, PrecedenceType, CallType, Distribution
    )
    from ..constants import SchedStrategy

    result = {
        "type": "LayeredNetwork",
        "name": model.name,
    }

    # Processors
    procs = []
    for p in model.processors:
        pj = {"name": p.name}
        if p.multiplicity is not None and np.isfinite(p.multiplicity) and int(p.multiplicity) != 1:
            pj["multiplicity"] = int(p.multiplicity)
        sched = p.sched_strategy
        if sched is not None:
            sname = sched.name if hasattr(sched, 'name') else str(sched)
            if sname != "INF":
                pj["scheduling"] = sname
        q = p.getQuantum()
        if q > 0:
            pj["quantum"] = q
        sf = p.getSpeedFactor()
        if sf != 1.0:
            pj["speedFactor"] = sf
        repl = p.getReplication()
        if repl > 1:
            pj["replication"] = repl
        procs.append(pj)
    result["processors"] = procs

    # Tasks
    tasks = []
    for t in model.tasks:
        tj = {"name": t.name}
        # Processor
        if t.processor is not None:
            tj["processor"] = t.processor.name
        # Multiplicity
        mult = t.multiplicity
        if mult is not None and np.isfinite(mult) and int(mult) != 1:
            tj["multiplicity"] = int(mult)
        # Scheduling
        sched = t.sched_strategy
        if sched is not None:
            sname = sched.name if hasattr(sched, 'name') else str(sched)
            tj["scheduling"] = sname
        # Think time
        if t.think_time is not None:
            tt_mean = _get_dist_mean_safe(t.think_time)
            if tt_mean > 0:
                tj["thinkTime"] = _layered_dist_to_json(t.think_time)
        # Fan in/out
        if t._fan_in:
            tj["fanIn"] = dict(t._fan_in)
        if t._fan_out:
            tj["fanOut"] = dict(t._fan_out)
        # Replication
        repl = t.getReplication()
        if repl > 1:
            tj["replication"] = repl
        tasks.append(tj)
    result["tasks"] = tasks

    # Entries
    entries = []
    for e in model.entries:
        ej = {"name": e.name}
        if e.task is not None:
            ej["task"] = e.task.name
        entries.append(ej)
    result["entries"] = entries

    # Activities
    activities = []
    for a in model.activities:
        aj = {"name": a.name}
        if a.task is not None:
            aj["task"] = a.task.name
        # Host demand
        if a.host_demand is not None:
            aj["hostDemand"] = _layered_dist_to_json(a.host_demand)
        # Bound to entry
        if a.bound_entry is not None:
            aj["boundTo"] = a.bound_entry.name
        # Replies to entry
        if a.reply_entry is not None:
            aj["repliesTo"] = a.reply_entry.name
        # Calls
        synch_calls = []
        asynch_calls = []
        for entry, mean_calls, call_type in a.calls:
            call_obj = {"entry": entry.name}
            if mean_calls != 1.0:
                call_obj["mean"] = mean_calls
            if call_type == CallType.SYNC:
                synch_calls.append(call_obj)
            elif call_type == CallType.ASYNC:
                asynch_calls.append(call_obj)
        if synch_calls:
            aj["synchCalls"] = synch_calls
        if asynch_calls:
            aj["asynchCalls"] = asynch_calls
        activities.append(aj)
    result["activities"] = activities

    # Precedences
    precedences = []
    for t in model.tasks:
        for p in t.precedences:
            pj = {"task": t.name}
            if p.prec_type == PrecedenceType.SERIAL:
                pj["type"] = "Serial"
                pj["activities"] = [a.name for a in p.activities]
            elif p.prec_type == PrecedenceType.PARALLEL:
                if len(p.pre_activities) == 1 and len(p.post_activities) > 1:
                    pj["type"] = "AndFork"
                    pj["activities"] = [p.pre_activities[0].name] + [a.name for a in p.post_activities]
                elif len(p.pre_activities) > 1 and len(p.post_activities) == 1:
                    pj["type"] = "AndJoin"
                    pj["activities"] = [a.name for a in p.pre_activities] + [p.post_activities[0].name]
                else:
                    pj["type"] = "Serial"
                    pj["activities"] = [a.name for a in (p.pre_activities + p.post_activities)]
            elif p.prec_type == PrecedenceType.CHOICE:
                if len(p.pre_activities) == 1 and len(p.post_activities) > 1:
                    pj["type"] = "OrFork"
                    pj["activities"] = [p.pre_activities[0].name] + [a.name for a in p.post_activities]
                    if p.probabilities:
                        pj["probabilities"] = list(p.probabilities)
                elif len(p.pre_activities) > 1 and len(p.post_activities) == 1:
                    pj["type"] = "OrJoin"
                    pj["activities"] = [a.name for a in p.pre_activities] + [p.post_activities[0].name]
                else:
                    pj["type"] = "Serial"
                    pj["activities"] = [a.name for a in (p.pre_activities + p.post_activities)]
            elif p.prec_type == PrecedenceType.LOOP:
                pj["type"] = "Loop"
                pj["activities"] = [a.name for a in p.activities]
                if p.count != 1.0:
                    pj["loopCount"] = p.count
            else:
                continue
            precedences.append(pj)
    if precedences:
        result["precedences"] = precedences

    return result


def _json_to_layered(data: Dict[str, Any]):
    """Reconstruct a LayeredNetwork from JSON data."""
    from ..layered import (
        LayeredNetwork, Processor, Task, Entry, Activity,
        ActivityPrecedence, PrecedenceType, CallType, Distribution as LDist
    )
    from ..constants import SchedStrategy

    model = LayeredNetwork(data.get("name", "model"))

    # Create processors
    proc_map = {}
    for pd in data.get("processors", []):
        name = pd["name"]
        mult = pd.get("multiplicity", 1)
        sched_str = pd.get("scheduling", "INF")
        sched = getattr(SchedStrategy, sched_str, SchedStrategy.INF)
        proc = Processor(model, name, mult, sched)
        q = pd.get("quantum", 0.0)
        if q > 0:
            proc.setQuantum(q)
        sf = pd.get("speedFactor", 1.0)
        if sf != 1.0:
            proc.setSpeedFactor(sf)
        repl = pd.get("replication")
        if repl is not None and repl > 1:
            proc.setReplication(repl)
        proc_map[name] = proc

    # Create tasks
    task_map = {}
    for td in data.get("tasks", []):
        name = td["name"]
        mult = td.get("multiplicity", 1)
        sched_str = td.get("scheduling", "INF")
        sched = getattr(SchedStrategy, sched_str, SchedStrategy.INF)
        proc_name = td.get("processor")
        proc = proc_map.get(proc_name) if proc_name else None
        task = Task(model, name, mult, sched)
        if proc is not None:
            task.on(proc)
        # Think time
        tt = td.get("thinkTime")
        if tt is not None:
            task.set_think_time(_json_to_layered_dist(tt))
        # Fan in/out
        fi = td.get("fanIn", {})
        for src, val in fi.items():
            task.setFanIn(src, val)
        fo = td.get("fanOut", {})
        for dst, val in fo.items():
            task.setFanOut(dst, val)
        # Replication
        repl = td.get("replication")
        if repl is not None and repl > 1:
            task.setReplication(repl)
        task_map[name] = task

    # Create entries
    entry_map = {}
    for ed in data.get("entries", []):
        name = ed["name"]
        task_name = ed.get("task")
        task = task_map.get(task_name) if task_name else None
        entry = Entry(model, name)
        if task is not None:
            entry.on(task)
        entry_map[name] = entry

    # Create activities
    act_map = {}
    for ad in data.get("activities", []):
        name = ad["name"]
        task_name = ad.get("task")
        task = task_map.get(task_name) if task_name else None
        hd = ad.get("hostDemand")
        host_demand = _json_to_layered_dist(hd) if hd else None
        act = Activity(model, name, host_demand)
        if task is not None:
            act.on(task)
        # Bound to
        bt = ad.get("boundTo")
        if bt and bt in entry_map:
            act.bound_to(entry_map[bt])
        # Replies to
        rt = ad.get("repliesTo")
        if rt and rt in entry_map:
            act.replies_to(entry_map[rt])
        # Calls
        for sc in ad.get("synchCalls", []):
            ename = sc["entry"]
            mean = sc.get("mean", 1.0)
            if ename in entry_map:
                act.synch_call(entry_map[ename], mean)
        for ac in ad.get("asynchCalls", []):
            ename = ac["entry"]
            mean = ac.get("mean", 1.0)
            if ename in entry_map:
                act.asynch_call(entry_map[ename], mean)
        act_map[name] = act

    # Create precedences
    for pd in data.get("precedences", []):
        task_name = pd.get("task")
        task = task_map.get(task_name)
        if task is None:
            continue
        ptype = pd.get("type", "Serial")
        act_names = pd.get("activities", [])
        acts = [act_map[n] for n in act_names if n in act_map]

        if ptype == "Serial" and len(acts) >= 2:
            task.add_precedence(ActivityPrecedence.Serial(acts))
        elif ptype == "AndFork" and len(acts) >= 2:
            task.add_precedence(ActivityPrecedence.AndFork(acts[0], acts[1:]))
        elif ptype == "AndJoin" and len(acts) >= 2:
            task.add_precedence(ActivityPrecedence.AndJoin(acts[:-1], acts[-1]))
        elif ptype == "OrFork" and len(acts) >= 2:
            probs = pd.get("probabilities", [])
            if not probs:
                n = len(acts) - 1
                probs = [1.0 / n] * n
            task.add_precedence(ActivityPrecedence.OrFork(acts[0], acts[1:], probs))
        elif ptype == "OrJoin" and len(acts) >= 2:
            task.add_precedence(ActivityPrecedence.OrJoin(acts[:-1], acts[-1]))
        elif ptype == "Loop" and len(acts) >= 2:
            count = pd.get("loopCount", 1.0)
            task.add_precedence(ActivityPrecedence.Loop(acts[0], acts[1:-1] if len(acts) > 2 else [acts[1]], acts[-1] if len(acts) > 2 else acts[0], count))

    return model


def _layered_dist_to_json(dist) -> Dict[str, Any]:
    """Convert a layered network distribution (dataclass or native) to JSON."""
    from ..layered import Distribution as LDist
    if isinstance(dist, LDist):
        # Dataclass distribution — use Exp with fitMean
        if dist.scv == 0:
            return {"type": "Det", "params": {"value": dist.mean}}
        return {"type": "Exp", "params": {"lambda": 1.0 / dist.mean if dist.mean > 0 else 1e8}}
    # Try native distribution
    return _dist_to_json(dist) or {"type": "Exp", "params": {"lambda": 1.0}}


def _json_to_layered_dist(d: Dict[str, Any]):
    """Convert JSON dist to a layered Distribution (dataclass)."""
    from ..layered import Distribution as LDist
    # Use native distribution conversion, then get mean
    native = _json_to_dist(d)
    if native is not None:
        try:
            mean = native.getMean()
            scv = native.getSCV()
            return LDist(mean=mean, scv=scv)
        except Exception:
            return LDist(mean=1.0, scv=1.0)
    return LDist(mean=1.0, scv=1.0)


def _get_dist_mean_safe(dist) -> float:
    """Get mean from any distribution safely."""
    if dist is None:
        return 0.0
    if isinstance(dist, (int, float)):
        return float(dist)
    if hasattr(dist, 'mean') and not callable(dist.mean):
        return dist.mean
    if hasattr(dist, 'getMean'):
        return dist.getMean()
    return 0.0


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _to_list(arr) -> list:
    """Convert numpy array or list to plain Python list."""
    if isinstance(arr, np.ndarray):
        return arr.tolist()
    if isinstance(arr, (list, tuple)):
        return list(arr)
    return [arr]


def _matrix_to_list(mat) -> list:
    """Convert 2D numpy array to nested list."""
    if isinstance(mat, np.ndarray):
        return mat.tolist()
    return mat


class _JSONEncoder(json.JSONEncoder):
    """Custom JSON encoder that handles numpy types."""
    def default(self, obj):
        if isinstance(obj, np.integer):
            return int(obj)
        if isinstance(obj, np.floating):
            return float(obj)
        if isinstance(obj, np.ndarray):
            return obj.tolist()
        if isinstance(obj, np.bool_):
            return bool(obj)
        return super().default(obj)


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------

def save_model(model, filename: str) -> None:
    """
    Save a LINE model (Network or LayeredNetwork) to a JSON file.

    The output conforms to the line-model.schema.json specification.

    Args:
        model: A Network or LayeredNetwork instance.
        filename: Output file path (should end in .json).

    Example:
        >>> model = Network('M/M/1')
        >>> # ... define model ...
        >>> save_model(model, 'mymodel.json')
    """
    from ..lang.network import Network
    from ..layered import LayeredNetwork

    if isinstance(model, LayeredNetwork):
        model_json = _layered_to_json(model)
    elif isinstance(model, Network):
        model_json = _network_to_json(model)
    else:
        raise TypeError(f"Unsupported model type: {type(model)}")

    doc = {
        "format": "line-model",
        "version": "1.0",
        "model": model_json
    }

    with open(filename, 'w') as f:
        json.dump(doc, f, indent=2, cls=_JSONEncoder)


def load_model(filename: str):
    """
    Load a LINE model from a JSON file.

    Returns a Network or LayeredNetwork depending on the model type
    in the JSON file.

    Args:
        filename: Path to a .json file conforming to line-model.schema.json.

    Returns:
        Network or LayeredNetwork instance.

    Example:
        >>> model = load_model('mymodel.json')
        >>> solver = SolverMVA(model)
        >>> print(solver.avg_table())
    """
    with open(filename, 'r') as f:
        doc = json.load(f)

    model_data = doc.get("model", {})
    mtype = model_data.get("type")

    if mtype == "Network":
        return _json_to_network(model_data)
    elif mtype == "LayeredNetwork":
        return _json_to_layered(model_data)
    else:
        raise ValueError(f"Unsupported model type: {mtype}")
