"""Synthetic PnP example. No camera, network, or real-world accuracy claim.
Run: python reproduce.py --output results.json
Dependencies: numpy, opencv-python-headless (cv2 >= 4.6).
"""
import argparse
import csv
import json
import platform
from pathlib import Path

import cv2
import numpy as np


def inputs():
    rng = np.random.default_rng(20261001)
    xyz = rng.uniform([-1.2, -0.9, -0.6], [1.2, 0.9, 0.6], (48, 3))
    K = np.array([[800., 0., 640.], [0., 800., 360.], [0., 0., 1.]])
    r = np.array([0.15, -0.10, 0.05], dtype=np.float64).reshape(3, 1)
    t = np.array([0.20, -0.10, 5.0], dtype=np.float64).reshape(3, 1)
    uv = cv2.projectPoints(xyz, r, t, K, None)[0].reshape(-1, 2)
    noisy = uv + rng.normal(0., 0.5, uv.shape)
    bad = noisy.copy()
    outlier_ids = rng.choice(len(xyz), 12, replace=False)
    bad[outlier_ids] = rng.uniform([0., 0.], [1280., 720.], (12, 2))
    good = np.ones(len(xyz), dtype=bool)
    good[outlier_ids] = False
    return xyz, K, r, t, noisy, bad, good


def estimate(xyz, uv, K, r_true, t_true, known_good, robust):
    cv2.setRNGSeed(20261001)
    if robust:
        ok, r, t, idx = cv2.solvePnPRansac(
            xyz, uv, K, None, iterationsCount=200, reprojectionError=3.,
            confidence=0.999, flags=cv2.SOLVEPNP_EPNP)
        if not ok or idx is None or len(idx) < 6:
            return {'status': 'failed'}
        idx = idx.ravel()
    else:
        ok, r, t = cv2.solvePnP(xyz, uv, K, None, flags=cv2.SOLVEPNP_EPNP)
        if not ok:
            return {'status': 'failed'}
        idx = np.arange(len(xyz))
    r, t = cv2.solvePnPRefineLM(xyz[idx], uv[idx], K, None, r, t)
    pred = cv2.projectPoints(xyz, r, t, K, None)[0].reshape(-1, 2)
    residual = np.linalg.norm(pred - uv, axis=1)
    R, _ = cv2.Rodrigues(r)
    Rt, _ = cv2.Rodrigues(r_true)
    rotation_error = np.degrees(np.arccos(np.clip((np.trace(R @ Rt.T)-1)/2, -1, 1)))
    center = (-R.T @ t).ravel()
    center_true = (-Rt.T @ t_true).ravel()
    depth = (xyz @ R.T + t.reshape(1, 3))[:, 2]
    return {
        'status': 'ok', 'fit_points': int(len(idx)),
        'fit_rmse_px': float(np.sqrt(np.mean(residual[idx]**2))),
        'known_good_rmse_px': float(np.sqrt(np.mean(residual[known_good]**2))),
        'all_rmse_px': float(np.sqrt(np.mean(residual**2))),
        'rotation_error_deg': float(rotation_error),
        'camera_center_error_m': float(np.linalg.norm(center-center_true)),
        'min_depth_m': float(depth.min()),
        'camera_center_world_m': center.tolist(),
        'fit_indices': idx.tolist(), 'residual_px': residual.tolist(),
    }


def run():
    xyz, K, r, t, noisy, bad, good = inputs()
    wrong_K = K.copy()
    wrong_K[0, 0] *= 1.15
    wrong_K[1, 1] *= 1.15
    cases = [
        ('clean_ransac', noisy, K, np.ones(len(xyz), dtype=bool), True),
        ('outliers_plain', bad, K, good, False),
        ('outliers_ransac', bad, K, good, True),
        ('wrong_focal_ransac', noisy, wrong_K, np.ones(len(xyz), dtype=bool), True),
    ]
    return {
        'environment': {'python': platform.python_version(), 'opencv': cv2.__version__, 'numpy': np.__version__},
        'settings': {'seed': 20261001, 'points': 48, 'outliers': 12, 'noise_sigma_px': 0.5,
                     'image_size_px': [1280, 720], 'K_true': K.tolist(),
                     'rvec_true': r.ravel().tolist(), 'tvec_true_m': t.ravel().tolist(),
                     'wrong_focal_scale': 1.15, 'distortion': 'none',
                     'ransac_threshold_px': 3., 'ransac_iterations': 200, 'confidence': 0.999},
        'cases': {name: estimate(xyz, uv, camera, r, t, mask, robust)
                  for name, uv, camera, mask, robust in cases},
    }


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--output', type=Path, default=Path('results.json'))
    parser.add_argument('--inputs', type=Path, help='Optionally export generated 3D/2D points as CSV')
    args = parser.parse_args()
    args.output.write_text(json.dumps(run(), indent=2, allow_nan=False)+'\n', encoding='utf-8')
    if args.inputs:
        xyz, _, _, _, noisy, bad, good = inputs()
        with args.inputs.open('w', newline='', encoding='utf-8') as f:
            writer = csv.writer(f)
            writer.writerow(['id','x_m','y_m','z_m','clean_u_px','clean_v_px','corrupt_u_px','corrupt_v_px','known_good'])
            for i in range(len(xyz)):
                writer.writerow([i,*xyz[i],*noisy[i],*bad[i],int(good[i])])
    print(args.output)