"""CPU connectome-subgraph controller, batched ramp environment and local web API."""

import os

os.environ.setdefault("OPENBLAS_NUM_THREADS", "1")
os.environ.setdefault("OMP_NUM_THREADS", "1")

import argparse
import atexit
import json
import subprocess
import threading
import time
from http.server import SimpleHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path

import cma
import numpy as np
import pyarrow.parquet as pq

ROOT = Path(__file__).parent
DT = 0.05
READOUT = 16
INPUTS = 18
ACTIONS = 2
PARAMS = (READOUT + 1) * ACTIONS
MODEL = "vercel-rigidbody-v1"
_NODE = None
_NODE_LOCK = threading.Lock()


def simulation(request):
    """Serialize RPC messages, while keeping independent batches for each rollout."""
    global _NODE
    with _NODE_LOCK:
        if _NODE is None:
            _NODE = subprocess.Popen(
                ["node", str(ROOT / "simulation.mjs")],
                stdin=subprocess.PIPE,
                stdout=subprocess.PIPE,
                text=True,
                bufsize=1,
            )
        _NODE.stdin.write(json.dumps(request, separators=(",", ":")) + "\n")
        _NODE.stdin.flush()
        line = _NODE.stdout.readline()
        if not line:
            raise RuntimeError("Reference physics worker exited. Run npm ci first.")
        result = json.loads(line)
        if "error" in result:
            raise RuntimeError(result["error"])
        return result


@atexit.register
def close_simulator():
    if _NODE is not None:
        _NODE.stdin.close()
        _NODE.wait()


class RampEnv:
    """Batched original rigid-body physics; actions select drive and lift/drop.

    The reference's fan, corner-nozzle attitude assist and contact solver are
    preserved. The neural policy chooses commands, not individual nozzle forces.
    """

    def __init__(self, count=1, seed=0, variation=True, conditions=1):
        rng = np.random.default_rng(seed)
        self.n = count
        samples = conditions if conditions > 1 else count
        starts = np.tile([-7.0, 0.0, np.pi / 2, 0.0], (samples, 1))
        if variation:
            starts += rng.uniform(
                [-0.15, -0.03, -0.005, 0], [0.15, 0.03, 0.005, 0.25], (samples, 4)
            )
        starts = np.tile(starts, (count // samples, 1))
        data = simulation({"op": "create", "starts": starts.tolist()})
        self.id = data["id"]
        self.update(data)

    def update(self, data):
        self.observations = np.asarray(data["observations"], dtype=np.float32)
        self.reward = np.asarray(data["reward"])
        self.done = np.asarray(data["done"], dtype=bool)
        self.success = np.asarray(data["success"], dtype=bool)
        self.metrics = data["metrics"]
        self.frame = data["frame"]

    def observe(self):
        return self.observations

    def step(self, action):
        self.update(
            simulation(
                {
                    "op": "step",
                    "id": self.id,
                    "actions": np.clip(action, -1, 1).tolist(),
                }
            )
        )
        return self.observe(), self.reward.copy(), self.done.copy()

    def close(self):
        simulation({"op": "close", "id": self.id})


class Connectome:
    def __init__(self):
        self.nodes = pq.read_table(ROOT / "data/nodes.parquet").to_pylist()
        self.edges = pq.read_table(ROOT / "data/edges.parquet").to_pylist()
        self.matrix = np.zeros((100, 100), dtype=np.float32)
        for edge in self.edges:
            self.matrix[edge["target_index"], edge["source_index"]] = edge[
                "synapse_count"
            ]
        self.matrix /= np.maximum(self.matrix.sum(axis=1, keepdims=True), 1)
        rng = np.random.default_rng(42)
        # Engineered sensory encoder, fixed across all policies and runs.
        self.encoder = rng.normal(0, 0.7, (INPUTS, 100)).astype(np.float32)
        self.output_nodes = np.arange(0, 100, 6)[:READOUT]
        # Sensory drive only enters other nodes; readout must traverse measured edges.
        self.encoder[:, self.output_nodes] = 0
        # Condition the readout on fixed synthetic sensory probes, without rewards
        # or car demonstrations. Otherwise correlated, low-amplitude graph outputs
        # make policy search mostly optimize constant action biases.
        probes = rng.uniform(-1, 1, (2048, INPUTS)).astype(np.float32)
        activity = np.zeros((2048, 100), dtype=np.float32)
        for _ in range(3):
            activity = 0.5 * activity + 0.5 * np.tanh(
                0.8 * (activity @ self.matrix.T) + probes @ self.encoder
            )
        values, vectors = np.linalg.eigh(np.cov(activity[:, self.output_nodes].T))
        self.whitener = (
            vectors @ np.diag(1 / np.sqrt(np.maximum(values, 1e-5))) @ vectors.T
        ).astype(np.float32)

    def forward(self, obs, state, weights, silenced=False):
        drive = obs @ self.encoder
        for _ in range(3):
            state = 0.5 * state + 0.5 * np.tanh(0.8 * (state @ self.matrix.T) + drive)
        if silenced:
            state *= 0
        features = np.column_stack(
            [state[:, self.output_nodes] @ self.whitener, np.ones(len(obs))]
        )
        action = np.tanh(
            np.einsum("bi,bij->bj", features, weights.reshape(-1, READOUT + 1, ACTIONS))
        )
        return action, state


BRAIN = Connectome()


def rollout(
    weights, seed=0, capture=False, variation=True, silenced=False, conditions=1
):
    weights = np.asarray(weights, dtype=np.float32).reshape(-1, PARAMS)
    env = RampEnv(len(weights), seed, variation, conditions)
    state = np.zeros((len(weights), 100), dtype=np.float32)
    frames = []
    try:
        while not env.done.all():
            action, state = BRAIN.forward(env.observe(), state, weights, silenced)
            env.step(action)
            if capture:
                frames.append({**env.frame, "activation": state[0].round(4).tolist()})
    finally:
        env.close()
    result = {
        "model": MODEL,
        "reward": env.reward.tolist(),
        "success": env.success.tolist(),
        "metrics": env.metrics,
    }
    if capture:
        result["frames"] = frames
    return result


def train(generations=400, population=24, seed=7, progress=None, preview=None):
    """Learn drive first, lift/drop next, then jointly optimize safe landings."""
    best = np.zeros(PARAMS)
    champion = best.copy()
    history = []
    start = time.perf_counter()
    drive_end = max(1, generations // 5)
    lift_end = max(drive_end + 1, generations * 4 // 5)
    best_score = -np.inf
    optimizer = None
    previous_stage = None
    for generation in range(generations):
        stage = 0 if generation < drive_end else 1 if generation < lift_end else 2
        joint = stage == 2
        head = min(stage, 1)
        objective = "drive" if stage == 0 else "landing"
        if stage != previous_stage:
            optimizer = cma.CMAEvolutionStrategy(
                best if joint else best.reshape(READOUT + 1, ACTIONS)[:, head],
                0.35 if stage == 0 else 0.08,
                {
                    "popsize": population,
                    "seed": seed + stage,
                    "verbose": -9,
                    "verb_log": 0,
                },
            )
            previous_stage = stage
        samples = optimizer.ask()
        # Evaluate the incumbent on the same starts as its challengers. In
        # particular, the lift stage must not discard a working launch merely
        # because its first random commands fire before the lip.
        samples[0] = (
            best.copy() if joint else best.reshape(READOUT + 1, ACTIONS)[:, head].copy()
        )
        candidates = np.repeat(best[None], population, axis=0)
        if joint:
            candidates[:] = samples
        else:
            candidates.reshape(population, READOUT + 1, ACTIONS)[:, :, head] = samples
        result = rollout(
            np.repeat(candidates, 2, axis=0), seed=seed + generation, conditions=2
        )
        rewards = np.asarray(result["reward"]).reshape(population, 2)
        fitness = (
            rewards
            if joint
            else np.asarray(result["metrics"][objective]).reshape(population, 2)
        )
        scores = fitness.mean(axis=1)
        optimizer.tell(samples, -scores)
        if joint:
            best = np.asarray(optimizer.result.xbest).copy()
        else:
            best.reshape(READOUT + 1, ACTIONS)[:, head] = optimizer.result.xbest
        if generation % 10 == 0 or generation == generations - 1:
            check = rollout(np.repeat(best[None], 8, axis=0), seed=1234)
            score = float(np.mean(check["reward"]))
            if score > best_score:
                best_score, champion = score, best.copy()
        entry = {
            "generation": generation + 1,
            "total": generations,
            "stage": objective,
            "episodes": (generation + 1) * population * 2,
            "mean": round(float(rewards.mean()), 2),
            "best": round(best_score, 2),
            "success": round(float(np.mean(result["success"])), 4),
            "seconds": round(time.perf_counter() - start, 2),
        }
        history.append(entry)
        if progress:
            progress(entry)
        if preview and generation % 10 == 0:
            preview(champion, generation + 1)
    best = champion
    evaluation = rollout(np.repeat(best[None], 100, axis=0), seed=98765)
    baseline = rollout(np.zeros((100, PARAMS)), seed=98765)
    ablation = rollout(np.repeat(best[None], 100, axis=0), seed=98765, silenced=True)
    return {
        "model": MODEL,
        "weights": best.tolist(),
        "history": history,
        "seed": seed,
        "generations": generations,
        "population": population,
        "trials_per_policy": 2,
        "evaluation": {
            "episodes": 100,
            "seed": 98765,
            "trained_success": sum(evaluation["success"]),
            "untrained_success": sum(baseline["success"]),
            "silenced_success": sum(ablation["success"]),
            "trained_mean_reward": float(np.mean(evaluation["reward"])),
        },
        "seconds": round(time.perf_counter() - start, 2),
        "method": "CMA-ES drive/lift readout; frozen measured 100-neuron MaleCNS graph; original 3D physics and attitude assist",
        "physics_source": "https://tesla-ramp.vercel.app/assets/main-LGR2XQC7.js",
    }


LOCK = threading.Lock()
JOB = {"running": False, "history": [], "error": None, "preview": None}


def checkpoint():
    for filename in ["data/latest.json", "data/checkpoint.json"]:
        path = ROOT / filename
        if path.exists():
            saved = json.loads(path.read_text())
            if saved.get("model") == MODEL:
                return saved
    return None


class Handler(SimpleHTTPRequestHandler):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, directory=str(ROOT / "web"), **kwargs)

    def send_json(self, data, status=200):
        body = json.dumps(data, allow_nan=False).encode()
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Cache-Control", "no-store")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def do_GET(self):
        from urllib.parse import parse_qs, urlparse

        url = urlparse(self.path)
        if url.path == "/api/info":
            saved = checkpoint()
            return self.send_json(
                {
                    "nodes": BRAIN.nodes,
                    "edges": BRAIN.edges,
                    "checkpoint": saved,
                    "cpu_count": os.cpu_count(),
                    "model": MODEL,
                }
            )
        if url.path == "/api/training":
            with LOCK:
                return self.send_json(JOB)
        if url.path == "/api/rollout":
            query = parse_qs(url.query)
            policy = query.get("policy", ["trained"])[0]
            if policy not in ("trained", "untrained"):
                return self.send_json({"error": "Unknown policy"}, 400)
            saved = checkpoint()
            if policy == "trained" and not saved:
                return self.send_json({"error": "Train a policy first."}, 409)
            weights = saved["weights"] if policy == "trained" else np.zeros(PARAMS)
            try:
                seed = int(query.get("seed", ["42"])[0])
                if not 0 <= seed <= 2**32 - 1:
                    raise ValueError()
            except ValueError:
                return self.send_json(
                    {"error": "Seed must be an unsigned 32-bit integer."}, 400
                )
            return self.send_json(rollout(weights, seed=seed, capture=True))
        if url.path.startswith("/api/"):
            return self.send_json({"error": "Unknown endpoint"}, 404)
        super().do_GET()

    def do_POST(self):
        if self.path != "/api/train":
            return self.send_json({"error": "Unknown endpoint"}, 404)
        with LOCK:
            if JOB["running"]:
                return self.send_json({"error": "Training already running"}, 409)
            JOB.update(running=True, history=[], error=None, preview=None)

        def work():
            try:

                def report(entry):
                    with LOCK:
                        JOB["history"].append(entry)

                def preview(weights, generation):
                    trial = rollout(weights, seed=42, capture=True)
                    with LOCK:
                        JOB["preview"] = {"generation": generation, "trial": trial}

                result = train(progress=report, preview=preview)
                temp = ROOT / "data/latest.tmp"
                temp.write_text(json.dumps(result))
                temp.replace(ROOT / "data/latest.json")
            except Exception as error:  # noqa: BLE001 - report worker failures to the UI
                with LOCK:
                    JOB["error"] = str(error)
            finally:
                with LOCK:
                    JOB["running"] = False

        threading.Thread(target=work, daemon=True).start()
        self.send_json({"running": True}, 202)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("command", choices=["serve", "train"])
    parser.add_argument("--port", type=int, default=8000)
    parser.add_argument("--generations", type=int, default=400)
    parser.add_argument("--population", type=int, default=24)
    args = parser.parse_args()
    if args.command == "train":
        result = train(
            args.generations,
            args.population,
            progress=lambda row: print(json.dumps(row), flush=True),
        )
        (ROOT / "data/checkpoint.json").write_text(json.dumps(result))
        print(
            json.dumps(
                {k: v for k, v in result.items() if k not in ["weights", "history"]}
            )
        )
    else:
        print(f"Roadster lab listening on port {args.port}", flush=True)
        ThreadingHTTPServer(("0.0.0.0", args.port), Handler).serve_forever()