#!/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 behavior or a lens model

SYNOPSIS

  $ mrcal-show-distortion-off-pinhole --vectorfield left.cameramodel

  ... a plot pops up showing the vector field of the difference from a pinhole
  projection

This tool is used to examine how a lens model behaves. Depending on the model,
the vectors could be very large or very small, and we can scale them by passing
'--vectorscale s'. By default we sample in a 60x40 grid, but this spacing can be
controlled by passing '--gridn w h'.

By default we render a heat map of the lens effects. We can also see the
vectorfield by passing in --vectorfield. Or we can see the radial distortion
curve by passing --radial

This tool treats a pinhole projection as a baseline, and visualizes deviations
from this baseline. So wide lenses will have a lot of reported "distortion".
Thus pinhole models aren't a good baseline, but pretty much every non-mrcal tool
in common use treats lenses as pinhole+distortion, and that makes this
visualization useful.

The plots produced with --radial show the magnitude of the radial distortion as
a function of the distance to the center. For opencv models the radial
distortion is plotted directly from the projection equations: these represent
radial distortion explicitly. For other models, we empirically sample the imager
from the optical center given by intrinsics[2:4] to the 4 corners and the 4
centers of the imager edges. Note that the splined models often shift the
projection center from intrinsics[2:4], causing unnatural-looking curves near
the center. This isn't a a bug: intrinsics[2:4] isn't really in the center,
which isn't a problem for anything except this visualization.

'''


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 report a 60x40 grid''')

    parser.add_argument('--vectorscale',
                        type=float,
                        default = 1.0,
                        help='''Scale the vectors by this factor. Default is 1.0 (report the truth), but this is often too small to see''')

    parser.add_argument('--radial',
                        action='store_true',
                        help='''Show the radial distortion scale factor instead of a colormap/vectorfield''')
    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('--show-fisheye-projections',
                        action='store_true',
                        help='''Applies to radial plots only. if given, the
                        radial plots include the behavior of common fisheye
                        projections, in addition to the behavior of THIS lens''')
    parser.add_argument('--cbmax',
                        type=float,
                        default=10,
                        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('--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 pass to gnuplotlib. May be given multiple
                        times''')
    parser.add_argument('--unset',
                        type=str,
                        action='append',
                        help='''Extra 'unset' directives to pass to gnuplotlib. May be given multiple
                        times''')

    parser.add_argument('model',
                        type=str,
                        help='''Input camera model. If "-' is given, we read standard input''')

    return parser.parse_args()

args = 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)

# 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)

if args.radial and args.vectorfield:
    sys.stderr.write("Usage error: at most one of --radial and --vectorfield can be given\n")
    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

if   args.radial:
    plot = \
        mrcal.show_distortion_off_pinhole_radial(model,
                                                 show_fisheye_projections = args.show_fisheye_projections,
                                                 hardcopy     = args.hardcopy,
                                                 terminal     = args.terminal,
                                                 **plotkwargs)
else:
    plot = mrcal.show_distortion_off_pinhole(model,
                                             vectorfield  = args.vectorfield,
                                             vectorscale  = args.vectorscale,
                                             cbmax        = args.cbmax,
                                             gridn_width  = args.gridn[0],
                                             gridn_height = args.gridn[1],

                                             hardcopy     = args.hardcopy,
                                             terminal     = args.terminal,
                                             **plotkwargs)

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