#!/usr/bin/env python3

# Copyright (c) 2017-2023 California Institute of Technology ("Caltech"). U.S.
# Government sponsorship acknowledged. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0

r'''Visualize the difference in projection between 2 camera pairs

SYNOPSIS

  $ mrcal-show-stereo-pair-diff      \
      --distance 2                   \
      before/camera-[01].cameramodel \
      after/camera-[01].cameramodel

  ... a plot pops up showing how these two pairs differ in their projections

THIS TOOL IS EXPERIMENTAL, AND LIKELY WILL BE CHANGED AND EXTENDED. MORE
COMPLETE DOCUMENTATION WILL BE WRITTEN LATER

It is useful to evaluate the stability and uncertainty of a multi-camera system,
looking at the intrinsics,extrinsics either separately or jointly. For instance,
when deployed in the field a calibration may shift in the extrinsics (camera
units moved) and/or intrinsics (lens moved), and it would be useful to quantify
them individually. Up to now (version 2.5) mrcal has been concerned mostly with
the intrinsics: the mean-pcam uncertainty method and
implied_Rt10__from_unprojections() try to compensate for and eliminate the
effect of extrinsic shifts in uncertainty and diff methods respectively.

This stereo_pair_diff() tool extends this in one way: in a multi-camera
calibration, it quantifies the joint intrinsics+extrinsics shift of a camera
pair. At a given distance this tool reprojects pixels in camera0 to camera1,
and compares the results of this reprojection. This is actually far simpler than
the earlier intrinsics-only methods: no implied transform or reference-frame
shift need to be considered. As of mrcal 2.5 we also have an extrinsics-only
stereo-pair diff in an unreleased analyses/extrinsics-stability.py tool. No
corresponding uncertainty method is implemented yet, but that is
straightforward. All the methods will be implemented in a future mrcal release.
They can then be studied to see what is useful to look at in practice.

This tool visualizes the results of mrcal.stereo_pair_diff()

'''

import sys
import argparse
import re
import os

def parse_args():

    parser = \
        argparse.ArgumentParser(description = __doc__,
                                formatter_class=argparse.RawDescriptionHelpFormatter)

    parser.add_argument('--gridn',
                        type=int,
                        default = (60,40),
                        nargs = 2,
                        help='''How densely we should sample the imager. By default we use a 60x40 grid''')
    parser.add_argument('--distance',
                        type=float,
                        help='''The projection difference varies depending on
                        the range to the observed world points, with the queried
                        range set in this argument. If omitted we look out to
                        infinity.''')

    parser.add_argument('--observations',
                        action='store_true',
                        default=False,
                        help='''If given, I show where the chessboard corners were observed at calibration
                        time. These should correspond to the low-diff regions.''')
    parser.add_argument('--valid-intrinsics-region',
                        action='store_true',
                        default=False,
                        help='''If given, I overlay the valid-intrinsics regions onto the plot''')
    parser.add_argument('--cbmax',
                        type=float,
                        default=4,
                        help='''Maximum range of the colorbar''')

    parser.add_argument('--title',
                        type=str,
                        default = None,
                        help='''Title string for the plot. Overrides the default
                        title. Exclusive with --extratitle''')
    parser.add_argument('--extratitle',
                        type=str,
                        default = None,
                        help='''Additional string for the plot to append to the
                        default title. Exclusive with --title''')

    parser.add_argument('--vectorfield',
                        action = 'store_true',
                        default = False,
                        help='''Plot the diff as a vector field instead of as a heat map. The vector field
                        contains more information (magnitude AND direction), but
                        is less clear at a glance''')

    parser.add_argument('--vectorscale',
                        type = float,
                        default = 1.0,
                        help='''If plotting a vectorfield, scale all the vectors by this factor. Useful to
                        improve legibility if the vectors are too small to
                        see''')


    parser.add_argument('--hardcopy',
                        type=str,
                        help='''Write the output to disk, instead of making an interactive plot''')
    parser.add_argument('--terminal',
                        type=str,
                        help=r'''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='''Extra 'set' directives to gnuplotlib. Can be given multiple times''')
    parser.add_argument('--unset',
                        type=str,
                        action='append',
                        help='''Extra 'unset' directives to gnuplotlib. Can be given multiple times''')

    parser.add_argument('models',
                        type=str,
                        nargs=4,
                        help='''4 camera models: 2 cameras from the first camera
                        pair, followed by 2 cameras from the other pair''')

    args = parser.parse_args()

    if len(args.models) != 4:
        print(f"I need exactly 4 models to diff. Instead got '{len(args.models)}'", file=sys.stderr)
        sys.exit(1)

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

    if args.vectorscale != 1.0 and not args.vectorfield:
        print("Error: --vectorscale only makes sense with --vectorfield",
              file = sys.stderr)
        sys.exit(1)

    return args

args = parse_args()

import mrcal
import numpy as np
import numpysane as nps


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

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

models = [openmodel(modelfilename) for modelfilename in args.models]

model_pairs = ( (models[0], models[1]),
                (models[2], models[3]) )


if args.observations:
    optimization_inputs = [ m.optimization_inputs() for model_pair in model_pairs for m in model_pair ]
    if any( oi is None for oi in optimization_inputs ):
        print("mrcal-show-stereo-pair-diff --observations requires optimization_inputs to be available for all models, but this is missing for some models",
              file=sys.stderr)
        sys.exit(1)

plot = mrcal.show_stereo_pair_diff(model_pairs,
                                   gridn_width             = args.gridn[0],
                                   gridn_height            = args.gridn[1],
                                   observations            = args.observations,
                                   valid_intrinsics_region = args.valid_intrinsics_region,
                                   distance                = args.distance,
                                   vectorfield             = args.vectorfield,
                                   vectorscale             = args.vectorscale,
                                   hardcopy                = args.hardcopy,
                                   terminal                = args.terminal,
                                   cbmax                   = args.cbmax,
                                   **plotkwargs_extra)[0]

if args.hardcopy is None:
    plot.wait()
