"""
Linear Admission Constraints for Finite Capacity Regions

Validates LINE's DES engine (Solver_ssj) enforcement of general linear
admission constraints An <= b against the Erlang fixed-point analytical
solution.

Uses JPype to invoke LINE's Java DES engine (which has linear constraint
support) directly.

Test cases:
  Case 1 — Per-class only (A=I): backward compatibility
  Case 2 — Global + per-class: identity rows plus an all-ones row
  Case 3 — General weighted: cross-class coupling
"""

import sys
import os
import subprocess
import numpy as np
import jpype
import jpype.imports

# Get line-dev.git root
LINE_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..', '..'))

# Add native Python line_solver for the Erlang FP function
sys.path.insert(0, os.path.join(LINE_ROOT, 'python'))
from line_solver.api.lossn.erlang import lossn_erlangfp  # noqa: E402


def _start_jvm():
    """Start JVM with jline.jar and all Maven dependencies."""
    if jpype.isJVMStarted():
        return

    jline_jar = os.path.join(LINE_ROOT, 'common', 'jline.jar')
    jmt_jar = os.path.join(LINE_ROOT, 'common', 'JMT.jar')

    # Get Maven classpath for transitive dependencies
    result = subprocess.run(
        ['mvn', '-f', os.path.join(LINE_ROOT, 'jar', 'pom.xml'),
         'dependency:build-classpath', '-DincludeScope=compile', '-q',
         '-Dmdep.outputFile=/dev/stdout'],
        capture_output=True, text=True, timeout=60
    )
    maven_cp = result.stdout.strip().split('\n')[-1] if result.returncode == 0 else ''

    classpath = [jline_jar]
    if os.path.isfile(jmt_jar):
        classpath.append(jmt_jar)
    if maven_cp:
        classpath.extend(maven_cp.split(':'))

    jpype.startJVM(classpath=classpath)


def build_and_solve(case_name, lambdas, mus, A, b, seed=23000, samples=500000):
    """Build an open loss network, solve with DES and Erlang FP, compare."""
    _start_jvm()

    K = len(lambdas)

    # Import Java classes
    from jline.lang import Network, OpenClass
    from jline.lang.nodes import Source, Delay, Sink
    from jline.lang.processes import Exp
    from jline.lang.constant import DropStrategy
    from jline.solvers.des import SolverDES
    from jline.util.matrix import Matrix

    # Build Java LINE network
    model = Network(f'FCR LinCon - {case_name}')

    source = Source(model, 'Source')
    delay = Delay(model, 'Delay')
    sink = Sink(model, 'Sink')

    classes = []
    for r in range(K):
        c = OpenClass(model, f'Class{r+1}')
        classes.append(c)

    for r in range(K):
        source.setArrival(classes[r], Exp.fitRate(lambdas[r]))
        delay.setService(classes[r], Exp.fitRate(mus[r]))

    P = model.initRoutingMatrix()
    for r in range(K):
        P.set(classes[r], classes[r], source, delay, 1.0)
        P.set(classes[r], classes[r], delay, sink, 1.0)
    model.link(P)

    # Add FCR with DROP policy and linear constraints
    import java.util
    node_list = java.util.ArrayList()
    node_list.add(delay)
    fcr = model.addRegion(node_list)
    for r in range(K):
        fcr.setClassMaxJobs(classes[r], 1000)
        fcr.setDropRule(classes[r], DropStrategy.Drop)
    fcr.setGlobalMaxJobs(1000)

    # Set linear constraints
    A_np = np.asarray(A, dtype=np.float64)
    b_np = np.asarray(b, dtype=np.float64).reshape(-1, 1)
    C_con, K_con = A_np.shape
    A_mat = Matrix(C_con, K_con)
    for i in range(C_con):
        for j in range(K_con):
            A_mat.set(i, j, float(A_np[i, j]))
    b_mat = Matrix(C_con, 1)
    for i in range(C_con):
        b_mat.set(i, 0, float(b_np[i, 0]))
    fcr.setLinearConstraints(A_mat, b_mat)

    # DES solve using LINE's Java DES engine
    solver = SolverDES(model, "seed", int(seed), "samples", int(samples))

    # getAvg() runs the solver and populates result
    solver_result = solver.getAvg()

    # Extract throughputs from the result object (TN is Matrix(M, K))
    tput_des = np.zeros(K)
    T = solver_result.TN

    if T is not None:
        # In Source-Delay-Sink: station 0 = Source, station 1 = Delay
        delay_idx = int(T.getNumRows()) - 1
        for r in range(K):
            tput_des[r] = T.get(delay_idx, r)

    # Analytical solve: nu = lambda/mu (offered load per class)
    nu = np.array(lambdas) / np.array(mus)
    qlen_fp, loss_fp, eblock_fp, niter = lossn_erlangfp(nu, A, b)

    # DES loss = 1 - Tput_des / lambda
    loss_des = 1.0 - tput_des / np.array(lambdas)

    # Analytical throughput = lambda * (1 - loss)
    tput_fp = np.array(lambdas) * (1.0 - loss_fp)

    return tput_des, tput_fp, loss_des, loss_fp


def run_case(case_name, lambdas, mus, A, b, tol=0.05, **kwargs):
    """Run a test case and report results."""
    print(f'\n{"="*60}')
    print(f'  {case_name}')
    print(f'{"="*60}')
    print(f'  A = {A.tolist()}')
    print(f'  b = {b.tolist()}')
    print(f'  lambda = {lambdas}')
    print(f'  mu     = {mus}')

    tput_des, tput_fp, loss_des, loss_fp = build_and_solve(
        case_name, lambdas, mus, A, b, **kwargs
    )

    print(f'\n  Throughput:')
    print(f'    DES:    {tput_des}')
    print(f'    Erlang: {tput_fp}')

    print(f'  Loss probability:')
    print(f'    DES:    {loss_des}')
    print(f'    Erlang: {loss_fp}')

    # Relative error on throughput
    rel_err = np.abs(tput_des - tput_fp) / np.maximum(tput_fp, 1e-10)
    print(f'  Relative error (throughput): {rel_err}')

    passed = np.all(rel_err < tol)
    status = 'PASS' if passed else 'FAIL'
    print(f'  Result: {status}')
    return passed


if __name__ == '__main__':
    all_pass = True

    # Case 1: Per-class only (A = I) — backward compatibility
    # Each class limited independently: n_1 <= 3, n_2 <= 3
    A1 = np.eye(2)
    b1 = np.array([3.0, 3.0])
    all_pass &= run_case(
        'Case 1: Per-class only (A=I)',
        lambdas=[0.3, 0.2], mus=[1.0, 0.8],
        A=A1, b=b1, samples=500000,
    )

    # Case 2: Global + per-class
    # n_1 <= 3, n_2 <= 3, n_1 + n_2 <= 5
    A2 = np.array([[1.0, 0.0],
                    [0.0, 1.0],
                    [1.0, 1.0]])
    b2 = np.array([3.0, 3.0, 5.0])
    all_pass &= run_case(
        'Case 2: Global + per-class',
        lambdas=[0.3, 0.2], mus=[1.0, 0.8],
        A=A2, b=b2, samples=500000,
    )

    # Case 3: General weighted constraints (cross-class coupling)
    # 2*n_1 + 1*n_2 <= 5, 1*n_1 + 3*n_2 <= 7
    A3 = np.array([[2.0, 1.0],
                    [1.0, 3.0]])
    b3 = np.array([5.0, 7.0])
    all_pass &= run_case(
        'Case 3: General weighted',
        lambdas=[0.3, 0.2], mus=[1.0, 0.8],
        A=A3, b=b3, samples=500000,
    )

    print(f'\n{"="*60}')
    if all_pass:
        print('  ALL CASES PASSED')
    else:
        print('  SOME CASES FAILED')
    print(f'{"="*60}')
