"""Direct, analytic neural-weight construction; no training or trajectories.

Run with Python + NumPy in the artifact directory. All controller computations
are compiled into the saved MLP. No functions in this file run at deployment.
"""
from pathlib import Path
import json
import numpy as np


class Affine:
    def __init__(self, terms=None, bias=0.0):
        self.terms = dict(terms or {})
        self.bias = float(bias)

    def __add__(self, other):
        other = aff(other)
        out = dict(self.terms)
        for key, value in other.terms.items():
            out[key] = out.get(key, 0.0) + value
        return Affine(out, self.bias + other.bias)

    __radd__ = __add__

    def __neg__(self):
        return self * -1.0

    def __sub__(self, other):
        return self + (-aff(other))

    def __rsub__(self, other):
        return aff(other) - self

    def __mul__(self, scalar):
        scalar = float(scalar)
        return Affine({k: scalar*v for k, v in self.terms.items()},
                      scalar*self.bias)

    __rmul__ = __mul__

    def __truediv__(self, scalar):
        return self * (1.0/float(scalar))


def aff(value):
    if isinstance(value, Affine):
        return value
    if isinstance(value, str):
        return Affine({value: 1.0})
    return Affine(bias=float(value))


class PWL:
    """An affine expression plus a finite sum of weighted ReLU hinges."""
    def __init__(self, base=0.0, hinges=None):
        self.base = aff(base)
        self.hinges = list(hinges or [])

    def __add__(self, other):
        other = pwl(other)
        return PWL(self.base + other.base, self.hinges + other.hinges)

    __radd__ = __add__

    def __neg__(self):
        return self * -1.0

    def __sub__(self, other):
        return self + (-pwl(other))

    def __rsub__(self, other):
        return pwl(other) - self

    def __mul__(self, scalar):
        scalar = float(scalar)
        return PWL(self.base*scalar,
                   [(scalar*k, argument) for k, argument in self.hinges])

    __rmul__ = __mul__


def pwl(value):
    return value if isinstance(value, PWL) else PWL(value)


def relu(argument):
    return PWL(hinges=[(1.0, aff(argument))])


def clip(argument, lower, upper):
    argument = aff(argument)
    return PWL(lower) + relu(argument-lower) - relu(argument-upper)


def absolute(argument):
    argument = aff(argument)
    return relu(argument) + relu(-argument)


def square(argument, radius=1.0, segments=256):
    """Uniform linear interpolation of z^2 on [-radius, radius].

    Positive hinge coefficients give an exactly even function with S(0)=0.
    No fitted coefficients: slope increments are analytically 2*h.
    """
    assert segments % 2 == 0
    argument = aff(argument)
    h = 2.0*radius/segments
    hinges = []
    for j in range(segments//2):
        coefficient = h if j == 0 else 2.0*h
        hinges.append((coefficient, argument-j*h))
        hinges.append((coefficient, -argument-j*h))
    return PWL(hinges=hinges)


def multiply(left, right, left_bound, right_bound):
    """Polarization identity, with analytic square interpolation.

    For |left|<=left_bound, |right|<=right_bound, the absolute error is
    at most left_bound*right_bound / 65536 (256 square segments).
    """
    u = aff(left)/left_bound
    v = aff(right)/right_bound
    return (left_bound*right_bound/4.0) * (
        square(u+v, radius=2.0) - square(u-v, radius=2.0))


def odd_table(argument, positive_knots, positive_values):
    """Odd, continuously piecewise affine interpolation, clamped at ends."""
    argument = aff(argument)
    knots = np.asarray(positive_knots, dtype=np.float64)
    values = np.asarray(positive_values, dtype=np.float64)
    assert knots[0] == 0.0 and values[0] == 0.0
    assert np.all(np.diff(knots) > 0.0)
    slopes = np.diff(values)/np.diff(knots)
    increments = np.concatenate(([slopes[0]], np.diff(slopes), [-slopes[-1]]))
    hinges = []
    for knot, coefficient in zip(knots, increments):
        hinges.append((float(coefficient), argument-float(knot)))
        hinges.append((-float(coefficient), -argument-float(knot)))
    return PWL(hinges=hinges)


class Compiler:
    """Registers are signed pairs of nonnegative ReLU outputs.

    Each stage expands analytic hinges in one layer, and combines them into
    signed registers in a second layer. Only dense numeric matrices are saved.
    """
    def __init__(self):
        self.names = ['o0', 'o1', 'o2', 'o3', 'o4']
        self.weights, self.biases, self.activations = [], [], []
        w = np.zeros((10, 5), dtype=np.float64)
        for j in range(5):
            w[2*j, j], w[2*j+1, j] = 1.0, -1.0
        self.append(w, np.zeros(10), 'relu')
        self.stages = []

    def append(self, w, b, activation):
        self.weights.append(np.asarray(w, dtype=np.float64))
        self.biases.append(np.asarray(b, dtype=np.float64))
        self.activations.append(activation)

    def vector(self, expression, width=None):
        expression = aff(expression)
        row = np.zeros(width if width is not None else 2*len(self.names))
        for name, coefficient in expression.terms.items():
            j = self.names.index(name)
            row[2*j] += coefficient
            row[2*j+1] -= coefficient
        return row, expression.bias

    def stage(self, definitions, keep=(), label=''):
        definitions = [(name, pwl(expression)) for name, expression in definitions]
        assert len(set(keep)) == len(keep)
        assert not (set(keep) & {name for name, _ in definitions})
        old_width = 2*len(self.names)
        hinge_count = sum(len(expression.hinges) for _, expression in definitions)
        hinge_width = old_width + hinge_count
        w1 = np.zeros((hinge_width, old_width))
        b1 = np.zeros(hinge_width)
        w1[:old_width, :] = np.eye(old_width)
        destinations = list(keep) + [name for name, _ in definitions]
        w2 = np.zeros((2*len(destinations), hinge_width))
        b2 = np.zeros(2*len(destinations))
        for j, name in enumerate(keep):
            source = self.names.index(name)
            w2[2*j, 2*source] = 1.0
            w2[2*j+1, 2*source+1] = 1.0
        cursor = old_width
        for j, (name, expression) in enumerate(definitions, start=len(keep)):
            row, bias = self.vector(expression.base, hinge_width)
            for coefficient, argument in expression.hinges:
                w1[cursor, :], b1[cursor] = self.vector(argument)
                row[cursor] += coefficient
                cursor += 1
            w2[2*j, :], w2[2*j+1, :] = row, -row
            b2[2*j], b2[2*j+1] = bias, -bias
        self.append(w1, b1, 'relu')
        self.append(w2, b2, 'relu')
        self.names = destinations
        self.stages.append((label, hinge_width, len(destinations)*2))

    def finish(self, expression):
        expression = pwl(expression)
        old_width = 2*len(self.names)
        width = old_width + len(expression.hinges)
        w1, b1 = np.zeros((width, old_width)), np.zeros(width)
        w1[:old_width, :] = np.eye(old_width)
        row, bias = self.vector(expression.base, width)
        for cursor, (coefficient, argument) in enumerate(expression.hinges,
                                                        start=old_width):
            w1[cursor, :], b1[cursor] = self.vector(argument)
            row[cursor] += coefficient
        self.append(w1, b1, 'relu')
        self.append(row.reshape(1, -1), np.array([bias]), 'tanh')

    def save(self, path):
        values = {'activations': np.asarray(self.activations, dtype='U6')}
        for j, (w, b) in enumerate(zip(self.weights, self.biases)):
            values['w'+str(j)], values['b'+str(j)] = w, b
        np.savez(path, **values)

    def forward(self, observations, dtype=np.float64):
        values = np.asarray(observations, dtype=dtype)
        for w, b, activation in zip(self.weights, self.biases, self.activations):
            values = values @ w.astype(dtype, copy=False).T + b.astype(dtype, copy=False)
            if activation == 'relu':
                values = np.maximum(values, 0)
            elif activation == 'tanh':
                values = np.tanh(values)
        return values


def main():
    # Only the supplied public task files are read. No simulator is imported.
    public = json.loads(Path('public_models/public.json').read_text())
    xml = Path('public_models/cartpole.xml').read_text()
    assert public['force_limit_N'] == 10.0
    assert public['pole_pivot_inertia_kg_m2'] == 1.0/30.0
    assert 'Direct-weight cart-pole' in xml

    # Algebraic pole placement for input a = cart acceleration.
    # delta is angle relative to upright; delta_dd = 1.5*a+14.715*delta-.3*w.
    desired_poles = np.array([-1.2, -1.6, -3.8, -4.2])
    polynomial = np.poly(desired_poles)
    gravity_rate = 14.715
    kx = polynomial[4]/gravity_rate
    kv = (polynomial[3] + 0.3*kx)/gravity_rate
    kw = (polynomial[1] - 0.3 + kv)/1.5
    kt = (polynomial[2] + gravity_rate + 0.3*kv + kx)/1.5
    energy_target, energy_gain = 1.12, 8.0
    cart_kp, cart_kd = 1.5, 2.2
    rail_reference, rail_rate = 2.0, 1.8
    acceleration_limit = 9.0

    net = Compiler()
    net.stage([
        ('x', clip(2.4*aff('o0'), -3.0, 3.0)),
        ('v', clip(3.0*aff('o1'), -6.0, 6.0)),
        ('s', clip(aff('o2'), -1.0, 1.0)),
        ('c', clip(aff('o3'), -1.0, 1.0)),
        ('w', clip(8.0*aff('o4'), -16.0, 16.0)),
    ], label='input scaling and defensive clipping')
    x, v, s, c, w = map(aff, ('x', 'v', 's', 'c', 'w'))
    asin_knots = np.linspace(0.0, 0.8, 65)
    net.stage([
        ('w2', 256.0*square(w/16.0, segments=512)),
        ('c2', square(c)),
        ('wc', multiply(w, c, 16.0, 1.0)),
        ('cs', multiply(c, s, 1.0, 1.0)),
        ('delta', -odd_table(s, asin_knots, np.arcsin(asin_knots))),
        ('absw', absolute(w)),
    ], keep=['x', 'v', 's', 'c', 'w'], label='kinematic scalar primitives')
    net.stage([
        ('E', aff('w2')/60.0 + 0.4905*(1.0-c)),
        ('sw2', multiply(s, aff('w2'), 1.0, 256.0)),
        ('gc', clip((-c-0.75)/0.15, 0.0, 1.0)),
        ('gw', clip((5.0-aff('absw'))/1.5, 0.0, 1.0)),
        ('bal', clip(kx*x+kv*v-kt*aff('delta')-kw*w,
                     -acceleration_limit, acceleration_limit)),
    ], keep=['x', 'v', 'c2', 'wc', 'cs'], label='energy, capture region, local feedback')
    net.stage([
        ('ep', multiply(aff('E')-energy_target, aff('wc'), 4.2, 16.0)),
        ('kick', 0.6*clip(1.0-aff('E')/0.1, 0.0, 1.0)),
        ('gate', PWL(aff('gc'))-relu(aff('gc')-aff('gw'))),
    ], keep=['x', 'v', 'c2', 'wc', 'cs', 'sw2', 'bal'],
       label='energy pumping product and capture gate')
    net.stage([
        ('pump', clip(energy_gain*aff('ep'), -6.5, 6.5)),
    ], keep=['x', 'v', 'c2', 'wc', 'cs', 'sw2', 'bal', 'kick', 'gate'],
       label='pump saturation')
    net.stage([
        ('sw', clip(aff('pump')-cart_kp*x-cart_kd*v+aff('kick'),
                    -acceleration_limit, acceleration_limit)),
    ], keep=['x', 'v', 'c2', 'wc', 'cs', 'sw2', 'bal', 'gate'],
       label='swing-up cart regulation')
    net.stage([
        ('blend', PWL(aff('sw'))+multiply(aff('gate'), aff('bal')-aff('sw'),
                                        1.0, 18.0)),
    ], keep=['x', 'v', 'c2', 'wc', 'cs', 'sw2'], label='continuous capture blend')
    lower = -(rail_rate**2)*(rail_reference+x)-2.0*rail_rate*v
    upper = (rail_rate**2)*(rail_reference-x)-2.0*rail_rate*v
    net.stage([
        ('lower_bounded', PWL(aff('blend'))+relu(lower-aff('blend'))),
    ], keep=['x', 'v', 'c2', 'wc', 'cs', 'sw2'], label='lower rail acceleration bound')
    net.stage([
        ('rail', PWL(aff('lower_bounded'))-relu(aff('lower_bounded')-upper)),
    ], keep=['v', 'c2', 'wc', 'cs', 'sw2'], label='upper rail acceleration bound')
    net.stage([
        ('a', clip(aff('rail'), -acceleration_limit, acceleration_limit)),
    ], keep=['v', 'c2', 'wc', 'cs', 'sw2'], label='acceleration feasibility clamp')
    net.stage([
        ('c2a', multiply(aff('c2'), aff('a'), 1.0, acceleration_limit)),
    ], keep=['v', 'wc', 'cs', 'sw2', 'a'], label='effective mass product')
    net.stage([
        ('F', 1.1*aff('a')-0.075*aff('c2a')+0.05*aff('v')
              -0.73575*aff('cs')-0.015*aff('wc')-0.05*aff('sw2')),
    ], label='compiled acceleration-to-force algebra')

    # A bounded, nearly hard-linear force limiter. Tabulate atanh analytically
    # at z=tanh(t), then apply the final tanh. There is no output clipping code.
    t = np.linspace(0.0, 3.0, 257)
    net.finish(odd_table(aff('F')/10.0, np.tanh(t), t))
    net.save('cartpole.npz')

    # Static checks only. None updates a physical state or generates a trajectory.
    shapes = [list(w.shape) for w in net.weights]
    assert 1 <= len(shapes) <= 64 and max(a for a, b in shapes) <= 4096
    assert shapes[0][1] == 5 and shapes[-1][0] == 1
    for i, (weight, bias) in enumerate(zip(net.weights, net.biases)):
        assert weight.ndim == 2 and bias.shape == (weight.shape[0],)
        assert np.isfinite(weight).all() and np.isfinite(bias).all()
        if i:
            assert weight.shape[1] == net.weights[i-1].shape[0]
    array_bytes = sum(w.nbytes+b.nbytes for w, b in zip(net.weights, net.biases))
    array_bytes += np.asarray(net.activations, dtype='U6').nbytes
    assert array_bytes < 128*1024*1024
    with np.load('cartpole.npz', allow_pickle=False) as saved:
        expected = {'activations'} | {f'{prefix}{i}' for i in range(len(shapes))
                                      for prefix in ('w', 'b')}
        assert set(saved.files) == expected
        assert saved['activations'].dtype.kind == 'U'
        assert all(saved[k].dtype.kind == 'f' for k in expected-{'activations'})

    # Independent points, not sequential state transitions. Randomness is used
    # solely to check finite, bounded network arithmetic, never to select weights.
    rng = np.random.default_rng(20261105)
    n = 96
    angles = rng.uniform(-np.pi, np.pi, n)
    checks = np.column_stack((rng.uniform(-1.0, 1.0, n),
                              rng.uniform(-2.0, 2.0, n),
                              np.sin(angles), np.cos(angles),
                              rng.uniform(-2.0, 2.0, n)))
    anchors = np.array([[0, 0, 0, 1, 0], [0, 0, 0, -1, 0],
                        [1e3]*5, [-1e3]*5], dtype=np.float64)
    checks = np.vstack([anchors, checks])
    out64 = net.forward(checks)
    out32 = net.forward(checks, dtype=np.float32)
    assert np.isfinite(out64).all() and np.isfinite(out32).all()
    assert np.max(np.abs(out64)) <= 1.0 and np.max(np.abs(out32)) <= 1.0
    assert abs(out64[1, 0]) < 1e-10

    # Static physical linearization and a numerical network Jacobian at upright.
    def upright_observation(state):
        xx, vv, delta, omega = state
        return np.array([xx/2.4, vv/3.0, -np.sin(delta), -np.cos(delta), omega/8.0])

    h = 1e-4
    jacobian = np.zeros(4)
    for j in range(4):
        unit = np.eye(4)[j]*h
        jacobian[j] = float(10.0*(net.forward(upright_observation(unit))[0]
                                  -net.forward(upright_observation(-unit))[0])/(2*h))
    inverse_mass = np.linalg.inv(np.array([[1.1, -0.05], [-0.05, 1.0/30.0]]))
    A = np.zeros((4, 4))
    A[0, 1], A[2, 3] = 1.0, 1.0
    A[[1, 3], 1] = inverse_mass @ np.array([-0.05, 0.0])
    A[[1, 3], 2] = inverse_mass @ np.array([0.0, 0.4905])
    A[[1, 3], 3] = inverse_mass @ np.array([0.0, -0.01])
    B = np.zeros(4)
    B[[1, 3]] = inverse_mass @ np.array([1.0, 0.0])
    closed_loop_eigenvalues = np.linalg.eigvals(A+np.outer(B, jacobian))
    assert np.all(np.real(closed_loop_eigenvalues) < 0.0)
    # These continuous-time local eigenvalues do not establish sampled,
    # saturated, nonlinear or hybrid-task success.

    report = {
        'affine_layers': len(shapes),
        'hidden_widths': [w.shape[0] for w in net.weights[:-1]],
        'max_width': max(w.shape[0] for w in net.weights),
        'numeric_parameters': sum(w.size+b.size for w, b in zip(net.weights, net.biases)),
        'uncompressed_array_bytes': array_bytes,
        'npz_file_bytes': Path('cartpole.npz').stat().st_size,
        'balance_acceleration_gains': {'x': kx, 'v': kv, 'delta': -kt, 'omega': -kw},
        'upright_force_jacobian': jacobian.tolist(),
        'local_continuous_eigenvalues': sorted(closed_loop_eigenvalues.real.tolist()),
        'synthetic_output_abs_max_float64': float(np.max(np.abs(out64))),
        'synthetic_output_abs_max_float32': float(np.max(np.abs(out32))),
        'synthetic_max_float32_difference': float(np.max(np.abs(out64-out32))),
        'downward_rest_output': float(out64[0, 0]),
        'upright_rest_output': float(out64[1, 0]),
        'trajectories_or_behavioral_evaluations': 0,
    }
    print(json.dumps(report, indent=2))

    rationale = f"""CART-POLE: FIRST DIRECT-WEIGHT SUBMISSION

Artifacts
  cartpole.npz: fixed numeric MLP; construction.py: reproducible analytic authoring.
  Only NumPy and standard Python are needed to construct the artifact.
  The public XML and JSON are the only external task information read.

Actual architecture
  {len(shapes)} affine layers; all hidden activations ReLU; final activation tanh.
  Input width 5, output width 1, maximum hidden width {report['max_width']}.
  Hidden widths: {report['hidden_widths']}
  {report['numeric_parameters']:,} numeric weight/bias entries, float64.
  Uncompressed array payload: {array_bytes:,} bytes.
  NPZ contains exactly activations (Unicode), w0,b0,... in runtime order.
  Sparse arithmetic connections are stored as ordinary dense numeric arrays.
  Signed-pair registers, hinge sums, clipping, products, and blending are all
  internal affine/ReLU computations, not external deployment code.

Physical derivation and designed law (SI units)
  s=sin(q), c=cos(q), w=qdot, v=xdot.
  E=w^2/60 + 0.4905*(1-c); upright energy is 0.981 J.
  From the supplied equations, E_dot=-0.05*w*c*a-0.01*w^2, where a=xdd.
  Therefore a proportional to (E-E_target)*w*c pumps energy below target.
  E_target=1.12 J deliberately allows margin for hinge damping, centering,
  and approximate arithmetic; it is not a measured or optimized target.
  pump=clip(8*(E-1.12)*w*c, -6.5, 6.5).
  kick=0.6*clip(1-E/0.1, 0, 1), a small low-energy symmetry-breaking bias.
  a_sw=clip(pump-1.5*x-2.2*v+kick, -9, 9).

  Inside the inverted hemisphere, delta=-asin(s) is upright angular error.
  Its table is clamped outside |s|<=0.8; the balance gate is zero before
  that angular ambiguity matters for physically consistent observations.
  a_bal=clip({kx:.12g}*x+{kv:.12g}*v-{kt:.12g}*delta-{kw:.12g}*w, -9, 9).
  Gains come from the linear acceleration-input model
      xdd=a, delta_dd=1.5*a+14.715*delta-0.3*w,
  with chosen poles -1.2,-1.6,-3.8,-4.2 per second, by polynomial matching.
  No gain fitting or behavior search was used.

  g=min(clip((-c-0.75)/0.15,0,1), clip((5-|w|)/1.5,0,1)).
  a_mix=a_sw+g*(a_bal-a_sw); no hysteresis, clock, or stored mode.
  For rail reference L=2 m and rate lambda=1.8/s, constrain a_mix between
      a_low=-lambda^2*(L+x)-2*lambda*v,
      a_high=lambda^2*(L-x)-2*lambda*v,
  then clip acceleration to [-9,9]. These are second-order barrier-style
  limits. With exact feasible acceleration they enforce
  h_dd+2*lambda*h_d+lambda^2*h>=0 for h=L +/- x. Force saturation,
  sampling, approximations, and model coupling preclude a safety guarantee.

  Elimination of qdd in the given equations gives the compiled force map
      F=(1.1-0.075*c^2)*a+0.05*v-0.73575*c*s-0.015*c*w-0.05*s*w^2.
  No inverse-dynamics or compensation routine runs outside the MLP.

Arithmetic compilation
  Squares are uniform linear interpolants with analytically known hinge
  coefficients, not regressions. w^2 uses 512 intervals after scaling w/16;
  c^2 uses 256. Products use polarization of 256-interval square splines
  on [-2,2]. Product error <= A*B/65536 for factors bounded by A,B.
  The asin table has 64 positive intervals, analytically sampled.
  Input scalings recover physical units; defensive clips are |x|<=3 m,
  |v|<=6 m/s, |w|<=16 rad/s, and |s|,|c|<=1. No statistical normalization.
  These ranges are construction assumptions, not measured state bounds.
  The output layer applies tanh to an odd analytic atanh interpolant:
  knots z=tanh(t), t=0,3/256,...,3, with odd extension and end clamping.
  It closely implements F/10 up to saturation near +/-tanh(3), rather
  than strongly compressing moderate forces. Final tanh guarantees an
  action in [-1,1], including if preceding interpolation has rounding error.

Checks actually performed
  Public parameter assertions, NPZ key/dtype/shape/size/finite checks.
  Independent synthetic inputs only: finite bounded outputs in float64
  and float32; no state transitions, task rollouts, or reward queries.
  Float64 max |output|: {report['synthetic_output_abs_max_float64']:.12g}.
  Float32 max |output|: {report['synthetic_output_abs_max_float32']:.12g}.
  Max float32/64 difference on those points: {report['synthetic_max_float32_difference']:.6g}.
  Static upright force Jacobian [x,v,delta,w]: {jacobian.tolist()}.
  Corresponding continuous-time linearized eigenvalues:
      {report['local_continuous_eigenvalues']}
  These checks are not behavioral feedback and do not show task success.

Limitations
  ZERO trajectories, training steps, reward searches, teacher fits, downloads,
  or pretrained policy weights. No simulator was imported or equations
  integrated. Coefficients and interpolation grids were chosen directly.
  Swing-up/capture/centering within 15 s, rail survival, sampled-data
  stability, and the continuous one-second success hold are UNTESTED.
  Energy control and local stabilization do not establish a global region
  of attraction. Gate switching, force limits, high angular speeds, and
  conflicting rail/balance demands can prevent capture or cause failure.
  Clipped velocities outside the declared ranges lose state information.
  The artifact is a genuine first submission, not a claim of success.
"""
    Path('rationale.txt').write_text(rationale)


if __name__ == '__main__':
    main()