#!/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'''Visualizes the surface represented in a splined lens model

SYNOPSIS

  $ mrcal-show-splined-model-correction cam.cameramodel
  ... a plot pops up showing the correction magnitude heatmap

Splined models are parametrized by flexible surfaces that define the projection
corrections (off some baseline model), and visualizing these corrections is
useful for understanding the projection behavior. Details of these models are
described in the documentation:

http://mrcal.secretsauce.net/lensmodels.html#splined-stereographic-lens-model

At this time LENSMODEL_SPLINED_STEREOGRAPHIC is the only splined model mrcal
has, so the baseline model is always LENSMODEL_STEREOGRAPHIC.

This tool produces a plot in the domain either of the input or the output of the
spline functions.

By default:
    The plot is presented based on the spline index. With
    LENSMODEL_SPLINED_STEREOGRAPHIC, this is the stereographic projection u.
    This is the "forward" direction, what the projection operation actually
    computes. In this view the knots form a regular grid, and the edge of the
    imager forms a (possibly very irregular) curve

if --imager-domain:
    The plot is presented based on the pixels in the imager. This is the
    backward direction: the domain is the OUTPUT of the splined functions. In
    this view the knot layout is (possibly highly) irregular. The edge of the
    imager is a perfect rectangle.

Separate from the domain, the data can be presented in 3 different ways:

- Magnitude heatmap. This is the default. We plot mag(deltauxy). This displays
  the deviation from the baseline model as a heat map.

- Individual heatmap. Selected by passing an "xy" argument. We plot deltaux or
  deltauy, depending on the value of xy. This displays the value of one of the
  two splined surfaces individually, as a heat map.

- Vector field. Selected by --vectorfield. Displays the correction (deltaux,
  deltauy) as a vector field.

The splined surfaces are defined by control points we call "knots". These knots
are arranged in a fixed grid (defined by the model configuration) with the value
at each knot set in the intrinsics vector.

The configuration selects the control point density and the expected field of
view of the lens. If the fov_x_deg configuration value is too big, many of the
knots will lie well outside the visible area, and will not be used. This is
wasteful. If fov_x_deg is too small, then some parts of the imager will lie
outside of the spline-in-bounds region, resulting in less-flexible projection
behavior at the edges of the imager. So the field of view should roughly match
the actual lens+camera we're using, and we can evaluate that with this tool.
This tool displays the spline-in-bounds region together with the usable
projection region (either the valid-intrinsics region or the imager bounds).
Ideally, the spline-in-bounds region is slightly bigger than the usable
projection region.

The usable projection region visualized by this tool is controlled by
--show-imager-bounds. If omitted, we display the valid-intrinsics region. This
is recommended, but keep in mind that this region is smaller than the full
imager, so a fov_x_deg that aligns well for one calibration may be too small in
a subsequent calibration of the same lens. If the subsequent calibration has
better coverage, and thus a bigger valid-intrinsics region. If
--show-imager-bounds: we use the imager bounds instead. The issue here is that
the projection near the edges of the imager is usually poorly-defined because
usually there isn't a lot of calibration data there. This makes the projection
behavior at the imager edges unknowable. Consequently, plotting the projection
at the imager edges is usually too alarming or not alarming enough. Passing
--show-imager-bounds is thus recommended only if we have very good calibration
coverage at the edge of the imager.

'''


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='''The density of the plotted grid. By default we
                        use a 60x40 grid''')

    parser.add_argument('--vectorfield',
                        action = 'store_true',
                        help='''Display the spline correction as a vector field.
                        if --vectorfield: the 'xy' argument MUST be omitted''')

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

    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('--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('--imager-domain',
                        action='store_true',
                        help='''By default, this produces a visualization in the domain of the spline-index
                        (normalized stereographic coordinates). Sometimes it's
                        more informative to look at the imager domain instead,
                        by passing this option''')
    parser.add_argument('--show-imager-bounds',
                        action='store_true',
                        help='''By default we communicate the usable projection
                        region to the user by displaying the valid-intrinsics
                        region. This isn't available in all models. To fall back
                        on the boundary of the full imager, pass
                        --show-imager-bounds. In the usual case of incomplete
                        calibration-time coverage at the edges, this results in
                        a very unrealistic representation of reality. Leaving
                        this at the default is recommended''')
    parser.add_argument('--observations',
                        action='store_true',
                        default=False,
                        help='''If given, I show where the chessboard corners were observed at calibration
                        time. This is useful to evaluate the reported unprojectable regions.''')
    parser.add_argument('model',
                        type=str,
                        help='''Input camera model. If "-' is given, we read standard input''')

    parser.add_argument('xy',
                        choices = ('x','y'),
                        nargs   = '?',
                        help='''Optional 'x' or 'y': which surface we're looking
                        at. MUST be omitted if --vectorfield. If omitted and not
                        --vectorfield: we plot the magnitude of the
                        (deltaux,deltauy) corretion vector''')

    args = parser.parse_args()

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

    return args

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




import numpy as np
import numpysane as nps

import mrcal



try:
    model = mrcal.cameramodel(args.model)
except Exception as e:
    print(f"Couldn't load camera model '{args.model}': {e}", file=sys.stderr)
    sys.exit(1)
lensmodel = model.intrinsics()[0]

if not re.match('LENSMODEL_SPLINED', lensmodel):
    print(f"This only makes sense with splined models. Input uses {lensmodel}",
          file = sys.stderr)
    sys.exit(1)

if not args.show_imager_bounds and \
   model.valid_intrinsics_region() is None:
    print("The given model has no valid-intrinsics region. Pass --show-imager-bounds",
          file = sys.stderr)
    sys.exit(1)

plotkwargs = {}
if args.set is not None:
    plotkwargs['set'] = args.set
if args.unset is not None:
    plotkwargs['unset'] = args.unset

if args.title is not None:
    plotkwargs['title'] = args.title
if args.extratitle is not None:
    plotkwargs['extratitle'] = args.extratitle

try:
    plot = mrcal.show_splined_model_correction( model,
                                                vectorfield             = args.vectorfield,
                                                xy                      = args.xy,
                                                imager_domain           = args.imager_domain,
                                                gridn_width             = args.gridn[0],
                                                gridn_height            = args.gridn[1],
                                                vectorscale             = args.vectorscale,
                                                valid_intrinsics_region = not args.show_imager_bounds,
                                                observations            = args.observations,
                                                hardcopy                = args.hardcopy,
                                                terminal                = args.terminal,
                                                **plotkwargs)
except Exception as e:
    print(f"Error: {e}", file=sys.stderr)
    sys.exit(1)

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