"""Private complete-worker screen; experimental CUDA source is bound separately."""
from __future__ import annotations

import argparse
from datetime import datetime, timezone
import hashlib
import json
import os
from pathlib import Path
import statistics
import subprocess
import sys
import time

ROOT = Path(__file__).resolve().parent
KERNEL = r"""
extern "C" __global__ void evaluate(
    const double* data, const int* frame, const double* cameras, int n, int covariance,
    double lx, double ly, double lz, double lw, double rx, double ry, double rz, double rw,
    double tx, double ty, double tz, double* out
) {
    int i = blockDim.x * blockIdx.x + threadIdx.x;
    if (i >= n) return;
    out[2*n+i] = 0.0;
    const double* p = data + 10*i;
    const double* c = cameras + 14*frame[i];
    double x=p[0], y=p[1], z=p[2];
    double ax=lw*x+lx*0.0+ly*z-lz*y;
    double ay=lw*y-lx*z+ly*0.0+lz*x;
    double az=lw*z+lx*y-ly*x+lz*0.0;
    double aw=lw*0.0-lx*x-ly*y-lz*z;
    double xx=aw*rx+ax*rw+ay*rz-az*ry+tx;
    double yy=aw*ry-ax*rz+ay*rw+az*rx+ty;
    double zz=aw*rz+ax*ry-ay*rx+az*rw+tz;
    if (!(zz > 1e-9)) return;
    double denominator=zz;
    if (c[0] != 0.0) {
        double norm=sqrt(xx*xx+yy*yy+zz*zz);
        if (c[0] == 1.0) denominator=zz+c[7]*norm;
        else {
            double zeta=c[7]*norm+zz;
            double second=sqrt(xx*xx+yy*yy+zeta*zeta);
            denominator=c[8]*second+(1.0-c[8])*zeta;
        }
        if (!(denominator > 1e-12)) return;
    }
    xx=xx/denominator; yy=yy/denominator;
    if (c[13] != 0.0) {
        double radius=xx*xx+yy*yy;
        if (isfinite(radius) && fabs(radius)>1e100) { out[2*n+i]=2.0; return; }
        double squared=radius*radius;
        double cubed=squared*radius;
        double radial=1.0+c[9]*radius+c[10]*squared+0.0*cubed;
        double dx=xx*radial+2.0*c[11]*xx*yy+c[12]*(radius+2.0*xx*xx);
        double dy=yy*radial+c[11]*(radius+2.0*yy*yy)+2.0*c[12]*xx*yy;
        xx=dx; yy=dy;
    }
    double u=c[1]*xx+c[3], v=c[2]*yy+c[4];
    if (!(u>=0.0 && u<c[5] && v>=0.0 && v<c[6])) return;
    double du=u-p[3], dv=v-p[4];
    double pixel=du*du+dv*dv;
    double squared_error=pixel;
    if (covariance) {
        double a=p[5], b=p[6], cc=p[7], d=p[8];
        squared_error=(d*du*du-(b+cc)*du*dv+a*dv*dv)/(a*d-b*cc);
    }
    out[i]=pixel;
    out[n+i]=sqrt(squared_error>0.0 ? squared_error : 0.0);
    out[2*n+i]=1.0;
}
"""


def sha(path: Path) -> str:
    return hashlib.sha256(path.read_bytes()).hexdigest()


def install_cuda():
    import math
    import numpy as np
    from calibrex.core.geometry import normalize_quaternion_xyzw, quaternion_conjugate_xyzw
    from calibrex.solvers import probabilistic_camera_lidar_refiner as module
    base = module._PreparedPoseEvaluation
    counters = {"calls": 0, "cuda_calls": 0, "fallback_calls": 0, "download_bytes": 0}

    class Prepared(base):
        def __init__(self, frames, settings):
            super().__init__(frames, settings)
            self.device = None
            if self.fallback or not self.frames:
                return
            for f in self.frames:
                intr = f.intrinsics
                if intr.distortion_model == "radial-tangential" and (
                    len(intr.distortion) not in {4, 5}
                    or (len(intr.distortion) == 5 and intr.distortion[4] != 0.0)
                ):
                    return
            import cupy as cp
            rows, frame_ids, cameras = [], [], []
            for index, f in enumerate(self.frames):
                intr = f.intrinsics
                rows.append(np.column_stack((f.points, f.means, f.covariance, f.weights)))
                frame_ids.append(np.full(len(f.points), index, dtype=np.int32))
                distortion = list(intr.distortion[:4]) + [0.0]*4
                cameras.append([
                    {"pinhole": 0.0, "mei": 1.0, "double_sphere": 2.0}[intr.projection],
                    intr.fx, intr.fy, intr.cx, intr.cy, intr.width, intr.height,
                    intr.xi or 0.0, intr.alpha or 0.0, *distortion[:4],
                    float(intr.distortion_model == "radial-tangential"),
                ])
            packed = np.concatenate(rows)
            self.host_weights = packed[:, 9].copy()
            self.n = len(packed)
            self.device = (
                cp.asarray(packed), cp.asarray(np.concatenate(frame_ids)),
                cp.asarray(np.asarray(cameras, dtype=np.float64)),
                cp.empty((3, self.n), dtype=cp.float64),
            )
            self.kernel = cp.RawKernel(
                KERNEL, "evaluate", options=("--std=c++11", "--fmad=false")
            )

        def evaluate(self, transform):
            counters["calls"] += 1
            if self.device is None:
                counters["fallback_calls"] += 1
                return super().evaluate(transform)
            import cupy as cp
            left = normalize_quaternion_xyzw(transform.rotation_quat_xyzw)
            right = quaternion_conjugate_xyzw(left)
            data, ids, cameras, out = self.device
            self.kernel(
                ((self.n+255)//256,), (256,),
                (data, ids, cameras, np.int32(self.n), np.int32(self.settings.use_covariance),
                 *(np.float64(v) for v in (*left, *right, *transform.translation_m)), out),
            )
            values = out.get()
            counters["cuda_calls"] += 1
            counters["download_bytes"] += values.nbytes
            if np.any(values[2] == 2.0):
                counters["fallback_calls"] += 1
                return super().evaluate(transform)
            valid = values[2] == 1.0
            count = int(np.count_nonzero(valid))
            if not count:
                return module.ProbabilisticPoseEvaluation(math.inf, None, None, 0, self.total)
            pixels, errors = values[0, valid], values[1, valid]
            weights = self.host_weights[valid]
            robust = np.log1p((errors/self.settings.cauchy_scale)**2)
            return module.ProbabilisticPoseEvaluation(
                float(np.sum(weights*robust)/np.sum(weights)),
                math.sqrt(float(np.sum(weights*pixels))/float(np.sum(weights))),
                float(np.sum(weights*errors)/np.sum(weights)), count, self.total-count,
            )

    module._PreparedPoseEvaluation = Prepared
    return counters


def worker(args):
    source = Path(args.source).resolve()
    sys.path.insert(0, str(source/"src"))
    from calibrex.evaluation.probabilistic_camera_lidar_run import run_probabilistic_camera_lidar_refinement
    from calibrex.solvers.probabilistic_camera_lidar_refiner import ProbabilisticCameraLidarRefinementOptions
    counters = install_cuda() if args.variant == "cuda" else {}
    output = Path(args.output)
    output.mkdir()
    begin = time.perf_counter()
    result = run_probabilistic_camera_lidar_refinement(
        args.correspondence, args.problem,
        options=ProbabilisticCameraLidarRefinementOptions(max_evaluations=400),
        command=["probabilistic-cuda-private-screen"],
    )
    result.save(output/"result.json")
    api = time.perf_counter()-begin
    raw = result.model_dump(mode="json")
    raw["provenance"].pop("created_at")
    canonical = json.dumps(raw, sort_keys=True, separators=(",", ":"), allow_nan=False)
    manifest = json.loads((source/"manifest.json").read_bytes())
    assert manifest == {p.relative_to(source/"src").as_posix():sha(p) for p in (source/"src").rglob("*.py")}
    for name, mod in list(sys.modules.items()):
        if (name == "calibrex" or name.startswith("calibrex.")) and getattr(mod, "__file__", None):
            path = Path(mod.__file__).resolve()
            assert path.is_relative_to(source/"src") and sha(path)==manifest[path.relative_to(source/"src").as_posix()]
    payload = {
        "variant": args.variant, "api_seconds": api, "counts": counters,
        "trace_evaluations": len(result.trace),
        "canonical_result_sha256": hashlib.sha256(canonical.encode()).hexdigest(),
        "provenance": {
            "recorded_at":datetime.now(timezone.utc).isoformat(),
            "source_commit":args.commit, "source_manifest_sha256":sha(source/"manifest.json"),
            "launcher_sha256":sha(Path(__file__)), "kernel_sha256":hashlib.sha256(KERNEL.encode()).hexdigest(),
            "correspondence_sha256":sha(Path(args.correspondence)), "problem_sha256":sha(Path(args.problem)),
        },
    }
    (output/"measurement.json").write_text(json.dumps(payload,indent=2)+"\n",encoding="utf-8")
    print(json.dumps(payload),flush=True)


def main():
    parser=argparse.ArgumentParser()
    parser.add_argument("--variant",choices=["numpy","cuda"])
    parser.add_argument("--source")
    parser.add_argument("--output")
    parser.add_argument("--commit")
    parser.add_argument("--correspondence")
    parser.add_argument("--problem")
    args=parser.parse_args()
    if args.variant:
        worker(args)
        return
    sys.path.insert(0,str(ROOT))
    from tools.benchmark_bspline_lidar_imu import _freeze
    output=ROOT/"outputs/probabilistic-cuda-screen-v1"
    output.mkdir()
    commit=subprocess.check_output(["git","rev-parse","HEAD"],text=True).strip()
    _freeze(output/"source",commit)
    fixture=ROOT/"outputs/probabilistic-cuda-screen-inputs-v1"
    protocol={
        "scope":"Exploratory complete artifact-backed refinement, public KITTI geometry and generated image means.",
        "order":["numpy","cuda","cuda","numpy"],"source_commit":commit,
        "launcher_sha256":sha(Path(__file__)),"kernel_sha256":hashlib.sha256(KERNEL.encode()).hexdigest(),
        "input_sha256":{name:sha(fixture/name) for name in ["correspondence.json","problem.json","fixture-protocol.json"]},
        "gates":{"whole_ratio_at_least":2.0,"canonical_result":"exact, creation time only normalized"},
        "provenance":{"recorded_at":datetime.now(timezone.utc).isoformat(),"activity":"unreserved host","blas_threads":1},
    }
    (output/"protocol.json").write_text(json.dumps(protocol,indent=2)+"\n",encoding="utf-8")
    measurements=[]
    for index,variant in enumerate(protocol["order"]):
        target=output/f"{index:02d}-{variant}"
        env=os.environ.copy()
        env["CUPY_CACHE_DIR"]=str(output/f"cache-{index:02d}")
        env["CUDA_CACHE_DISABLE"]="1"
        assert not Path(env["CUPY_CACHE_DIR"]).exists()
        command=[sys.executable,str(Path(__file__)),"--variant",variant,"--source",str(output/"source"),
                 "--output",str(target),"--commit",commit,"--correspondence",str(fixture/"correspondence.json"),
                 "--problem",str(fixture/"problem.json")]
        begin=time.perf_counter()
        with (output/f"{index:02d}-{variant}.stdout.log").open("w",encoding="utf-8") as stdout:
            with (output/f"{index:02d}-{variant}.stderr.log").open("w",encoding="utf-8") as stderr:
                completed=subprocess.run(command,env=env,stdout=stdout,stderr=stderr)
        whole=time.perf_counter()-begin
        if completed.returncode:
            print(json.dumps({"index":index,"variant":variant,"returncode":completed.returncode}),flush=True)
            raise RuntimeError(f"Worker failed; see retained {variant} stderr.")
        result=json.loads((target/"measurement.json").read_bytes())
        result["whole_seconds"]=whole
        measurements.append(result)
        (output/"attempts.json").write_text(json.dumps(measurements,indent=2)+"\n",encoding="utf-8")
        print(json.dumps({"index":index,"variant":variant,"whole_seconds":whole,"api_seconds":result["api_seconds"],
                          "trace_evaluations":result["trace_evaluations"],"counts":result["counts"]}),flush=True)
    exact=len({m["canonical_result_sha256"] for m in measurements})==1
    ratio=statistics.median(m["whole_seconds"] for m in measurements if m["variant"]=="numpy")/statistics.median(
        m["whole_seconds"] for m in measurements if m["variant"]=="cuda")
    summary={"exact_complete_results":exact,"whole_ratio":ratio,"passes_screen":exact and ratio>=2.0,
             "protocol_sha256":sha(output/"protocol.json"),"provenance":protocol["provenance"]}
    (output/"summary.json").write_text(json.dumps(summary,indent=2)+"\n",encoding="utf-8")
    print(json.dumps(summary),flush=True)


if __name__=="__main__":
    main()
