r"""Converts a camera model from one lens model to another

SYNOPSIS

  $ mrcal-convert-lensmodel
      --viz LENSMODEL_OPENCV4 left.cameramodel

  ... lots of output as the solve runs ...
  RMS error of the solution: 3.40256580058 pixels.
  Wrote 'left-LENSMODEL_OPENCV4.cameramodel'

  ... a plot pops up showing the differences ...

Given one camera model, this tool computes another camera model that represents
the same camera, but utilizes a different lens model. While lens models all
exist to solve the same problem, the different representations don't map to one
another perfectly, and this tool finds the best-fitting parameters of the target
lens model. Two different methods are implemented:

1. If the given cameramodel file contains optimization_inputs, then we have all
   the data that was used to compute this model in the first place, and we can
   re-run the original optimization, using the new lens model. This is the
   default behavior, and is the preferred choice. However it can only work with
   models that were computed by mrcal originally. We re-run the full original
   solve, even it contained multiple cameras, unless --monocular is given. With
   that option, we re-solve only the subset of the images observed by the one
   requested camera

2. We can sample a grid of points on the imager, unproject them to observation
   vectors in the camera coordinate system, and then fit a new camera model that
   reprojects these vectors as closely to the original pixel coordinates as
   possible. This can be applied to models that didn't come from mrcal. Select
   this mode by passing --sampled.

Since camera models (lens parameters AND geometry) are computed off real pixel
observations, the confidence of the projections varies greatly across the imager
and across observation distances. The first method uses the original data, so it
implicitly respects these uncertainties: uncertain areas in the original model
will be uncertain in the new model as well. The second method, however, doesn't
have this information: it doesn't know which parts of the imager and space are
reliable, so the results suffer.

As always, the intrinsics have some baked-in geometry information. Both methods
optimize intrinsics AND extrinsics, and output cameramodels with updated
versions of both. If --sampled: we can request that only the intrinsics be
optimized by passing --intrinsics-only.

Also, if --sampled and not --intrinsics-only: we fit the extrinsics off 3D
points, not just observation directions. The distance from the camera to the
points is set by --distance. This can take a comma-separated list of distances
to use. It's STRONGLY recommended to ask for two different distances:

- A "near" distance: where we expect the intrinsics to have the most accuracy.
  At the range of the chessboards, usually

- A "far" distance: at "infinity". A few km is good usually.

The reason for this is that --sampled solves at a single distance aren't
sufficiently constrained. If we ask for a single far distance: "--distance 1000"
for instance, we can easily get an extrinsics shift of 100m. This is aphysical:
changing the intrinsics could shift the camera origin by a few mm, but not 100m.
Conceptually we want to perform a rotation-only extrinsics solve, but this isn't
yet implemented. Passing both a near and far distance appears to constrain the
extrinsics well in practice. The computed extrinsics transform is printed on the
console, with a warning if an aphysical shift was computed. Do pay attention to
the console output.

Sampled solves are sometimes sensitive to the optimization seed. To control for
this, pass --num-trials to evaluate the solve multiple times from different
random seeds, and to pick the best one. These solves are usually quick, so
there's no harm in passing something like "--num-trials 10".

We need to consider that the model we're trying to fit may not fit the original
model in all parts of the imager. Usually this is a factor when converting
wide-angle cameramodels to use a leaner model: a decent fit will be possible at
the center, with more and more divergence as we move towards the edges. We
handle this with the --where and --radius options to allow the user to select
the area of the imager that is used for the fit: observations outside the
selected area are thrown out. This region is centered on the point given by
--where (or at the center of the imager, if --where is omitted). The radius of
this region is given by --radius. If '--radius 0' then we use ALL the data. A
radius<0 can be used to set the size of the no-data margin at the corners; in
this case I'll use

    r = sqrt(width^2 + height^2)/2. - abs(radius)

There's a balance to strike here. A larger radius means that we'll try to fit as
well as we can in a larger area. This might mean that we won't fit well
anywhere, but we won't do terribly anywhere, either. A smaller area means that
we give up on the outer regions entirely (resulting in very poor fits there),
but we'll be able to fit much better in the areas that remain. Generally
empirical testing is required to find a good compromise: pass --viz to see the
resulting differences. Note that --radius and --where applies only if we're
optimizing sampled reprojections; if we're using the original optimization
inputs, the options are illegal.

The output is written to a file on disk, with the same filename as the input
model, but with the new lensmodel added to the filename.

"""

import sys
import argparse
from typing import Any
import re
import os
import numpy as np
import numpysane as nps
import time
import copy
import shlex

import drcal


def parse_args():
    parser = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )

    parser.add_argument(
        "--sampled",
        action="store_true",
        help="""Instead of solving the original calibration problem using the new lens model,
                        use sampled imager points. This produces biased results,
                        but can be used even if the original optimization_inputs
                        aren't available""",
    )
    parser.add_argument(
        "--gridn",
        type=int,
        default=(30, 20),
        nargs=2,
        help="""Used if --sampled. How densely we should sample the imager. By default we use
                        a 30x20 grid""",
    )
    parser.add_argument(
        "--distance",
        type=str,
        help="""Required if --sampled and not --intrinsics-only.
                        A sampled solve fits the intrinsics and extrinsics to
                        match up reprojections of a grid of observed pixels. The
                        points being projected are set a particular distance
                        (set by this argument) from the camera. Set this to the
                        distance that is expected to be most confident for the
                        given cameramodel. Points at infinity aren't supported
                        yet: specify a high distance instead. We can fit
                        multiple distances at the same time by passing them here
                        in a comma-separated, whitespace-less list. If multiple
                        distances are given, we fit the model using ALL the
                        distances, but --viz displays the difference for the
                        FIRST distance given. See the description above. Without
                        --sampled, this is used for the visualization only""",
    )
    parser.add_argument(
        "--intrinsics-only",
        action="store_true",
        help="""Used if --sampled. By default I optimize the
                        intrinsics and extrinsics to find the closest
                        reprojection. If for whatever reason we know that the
                        camera coordinate system was already right, or we need
                        to keep the original extrinsics, pass --intrinsics-only.
                        The resulting extrinsics will be the same, but the fit
                        will not be as good. In many cases, optimizing
                        extrinsics is required to get a usable fit, so
                        --intrinsics-only may not be an option if accurate
                        results are required.""",
    )

    parser.add_argument(
        "--where",
        type=float,
        nargs=2,
        help="""Used with or without --sampled. I use a subset
                        of the imager to compute the fit. The active region is a
                        circle centered on this point. If omitted, we will focus
                        on the center of the imager""",
    )
    parser.add_argument(
        "--radius",
        type=float,
        help="""Used with or without --sampled. I use a subset
                        of the imager to compute the fit. The active region is a
                        circle with a radius given by this parameter. If radius
                        == 0, I'll use the whole imager for the fit. If radius <
                        0, this parameter specifies the width of the region at
                        the corners that I should ignore: I will use
                        sqrt(width^2 + height^2)/2. - abs(radius). This is valid
                        ONLY if we're focusing at the center of the imager. By
                        default I ignore a large-ish chunk area at the corners.""",
    )

    parser.add_argument(
        "--monocular",
        action="store_true",
        help="""Used if not --sampled. By default we re-solve
                        the full calibration problem that was used to originally
                        obtain the input model, even if it contained multiple
                        cameras. If --monocular, we re-solve only a subset of
                        the original problem, using ONLY the observations made
                        by THIS camera""",
    )

    parser.add_argument(
        "--viz",
        action="store_true",
        help="""Visualize the differences between the input and
                        output models. If we have --distance, the FIRST given
                        distance is used to display the fit error""",
    )
    parser.add_argument(
        "--cbmax", type=float, default=4, help="""Maximum range of the colorbar"""
    )
    parser.add_argument(
        "--title",
        type=str,
        default=None,
        help="""Used if --viz. Title string for the diff plot.
                        Overrides the default title. Exclusive with
                        --extratitle""",
    )
    parser.add_argument(
        "--extratitle",
        type=str,
        default=None,
        help="""Used if --viz. Additional string for the plot to
                        append to the default title. Exclusive with --title""",
    )
    parser.add_argument(
        "--hardcopy",
        type=str,
        help="""Used if --viz. Write the diff output to disk,
                        instead of making an interactive plot""",
    )
    parser.add_argument(
        "--terminal",
        type=str,
        help=r"""Used if --viz. gnuplotlib terminal. The default
                        is good almost always, so most people don't need this
                        option""",
    )
    parser.add_argument(
        "--set",
        type=str,
        action="append",
        help="""Used if --viz. Extra 'set' directives to
                        gnuplotlib. Can be given multiple times""",
    )
    parser.add_argument(
        "--unset",
        type=str,
        action="append",
        help="""Used if --viz. Extra 'unset' directives to
                        gnuplotlib. Can be given multiple times""",
    )

    parser.add_argument(
        "--force",
        "-f",
        action="store_true",
        default=False,
        help="""By default existing models on disk are not
                        overwritten. Pass --force to overwrite them without
                        complaint""",
    )
    parser.add_argument(
        "--outdir",
        type=lambda d: d
        if os.path.isdir(d)
        else parser.error(
            f"--outdir requires an existing directory as the arg, but got '{d}'"
        ),
        help="""Directory to write the output into. If omitted,
                        we use the directory of the input model""",
    )

    parser.add_argument(
        "--num-trials",
        type=int,
        default=1,
        help="""If given, run the solve more than once. Useful
                        in case random initialization produces noticeably
                        different results. By default we run just one trial,
                        which is enough most of the time""",
    )

    parser.add_argument("to", type=str, help="The target lens model")

    parser.add_argument(
        "model",
        default="-",
        nargs="?",
        type=str,
        help="""Input camera model. If omitted or "-", we read
                        standard input and write to standard output""",
    )

    args = parser.parse_args()

    if args.title is not None and args.extratitle is not None:
        print("Error: --title and --extratitle are exclusive", file=sys.stderr)
        sys.exit(1)

    if args.distance is not None:
        import math

        try:
            args.distance = [float(d) for d in args.distance.split(",")]
        except:
            print(
                "Error: distances must be given a comma-separated list of floats in --distance",
                file=sys.stderr,
            )
            sys.exit(1)

        if any(not math.isfinite(x) or x <= 0 for x in args.distance):
            print("All values in --distances must be finite and > 0", file=sys.stderr)
            sys.exit(1)

    return args


def main():
    args = parse_args()

    # arg-parsing is done before the imports so that --help works without building
    # stuff, so that I can generate the manpages and README

    lensmodel_to = args.to

    try:
        meta = drcal.lensmodel_metadata_and_config(lensmodel_to)
    except Exception as e:
        print(
            f"Invalid lens model '{lensmodel_to}': couldn't get the metadata: {e}",
            file=sys.stderr,
        )
        sys.exit(1)
    if not meta["has_gradients"]:
        print(
            f"lens model {lensmodel_to} is not supported at this time: its gradients aren't implemented",
            file=sys.stderr,
        )
        sys.exit(1)

    try:
        Ndistortions = drcal.lensmodel_num_params(lensmodel_to) - 4
    except:
        print(f"Unknown lens model: '{lensmodel_to}'", file=sys.stderr)
        sys.exit(1)

    try:
        m = drcal.cameramodel(args.model)
    except Exception as e:
        print(f"Couldn't load camera model '{args.model}': {e}", file=sys.stderr)
        sys.exit(1)

    if args.model == "-":
        # Input read from stdin. Write output to stdout
        file_output = sys.stdout
    else:
        if args.outdir is None:
            filename_base, extension = os.path.splitext(args.model)
            file_output = f"{filename_base}-{lensmodel_to}{extension}"
        else:
            f, extension = os.path.splitext(args.model)
            directory, filename_base = os.path.split(f)
            file_output = f"{args.outdir}/{filename_base}-{lensmodel_to}{extension}"

        if os.path.exists(file_output) and not args.force:
            print(
                f"ERROR: '{file_output}' already exists. Not overwriting this file. Pass -f to overwrite",
                file=sys.stderr,
            )
            sys.exit(1)

    lensmodel_from = m.assume_intrinsics()[0]

    if lensmodel_from == lensmodel_to:
        print(
            f"Input and output have the same lens model {lensmodel_from}. Nothing to do",
            file=sys.stderr,
        )
        print("RMS error of the solution: 0 pixels.", file=sys.stderr)
        sys.exit(0)

    if not args.sampled:
        if args.intrinsics_only:
            print("--intrinsics-only requires --sampled", file=sys.stderr)
            sys.exit(1)
        if not (
            drcal.lensmodel_metadata_and_config(lensmodel_from)["has_core"]
            and drcal.lensmodel_metadata_and_config(lensmodel_to)["has_core"]
        ):
            print(
                "Without --sampled, the TO and FROM models must support contain an intrinsics core. It COULD work otherwise, but somebody needs to implement it",
                file=sys.stderr,
            )
            sys.exit(1)

        optimization_inputs_maybe = m.optimization_inputs()
        if optimization_inputs_maybe is None:
            print(
                "optimization_inputs not available in this model, so only sampled fits are possible. Pass --sampled",
                file=sys.stderr,
            )
            sys.exit(1)

        optimization_inputs = optimization_inputs_maybe

        icam_intrinsics = m.icam_intrinsics()

        if args.monocular:
            if (
                optimization_inputs["indices_point_camintrinsics_camextrinsics"]
                is not None
                and optimization_inputs[
                    "indices_point_camintrinsics_camextrinsics"
                ].size
                > 0
            ):
                print(
                    "--monocular can be used ONLY to re-optimize vanilla calibration problems. This one has points",
                    file=sys.stderr,
                )
                sys.exit(1)

            intrinsics = optimization_inputs["intrinsics"]
            imagersizes = optimization_inputs["imagersizes"]
            extrinsics_rt_fromref = optimization_inputs["extrinsics_rt_fromref"]
            frames_rt_toref = optimization_inputs["frames_rt_toref"]
            indices_frame_camintrinsics_camextrinsics = optimization_inputs[
                "indices_frame_camintrinsics_camextrinsics"
            ]
            observations_board = optimization_inputs["observations_board"]

            imagepaths = None
            if "imagepaths" in optimization_inputs:
                imagepaths = optimization_inputs["imagepaths"]

            mask_observations = (
                indices_frame_camintrinsics_camextrinsics[:, 1] == icam_intrinsics
            )
            indices_frame_camintrinsics_camextrinsics = (
                indices_frame_camintrinsics_camextrinsics[mask_observations]
            )
            observations_board = observations_board[mask_observations]

            if "imagepaths" in optimization_inputs:
                assert imagepaths is not None
                imagepaths = imagepaths[mask_observations]

            ### Now I must cull the extrinsics and frames to only use these
            ### observations

            # intrinsics
            indices_frame_camintrinsics_camextrinsics[:, 1] = 0
            intrinsics = intrinsics[(icam_intrinsics,), :]
            imagersizes = imagersizes[(icam_intrinsics,), :]

            # extrinsics
            icam_extrinsics = indices_frame_camintrinsics_camextrinsics[0, 2]
            if not np.all(
                indices_frame_camintrinsics_camextrinsics[:, 2] == icam_extrinsics
            ):
                print(
                    "Error: --monocular can work ONLY if the calibration-time cameras are stationary. Here the requested camera was moving",
                    file=sys.stderr,
                )
                sys.exit(1)
            indices_frame_camintrinsics_camextrinsics[:, 2] = -1
            if icam_extrinsics >= 0:
                # I'm moving the reference to lie at this one new camera. Transform
                # everything
                rt_cam_reforig = extrinsics_rt_fromref[icam_extrinsics]
                rt_refnew_reforig = rt_cam_reforig

                frames_rt_toref = drcal.compose_rt(rt_refnew_reforig, frames_rt_toref)
                extrinsics_rt_fromref = np.empty((0, 6), dtype=float)

            # frames
            frames_rt_toref = frames_rt_toref[
                indices_frame_camintrinsics_camextrinsics[:, 0], :
            ]
            Nframes = len(indices_frame_camintrinsics_camextrinsics)
            indices_frame_camintrinsics_camextrinsics[:, 0] = np.arange(Nframes)

            # store everythin back into the inputs
            optimization_inputs["intrinsics"] = intrinsics
            optimization_inputs["imagersizes"] = imagersizes
            optimization_inputs["extrinsics_rt_fromref"] = extrinsics_rt_fromref
            optimization_inputs["frames_rt_toref"] = frames_rt_toref
            optimization_inputs["indices_frame_camintrinsics_camextrinsics"] = (
                indices_frame_camintrinsics_camextrinsics
            )
            optimization_inputs["observations_board"] = observations_board
            if "imagepaths" in optimization_inputs:
                optimization_inputs["imagepaths"] = imagepaths

            icam_intrinsics = 0

        optimization_inputs_before = copy.deepcopy(optimization_inputs)

        intrinsics_from = optimization_inputs["intrinsics"]
        intrinsics_from_core = intrinsics_from[..., :4]
        Ncameras = len(intrinsics_from)

        # Ignore observations in the corners, as requested
        if not (args.radius is None or args.radius == 0):
            for icam in range(Ncameras):
                dims = optimization_inputs["imagersizes"][icam].astype(float)

                if args.where is None:
                    focus_center = (dims - 1) / 2
                else:
                    focus_center = args.where

                if args.radius > 0:
                    r = args.radius
                else:
                    if nps.norm2(focus_center - (dims - 1.0) / 2) > 1e-3:
                        print(
                            "A radius <0 is only implemented if we're focusing on the imager center",
                            file=sys.stderr,
                        )
                        sys.exit(1)
                    r = nps.mag(dims) / 2.0 + args.radius

                indices_cam = optimization_inputs[
                    "indices_frame_camintrinsics_camextrinsics"
                ][:, 1]
                observations = optimization_inputs["observations_board"]

                mask_thiscam = indices_cam == icam
                mask_past_radius = (
                    nps.norm2(observations[..., :2] - focus_center) > r * r
                )

                observations[nps.mv(mask_thiscam, -1, -3) * mask_past_radius, 2] = -1

        optimization_inputs["verbose"] = False

        optimization_inputs["lensmodel"] = lensmodel_to
        optimization_inputs["intrinsics"] = np.zeros(
            (Ncameras, Ndistortions + 4), dtype=float
        )
        optimization_inputs["intrinsics"][:, :4] = intrinsics_from_core

        # I do this in stages, similar to how mrcal-calibrate-cameras does it. First
        # just the frames and extrinsics. Assuming a core-only intrinsics
        optimization_inputs["do_optimize_intrinsics_core"] = False
        optimization_inputs["do_optimize_intrinsics_distortions"] = False
        optimization_inputs["do_optimize_calobject_warp"] = False
        optimization_inputs["do_apply_outlier_rejection"] = False
        optimization_inputs["do_apply_regularization"] = False
        stats = drcal.optimize(**optimization_inputs)

        # Then I optimize the core also
        optimization_inputs["do_optimize_intrinsics_core"] = True
        stats = drcal.optimize(**optimization_inputs)

        # Then the intrinsics too
        optimization_inputs["do_optimize_intrinsics_distortions"] = True
        optimization_inputs["do_apply_outlier_rejection"] = True
        optimization_inputs["do_apply_regularization"] = True
        if re.match("LENSMODEL_SPLINED_STEREOGRAPHIC_", lensmodel_to):
            # splined models have a core, but those variables are largely redundant
            # with the spline parameters. So I lock down the core when targeting
            # splined models
            optimization_inputs["do_optimize_intrinsics_core"] = False
        # stolen expand_intrinsics() in mrcal-calibrate-extrinsics. Please consolidate
        optimization_inputs["intrinsics"][:, 4:] = (
            (np.random.random((Ncameras, Ndistortions)) - 0.5) * 2.0 * 1e-6
        )
        modelmatch = re.search("OPENCV([0-9]+)", lensmodel_to)
        if modelmatch:
            Nd = int(modelmatch.group(1))
            if Nd >= 8:
                # Push down the rational components of the seed. I'd like these all to
                # sit at 0 ideally. The radial distortion in opencv is x_distorted =
                # x*scale where r2 = norm2(xy - xyc) and
                #
                # scale = (1 + k0 r2 + k1 r4 + k4 r6)/(1 + k5 r2 + k6 r4 + k7 r6)
                #
                # Note that k2,k3 are tangential (NOT radial) distortion components.
                # Note that the r6 factor in the numerator is only present for
                # >=LENSMODEL_OPENCV5. Note that the denominator is only present for >=
                # LENSMODEL_OPENCV8. The danger with a rational model is that it's
                # possible to get into a situation where scale ~ 0/0 ~ 1. This would
                # have very poorly behaved derivatives. If all the rational coefficients
                # are ~0, then the denominator is always ~1, and this problematic case
                # can't happen. I favor that.
                optimization_inputs["intrinsics"][:, 4:][:, 5:8] *= 1e-3
        stats = drcal.optimize(**optimization_inputs)

        # And finally I do the calobject_warp too
        optimization_inputs["do_optimize_calobject_warp"] = True
        stats = drcal.optimize(**optimization_inputs)

        # I pick the inlier set using the post-solve inliers/outliers. The solve
        # could add some outliers, but it won't take any away, so the union of the
        # before/after inlier sets is just the post-solve set
        idx_inliers_after = optimization_inputs["observations_board"][..., 2] > 0

        calobject_camframe_before = drcal.hypothesis_board_corner_positions(
            icam_intrinsics, **optimization_inputs_before, idx_inliers=idx_inliers_after
        )[-2]
        calobject_camframe_after = drcal.hypothesis_board_corner_positions(
            icam_intrinsics, **optimization_inputs, idx_inliers=idx_inliers_after
        )[-2]

        # This is a procrustes-based transformation. This transform is usable, but
        # isn't based on camera observations, and often produces poor diffs if used
        # directly. Let mrcal.projection_diff() figure out the implied_Rt10 instead
        # of using this
        implied_Rt10 = drcal.align_procrustes_points_Rt01(
            calobject_camframe_after, calobject_camframe_before
        )

        # For the purposes of visualization I use what I was given. Or the mean
        # observation distance if I was given nothing
        if args.distance is None:
            distance_for_diff = np.mean(nps.mag(calobject_camframe_after))
        else:
            distance_for_diff = np.array(args.distance)

        print(
            f"RMS error of the solution: {stats['rms_reproj_error__pixels']} pixels.",
            file=sys.stderr,
        )
        m_to = drcal.cameramodel(
            optimization_inputs=optimization_inputs, icam_intrinsics=icam_intrinsics
        )

    else:
        # Sampled solve. I grid the imager, unproject and fit

        intrinsics_from = m.assume_intrinsics()

        if args.distance is None:
            if args.intrinsics_only:
                # We're only looking at the intrinsics, so the distance doesn't
                # matter. I pick an arbitrary value
                args.distance = [10.0]
            else:
                print(
                    "--sampled without --intrinsics-only REQUIRES --distance",
                    file=sys.stderr,
                )
                sys.exit(1)

        if args.monocular:
            print("--monocular doesn't work with --sampled", file=sys.stderr)
            sys.exit(1)

        dims = m.imagersize()
        if dims is None:
            print(
                "Warning: imager size not available. Using centerpixel*2",
                file=sys.stderr,
            )
            dims = intrinsics_from[1][2:4] * 2

        if args.radius is None:
            # By default use 1/4 of the smallest dimension
            args.radius = -np.min(m.imagersize()) // 4
            print(
                f"Default radius: {args.radius}. We're ignoring the regions {-args.radius} pixels from each corner",
                file=sys.stderr,
            )
            if (
                args.where is not None
                and nps.norm2(args.where - (dims - 1.0) / 2) > 1e-3
            ):
                print(
                    "A radius <0 is only implemented if we're focusing on the imager center: use an explicit --radius, or omit --where",
                    file=sys.stderr,
                )
                sys.exit(1)

        # Alrighty. Let's actually do the work. I do this:
        #
        # 1. Sample the imager space with the known model
        # 2. Unproject to get the 3d observation vectors
        # 3. Solve a new model that fits those vectors to the known observations, but
        #    using the new model

        ### I sample the pixels in an NxN grid
        Nx, Ny = args.gridn

        qx = np.linspace(0, dims[0] - 1, Nx)
        qy = np.linspace(0, dims[1] - 1, Ny)

        # q is (Ny*Nx, 2). Each slice of q[:] is an (x,y) pixel coord
        q = np.ascontiguousarray(
            nps.transpose(nps.clump(nps.cat(*np.meshgrid(qx, qy)), n=-2))
        )
        if args.radius != 0:
            # we use a subset of the input data for the fit
            if args.where is None:
                focus_center = (dims - 1.0) / 2.0
            else:
                focus_center = args.where

            if args.radius > 0:
                r = args.radius
            else:
                if nps.norm2(focus_center - (dims - 1.0) / 2) > 1e-3:
                    print(
                        "A radius <0 is only implemented if we're focusing on the imager center",
                        file=sys.stderr,
                    )
                    sys.exit(1)
                r = nps.mag(dims) / 2.0 + args.radius

            grid_off_center = q - focus_center
            i = nps.norm2(grid_off_center) < r * r
            q = q[i, ...]

        # To visualize the sample grid:
        # import gnuplotlib as gp
        # gp.plot(q[:,0], q[:,1], _with='points pt 7 ps 2', xrange=[0,3904],yrange=[3904,0], wait=1, square=1)
        # sys.exit()

        intrinsics_from = m.assume_intrinsics()

        ### I unproject this, with broadcasting
        # shape (Ndistances,)
        d = np.array(args.distance)
        # shape (Ndistances, Ny*Nx, 2)
        q = q * np.ones((d.size, 1, 1))
        # shape (Ndistances, Ny*Nx, 3)
        p = drcal.unproject(q, *intrinsics_from, normalize=True) * nps.mv(d, -1, -3)

        # shape (Ndistances*Ny*Nx, 2)
        q = nps.clump(q, n=2)
        # shape (Ndistances*Ny*Nx, 3)
        p = nps.clump(p, n=2)

        # The list of distances. The meaning is the same as expected by
        # mrcal.show_projection_diff(): we visualize the diff of the FIRST distance
        distance_for_diff = d

        # Ignore any failed unprojections
        i_finite = np.isfinite(p[:, 0])
        p = p[i_finite]
        q = q[i_finite]
        Npoints = len(q)
        weights = np.ones((Npoints,), dtype=float)

        ### Solve!

        ### I solve the optimization a number of times with different random seed
        ### values, taking the best-fitting results. This is required for the richer
        ### models such as LENSMODEL_OPENCV8
        err_rms_best = 1e10
        intrinsics_data_best = None
        extrinsics_rt_fromref_best = None

        for i in range(args.num_trials):
            # random seed for the new intrinsics
            intrinsics_core = intrinsics_from[1][:4]
            distortions = (
                np.random.rand(Ndistortions) - 0.5
            ) * 1e-3  # random initial seed
            intrinsics_to_values = nps.dummy(
                nps.glue(intrinsics_core, distortions, axis=-1), axis=-2
            )
            # each point has weight 1.0
            observations_points = nps.glue(q, nps.transpose(weights), axis=-1)
            observations_points = np.ascontiguousarray(
                observations_points
            )  # must be contiguous. mrcal.optimize() should really be more lax here

            # Which points we're observing. This is dense and kinda silly for this
            # application. Each slice is (i_point,i_camera,i_camera-1). Initially O
            # do everything in camera-0 coordinates, and I do not move the
            # extrinsics
            indices_point_camintrinsics_camextrinsics = np.zeros(
                (Npoints, 3), dtype=np.int32
            )
            indices_point_camintrinsics_camextrinsics[:, 0] = np.arange(
                Npoints, dtype=np.int32
            )
            indices_point_camintrinsics_camextrinsics[:, 1] = 0
            indices_point_camintrinsics_camextrinsics[:, 2] = -1

            optimization_inputs: dict[str, Any] = dict(
                intrinsics=intrinsics_to_values,
                extrinsics_rt_fromref=None,
                frames_rt_toref=None,  # no frames. Just points
                points=p,
                observations_board=None,  # no board observations
                indices_frame_camintrinsics_camextrinsics=None,  # no board observations
                observations_point=observations_points,
                indices_point_camintrinsics_camextrinsics=indices_point_camintrinsics_camextrinsics,
                lensmodel=lensmodel_to,
                imagersizes=nps.atleast_dims(dims, -2),
                # I'm not optimizing the point positions (frames), so these
                # need to be set to be inactive, and to include the ranges I do
                # have
                point_min_range=np.min(d) * 0.9,
                point_max_range=np.max(d) * 1.1,
                # I optimize the lens parameters. That's the whole point
                do_optimize_intrinsics_core=True,
                do_optimize_intrinsics_distortions=True,
                do_optimize_extrinsics=False,
                # NOT optimizing the observed point positions
                do_optimize_frames=False,
            )

            if re.match("LENSMODEL_SPLINED_STEREOGRAPHIC_", lensmodel_to):
                # splined models have a core, but those variables are largely redundant
                # with the spline parameters. So I lock down the core when targetting
                # splined models
                optimization_inputs["do_optimize_intrinsics_core"] = False

            stats = drcal.optimize(
                **optimization_inputs,
                # No outliers. I have the points that I have
                do_apply_outlier_rejection=False,
                verbose=False,
            )

            if not args.intrinsics_only:
                # go again, but refine this solution, allowing us to fit the
                # extrinsics too
                optimization_inputs["indices_point_camintrinsics_camextrinsics"][
                    :, 2
                ] = np.zeros((Npoints,), dtype=np.int32)
                optimization_inputs["extrinsics_rt_fromref"] = (
                    np.random.rand(1, 6) - 0.5
                ) * 1e-6
                optimization_inputs["do_optimize_extrinsics"] = True

                stats = drcal.optimize(
                    **optimization_inputs,
                    # No outliers. I have the points that I have
                    do_apply_outlier_rejection=False,
                    verbose=False,
                )

            err_rms = stats["rms_reproj_error__pixels"]
            print(f"RMS error of this solution: {err_rms} pixels.", file=sys.stderr)
            if err_rms < err_rms_best:
                err_rms_best = err_rms
                intrinsics_data_best = optimization_inputs["intrinsics"][0, :].copy()
                if not args.intrinsics_only:
                    extrinsics_rt_fromref_best = optimization_inputs[
                        "extrinsics_rt_fromref"
                    ][0, :].copy()

        if intrinsics_data_best is None:
            print("No valid intrinsics found!", file=sys.stderr)
            sys.exit(1)

        if args.num_trials > 1:
            print(
                f"RMS error of the BEST solution: {err_rms_best} pixels.",
                file=sys.stderr,
            )

        m_to = drcal.cameramodel(
            intrinsics=(lensmodel_to, intrinsics_data_best.ravel()), imagersize=dims
        )

        if args.intrinsics_only:
            implied_Rt10 = drcal.identity_Rt()
        else:
            implied_Rt10 = drcal.Rt_from_rt(extrinsics_rt_fromref_best)

    if not (args.sampled and args.intrinsics_only):
        implied_rt10 = drcal.rt_from_Rt(implied_Rt10)
        print(
            f"Transformation cam_fitted <-- cam_input:  rotation: {nps.mag(implied_rt10[:3]) * 180.0 / np.pi:.03f} degrees, translation: {np.array2string(implied_rt10[3:], precision=2)} m",
            file=sys.stderr,
        )

        dist_shift = nps.mag(implied_rt10[3:])
        if dist_shift > 0.01:
            msg = f"## WARNING: fitted camera moved by {dist_shift:.03f}m. This is probably aphysically high, and something is wrong."
            if args.distance is not None and args.sampled:
                print(msg + " Pass both a high and a low --distance?", file=sys.stderr)
            else:
                print(msg, file=sys.stderr)

    note = "generated on {} with   {}\n".format(
        time.strftime("%Y-%m-%d %H:%M:%S"), " ".join(shlex.quote(s) for s in sys.argv)
    )

    m_to.extrinsics_Rt_fromref(
        drcal.compose_Rt(implied_Rt10, drcal.Rt_from_rt(m.extrinsics_rt_fromref()))
    )
    m_to.write(file_output, note=note)

    if isinstance(file_output, str):
        print(f"Wrote '{file_output}'", file=sys.stderr)

    if args.viz:
        plotkwargs_extra = {}
        if args.set is not None:
            plotkwargs_extra["set"] = args.set
        if args.unset is not None:
            plotkwargs_extra["unset"] = args.unset

        if args.title is not None:
            plotkwargs_extra["title"] = args.title
        if args.extratitle is not None:
            plotkwargs_extra["extratitle"] = args.extratitle

        # If this is a sampled solve, we don't have projection uncertainties. But we
        # do have an implied_Rt10 based on observations, and we can feed that to the
        # diff directly.
        #
        # If this is NOT a sampled solve, then the implied_Rt10 comes from a
        # Procrustes fit (minimizing error in 3D space), which is usually not
        # close-enough to get a diff. We let the diff method figure out the
        # implied_Rt10, and we use the uncertainties that we have to do it
        if args.sampled:
            use_uncertainties = False
        else:
            implied_Rt10 = None
            use_uncertainties = True

        plot, _ = drcal.show_projection_diff(
            (m, m_to),
            implied_Rt10=implied_Rt10,
            distance=distance_for_diff,
            cbmax=args.cbmax,
            gridn_width=None if args.gridn is None else args.gridn[0],
            gridn_height=None if args.gridn is None else args.gridn[1],
            use_uncertainties=use_uncertainties,
            hardcopy=args.hardcopy,
            terminal=args.terminal,
            **plotkwargs_extra,
        )
        if args.hardcopy is None:
            plot.wait()  # type: ignore


if __name__ == "__main__":
    main()
