"""Experimentally determines a camera's rolling shutter skew.

See the accompanying PDF for instructions on how to use this test.
"""
from __future__ import division
from __future__ import print_function

import argparse
import glob
import math
import os
import sys
import tempfile

import cv2
import its.caps
import its.device
import its.image
import its.objects
import numpy as np

DEBUG = False

# Constants for which direction the camera is facing.
FACING_FRONT = 0
FACING_BACK = 1
FACING_EXTERNAL = 2

# Camera capture defaults.
FPS = 30
WIDTH = 640
HEIGHT = 480
TEST_LENGTH = 1

# Each circle in a cluster must be within this many pixels of some other circle
# in the cluster.
CLUSTER_DISTANCE = 50.0 / HEIGHT
# A cluster must consist of at least this percentage of the total contours for
# it to be allowed into the computation.
MAJORITY_THRESHOLD = 0.7

# Constants to make sure the slope of the fitted line is reasonable.
SLOPE_MIN_THRESHOLD = 0.5
SLOPE_MAX_THRESHOLD = 1.5

# To improve readability of unit conversions.
SEC_TO_NSEC = float(10**9)
MSEC_TO_NSEC = float(10**6)
NSEC_TO_MSEC = 1.0 / float(10**6)


class RollingShutterArgumentParser(object):
  """Parses command line arguments for the rolling shutter test."""

  def __init__(self):
    self.__parser = argparse.ArgumentParser(
        description='Run rolling shutter test')
    self.__parser.add_argument(
        '-d', '--debug',
        action='store_true',
        help='print and write data useful for debugging')
    self.__parser.add_argument(
        '-f', '--fps',
        type=int,
        help='FPS to capture with during the test (defaults to 30)')
    self.__parser.add_argument(
        '-i', '--img_size',
        help=('comma-separated dimensions of captured images (defaults '
              'to 640x480). Example: --img_size=<width>,<height>'))
    self.__parser.add_argument(
        '-l', '--led_time',
        type=float,
        required=True,
        help=('how many milliseconds each column of the LED array is '
              'lit for'))
    self.__parser.add_argument(
        '-p', '--panel_distance',
        type=float,
        help='how far the LED panel is from the camera (in meters)')
    self.__parser.add_argument(
        '-r', '--read_dir',
        help=('read existing test data from specified directory.  If '
              'not specified, new test data is collected from the '
              'device\'s camera)'))
    self.__parser.add_argument(
        '--device_id',
        help=('device ID for device being tested (can also use '
              '\'device=<DEVICE ID>\')'))
    self.__parser.add_argument(
        '-t', '--test_length',
        type=int,
        help=('how many seconds the test should run for (defaults to 1 '
              'second)'))
    self.__parser.add_argument(
        '-o', '--debug_dir',
        help=('write debugging information in a folder in the '
              'specified directory.  Otherwise, the system\'s default '
              'location for temporary folders is used.  --debug must '
              'be specified along with this argument.'))

  def parse_args(self):
    """Returns object containing parsed values from the command line."""
    # Don't show argparse the 'device' flag, since it's in a different
    # format than the others (to maintain CameraITS conventions) and it will
    # complain.
    filtered_args = [arg for arg in sys.argv[1:] if 'device=' not in arg]
    args = self.__parser.parse_args(filtered_args)
    if args.device_id:
      # If argparse format is used, convert it to a format its.device can
      # use later on.
      sys.argv.append('device=%s' % args.device_id)
    return args


def main():
  global DEBUG
  global CLUSTER_DISTANCE

  parser = RollingShutterArgumentParser()
  args = parser.parse_args()

  DEBUG = args.debug
  if not DEBUG and args.debug_dir:
    print('argument --debug_dir requires --debug', file=sys.stderr)
    sys.exit()

  if args.read_dir is None:
    # Collect new data.
    raw_caps, reported_skew = collect_data(args)
    frames = [its.image.convert_capture_to_rgb_image(c) for c in raw_caps]
  else:
    # Load existing data.
    frames, reported_skew = load_data(args.read_dir)

  # Make the cluster distance relative to the height of the image.
  (frame_h, _, _) = frames[0].shape
  CLUSTER_DISTANCE = frame_h * CLUSTER_DISTANCE
  debug_print('Setting cluster distance to %spx.' % CLUSTER_DISTANCE)

  if DEBUG:
    debug_dir = setup_debug_dir(args.debug_dir)
    # Write raw frames.
    for i, img in enumerate(frames):
      its.image.write_image(img, '%s/raw/%03d.png' % (debug_dir, i))
  else:
    debug_dir = None

  avg_shutter_skew, num_frames_used = find_average_shutter_skew(
      frames, args.led_time, debug_dir)
  if debug_dir:
    # Write the reported skew with the raw images, so the directory can also
    # be used to read from.
    with open(debug_dir + '/raw/reported_skew.txt', 'w') as f:
      f.write('%sms\n' % reported_skew)

  if avg_shutter_skew is None:
    print('Could not find usable frames.')
  else:
    print('Device reported shutter skew of %sms.' % reported_skew)
    print('Measured shutter skew is %sms (averaged over %s frames).' %
          (avg_shutter_skew, num_frames_used))


def collect_data(args):
  """Capture a new set of frames from the device's camera.

  Args:
      args: Parsed command line arguments.

  Returns:
      A list of RGB images as numpy arrays.
  """
  fps = args.fps if args.fps else FPS
  if args.img_size:
    w, h = map(int, args.img_size.split(','))
  else:
    w, h = WIDTH, HEIGHT
  test_length = args.test_length if args.test_length else TEST_LENGTH

  with its.device.ItsSession() as cam:
    props = cam.get_camera_properties()
    its.caps.skip_unless(its.caps.manual_sensor(props))
    facing = props['android.lens.facing']
    if facing != FACING_FRONT and facing != FACING_BACK:
      print('Unknown lens facing %s' % facing)
      assert 0

    fmt = {'format': 'yuv', 'width': w, 'height': h}
    s, e, _, _, _ = cam.do_3a(get_results=True, do_af=False)
    req = its.objects.manual_capture_request(s, e)
    req['android.control.aeTargetFpsRange'] = [fps, fps]

    # Convert from milliseconds to nanoseconds.  We only want enough
    # exposure time to saturate approximately one column.
    exposure_time = (args.led_time / 2.0) * MSEC_TO_NSEC
    print('Using exposure time of %sns.' % exposure_time)
    req['android.sensor.exposureTime'] = exposure_time
    req['android.sensor.frameDuration'] = int(SEC_TO_NSEC / fps)

    if args.panel_distance is not None:
      # Convert meters to diopters and use that for the focus distance.
      req['android.lens.focusDistance'] = 1 / args.panel_distance
    print('Starting capture')
    raw_caps = cam.do_capture([req]*fps*test_length, fmt)
    print('Finished capture')

    # Convert from nanoseconds to milliseconds.
    shutter_skews = {c['metadata']['android.sensor.rollingShutterSkew'] *
                     NSEC_TO_MSEC for c in raw_caps}
    # All frames should have same rolling shutter skew.
    assert len(shutter_skews) == 1
    shutter_skew = list(shutter_skews)[0]

    return raw_caps, shutter_skew


def load_data(dir_name):
  """Reads camera frame data from an existing directory.

  Args:
    dir_name: Name of the directory to read data from.

  Returns:
    A list of RGB images as numpy arrays.
  """
  frame_files = glob.glob('%s/*.png' % dir_name)
  frames = []
  for frame_file in sorted(frame_files):
    frames.append(its.image.load_rgb_image(frame_file))
  with open('%s/reported_skew.txt' % dir_name, 'r') as f:
    reported_skew = f.readline()[:-2]  # Strip off 'ms' suffix
  return frames, reported_skew


def find_average_shutter_skew(frames, led_time, debug_dir=None):
  """Finds the average shutter skew using the given frames.

  Frames without enough information will be discarded from the average to
  improve overall accuracy.

  Args:
      frames:    List of RGB images from the camera being tested.
      led_time:  How long a single LED column is lit for (in milliseconds).
      debug_dir: (optional) Directory to write debugging information to.

  Returns:
      The average calculated shutter skew and the number of frames used to
      calculate the average.
  """
  avg_shutter_skew = 0.0
  avg_slope = 0.0
  weight = 0.0
  num_frames_used = 0

  for i, frame in enumerate(frames):
    debug_print('------------------------')
    debug_print('| PROCESSING FRAME %03d |' % i)
    debug_print('------------------------')
    shutter_skew, confidence, slope = calculate_shutter_skew(
        frame, led_time, i, debug_dir=debug_dir)
    if shutter_skew is None:
      debug_print('Skipped frame.')
    else:
      debug_print('Shutter skew is %sms (confidence: %s).' %
                  (shutter_skew, confidence))
      # Use the confidence to weight the average.
      avg_shutter_skew += shutter_skew * confidence
      avg_slope += slope * confidence
      weight += confidence
      num_frames_used += 1

  debug_print('\n')
  if num_frames_used == 0:
    return None, None
  else:
    avg_shutter_skew /= weight
    avg_slope /= weight
    slope_err_str = ('The average slope of the fitted line was too %s '
                     'to get an accurate measurement (slope was %s).  '
                     'Try making the LED panel %s.')
    if avg_slope < SLOPE_MIN_THRESHOLD:
      print(slope_err_str % ('flat', avg_slope, 'slower'),
            file=sys.stderr)
    elif avg_slope > SLOPE_MAX_THRESHOLD:
      print(slope_err_str % ('steep', avg_slope, 'faster'),
            file=sys.stderr)
    return avg_shutter_skew, num_frames_used


def calculate_shutter_skew(frame, led_time, frame_num=None, debug_dir=None):
  """Calculates the shutter skew of the camera being used for this test.

  Args:
      frame:     A single RGB image captured by the camera being tested.
      led_time:  How long a single LED column is lit for (in milliseconds).
      frame_num: (optional) Number of the given frame.
      debug_dir: (optional) Directory to write debugging information to.

  Returns:
      The shutter skew (in milliseconds), the confidence in the accuracy of
      the measurement (useful for weighting averages), and the slope of the
      fitted line.
  """
  contours, scratch_img, contour_img, mono_img = find_contours(frame.copy())
  if debug_dir is not None:
    cv2.imwrite('%s/contour/%03d.png' % (debug_dir, frame_num), contour_img)
    cv2.imwrite('%s/mono/%03d.png' % (debug_dir, frame_num), mono_img)

  largest_cluster, cluster_percentage = find_largest_cluster(contours,
                                                             scratch_img)
  if largest_cluster is None:
    debug_print('No majority cluster found.')
    return None, None, None
  elif len(largest_cluster) <= 1:
    debug_print('Majority cluster was too small.')
    return None, None, None
  debug_print('%s points in the largest cluster.' % len(largest_cluster))

  np_cluster = np.array([[c.x, c.y] for c in largest_cluster])
  [vx], [vy], [x0], [y0] = cv2.fitLine(np_cluster, cv2.cv.CV_DIST_L2,
                                       0, 0.01, 0.01)
  slope = vy / vx
  debug_print('Slope is %s.' % slope)
  (frame_h, frame_w, _) = frame.shape
  # Draw line onto scratch frame.
  pt1 = tuple(map(int, (x0 - vx * 1000, y0 - vy * 1000)))
  pt2 = tuple(map(int, (x0 + vx * 1000, y0 + vy * 1000)))
  cv2.line(scratch_img, pt1, pt2, (0, 255, 255), thickness=3)

  # We only need the width of the cluster.
  _, _, cluster_w, _ = find_cluster_bounding_rect(largest_cluster,
                                                  scratch_img)

  num_columns = find_num_columns_spanned(largest_cluster)
  debug_print('%s columns spanned by cluster.' % num_columns)
  # How long it takes for a column to move from the left of the bounding
  # rectangle to the right.
  left_to_right_time = led_time * num_columns
  milliseconds_per_x_pixel = left_to_right_time / cluster_w
  # The distance between the line's intersection at the top of the frame and
  # the intersection at the bottom.
  x_range = frame_h / slope
  shutter_skew = milliseconds_per_x_pixel * x_range
  # If the aspect ratio is different from 4:3 (the aspect ratio of the actual
  # sensor), we need to correct, because it will be cropped.
  shutter_skew *= (float(frame_w) / float(frame_h)) / (4.0 / 3.0)

  if debug_dir is not None:
    cv2.imwrite('%s/scratch/%03d.png' % (debug_dir, frame_num),
                scratch_img)

  return shutter_skew, cluster_percentage, slope


def find_contours(img):
  """Finds contours in the given image.

  Args:
    img: Image in Android camera RGB format.

  Returns:
    OpenCV-formatted contours, the original image in OpenCV format, a
    thresholded image with the contours drawn on, and a grayscale version of
    the image.
  """
  # Convert to format OpenCV can work with (BGR ordering with byte-ranged
  # values).
  img *= 255
  img = img.astype(np.uint8)
  img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)

  # Since the LED colors for the panel we're using are red, we can get better
  # contours for the LEDs if we ignore the green and blue channels.  This also
  # makes it so we don't pick up the blue control screen of the LED panel.
  red_img = img[:, :, 2]
  _, thresh = cv2.threshold(red_img, 0, 255, cv2.THRESH_BINARY +
                            cv2.THRESH_OTSU)

  # Remove noise before finding contours by eroding the thresholded image and
  # then re-dilating it.  The size of the kernel represents how many
  # neighboring pixels to consider for the result of a single pixel.
  kernel = np.ones((3, 3), np.uint8)
  opening = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel, iterations=2)

  if DEBUG:
    # Need to convert it back to BGR if we want to draw colored contours.
    contour_img = cv2.cvtColor(opening, cv2.COLOR_GRAY2BGR)
  else:
    contour_img = None
  cv2_version = cv2.__version__
  if cv2_version.startswith('3.'):  # OpenCV 3.x
    _, contours, _ = cv2.findContours(
        opening, cv2.cv.CV_RETR_EXTERNAL, cv2.cv.CV_CHAIN_APPROX_NONE)
  else:  # OpenCV 2.x and 4.x
    contours, _ = cv2.findContours(
        opening, cv2.cv.CV_RETR_EXTERNAL, cv2.cv.CV_CHAIN_APPROX_NONE)
  if DEBUG:
    cv2.drawContours(contour_img, contours, -1, (0, 0, 255), thickness=2)
  return contours, img, contour_img, red_img


def convert_to_circles(contours):
  """Converts given contours into circle objects.

  Args:
    contours: Contours generated by OpenCV.

  Returns:
    A list of circles.
  """

  class Circle(object):
    """Holds data to uniquely define a circle."""

    def __init__(self, contour):
      self.x = int(np.mean(contour[:, 0, 0]))
      self.y = int(np.mean(contour[:, 0, 1]))
      # Get diameters of each axis then half it.
      x_r = (np.max(contour[:, 0, 0]) - np.min(contour[:, 0, 0])) / 2.0
      y_r = (np.max(contour[:, 0, 1]) - np.min(contour[:, 0, 1])) / 2.0
      # Average x radius and y radius to get the approximate radius for
      # the given contour.
      self.r = (x_r + y_r) / 2.0
      assert self.r > 0.0

    def distance_to(self, other):
      return (math.sqrt((other.x - self.x)**2 + (other.y - self.y)**2) -
              self.r - other.r)

    def intersects(self, other):
      return self.distance_to(other) <= 0.0

  return list(map(Circle, contours))


def find_largest_cluster(contours, frame):
  """Finds the largest cluster in the given contours.

  Args:
    contours: Contours generated by OpenCV.
    frame:    For drawing debugging information onto.

  Returns:
    The cluster with the most contours in it and the percentage of all
    contours that the cluster contains.
  """
  clusters = proximity_clusters(contours)

  if not clusters:
    return None, None  # No clusters found.

  largest_cluster = max(clusters, key=len)
  cluster_percentage = len(largest_cluster) / len(contours)

  if cluster_percentage < MAJORITY_THRESHOLD:
    return None, None

  if DEBUG:
    # Draw largest cluster on scratch frame.
    for circle in largest_cluster:
      cv2.circle(frame, (int(circle.x), int(circle.y)), int(circle.r),
                 (0, 255, 0), thickness=2)

  return largest_cluster, cluster_percentage


def proximity_clusters(contours):
  """Sorts the given contours into groups by distance.

  Converts every given contour to a circle and clusters by adding a circle to
  a cluster only if it is close to at least one other circle in the cluster.

  TODO: Make algorithm faster (currently O(n**2)).

  Args:
    contours: Contours generated by OpenCV.

  Returns:
    A list of clusters, where each cluster is a list of the circles
    contained in the cluster.
  """
  circles = convert_to_circles(contours)

  # Use disjoint-set data structure to store assignments.  Start every point
  # in their own cluster.
  cluster_assignments = [-1 for i in range(len(circles))]

  def get_canonical_index(i):
    if cluster_assignments[i] >= 0:
      index = get_canonical_index(cluster_assignments[i])
      # Collapse tree for better runtime.
      cluster_assignments[i] = index
      return index
    else:
      return i

  def get_cluster_size(i):
    return -cluster_assignments[get_canonical_index(i)]

  for i, curr in enumerate(circles):
    close_circles = [j for j, p in enumerate(circles) if i != j and
                     curr.distance_to(p) < CLUSTER_DISTANCE]
    if close_circles:
      # Note: largest_cluster is an index into cluster_assignments.
      largest_cluster = min(close_circles, key=get_cluster_size)
      largest_size = get_cluster_size(largest_cluster)
      curr_index = get_canonical_index(i)
      curr_size = get_cluster_size(i)
      if largest_size > curr_size:
        # largest_cluster is larger than us.
        target_index = get_canonical_index(largest_cluster)
        # Add our cluster size to the bigger one.
        cluster_assignments[target_index] -= curr_size
        # Reroute our group to the bigger one.
        cluster_assignments[curr_index] = target_index
      else:
        # We're the largest (or equal to the largest) cluster.  Reroute
        # all groups to us.
        for j in close_circles:
          smaller_size = get_cluster_size(j)
          smaller_index = get_canonical_index(j)
          if smaller_index != curr_index:
            # We only want to modify clusters that aren't already in
            # the current one.

            # Add the smaller cluster's size to ours.
            cluster_assignments[curr_index] -= smaller_size
            # Reroute their group to us.
            cluster_assignments[smaller_index] = curr_index

  # Convert assignments list into list of clusters.
  clusters_dict = {}
  for i in range(len(cluster_assignments)):
    canonical_index = get_canonical_index(i)
    if canonical_index not in clusters_dict:
      clusters_dict[canonical_index] = []
    clusters_dict[canonical_index].append(circles[i])
  return clusters_dict.values()


def find_cluster_bounding_rect(cluster, scratch_frame):
  """Finds the minimum rectangle that bounds the given cluster.

  The bounding rectangle will always be axis-aligned.

  Args:
    cluster:       Cluster being used to find the bounding rectangle.
    scratch_frame: Image that rectangle is drawn onto for debugging
                   purposes.

  Returns:
    The leftmost and topmost x and y coordinates, respectively, along with
    the width and height of the rectangle.
  """
  avg_distance = find_average_neighbor_distance(cluster)
  debug_print('Average distance between points in largest cluster is %s '
              'pixels.' % avg_distance)

  c_x = min(cluster, key=lambda c: c.x - c.r)
  c_y = min(cluster, key=lambda c: c.y - c.r)
  c_w = max(cluster, key=lambda c: c.x + c.r)
  c_h = max(cluster, key=lambda c: c.y + c.r)

  x = c_x.x - c_x.r - avg_distance
  y = c_y.y - c_y.r - avg_distance
  w = (c_w.x + c_w.r + avg_distance) - x
  h = (c_h.y + c_h.r + avg_distance) - y

  if DEBUG:
    points = np.array([[x, y], [x + w, y], [x + w, y + h], [x, y + h]],
                      np.int32)
    cv2.polylines(scratch_frame, [points], True, (255, 0, 0), thickness=2)

  return x, y, w, h


def find_average_neighbor_distance(cluster):
  """Finds the average distance between every circle and its closest neighbor.

  Args:
    cluster: List of circles

  Returns:
    The average distance.
  """
  avg_distance = 0.0
  for a in cluster:
    closest_point = None
    closest_dist = None
    for b in cluster:
      if a is b:
        continue
      curr_dist = a.distance_to(b)
      if closest_point is None or curr_dist < closest_dist:
        closest_point = b
        closest_dist = curr_dist
    avg_distance += closest_dist
  avg_distance /= len(cluster)
  return avg_distance


def find_num_columns_spanned(circles):
  """Finds how many columns of the LED panel are spanned by the given circles.

  Args:
    circles: List of circles (assumed to be from the LED panel).

  Returns:
    The number of columns spanned.
  """
  if not circles:
    return 0

  def x_intersects(c_a, c_b):
    return abs(c_a.x - c_b.x) < (c_a.r + c_b.r)

  circles = sorted(circles, key=lambda c: c.x)
  last_circle = circles[0]
  num_columns = 1
  for circle in circles[1:]:
    if not x_intersects(circle, last_circle):
      last_circle = circle
      num_columns += 1

  return num_columns


def setup_debug_dir(dir_name=None):
  """Creates a debug directory and required subdirectories.

  Each subdirectory contains images from a different step in the process.

  Args:
    dir_name: The directory to create.  If none is specified, a temp
    directory is created.

  Returns:
    The name of the directory that is used.
  """
  if dir_name is None:
    dir_name = tempfile.mkdtemp()
  else:
    force_mkdir(dir_name)
  print('Saving debugging files to "%s"' % dir_name)
  # For original captured images.
  force_mkdir(dir_name + '/raw', clean=True)
  # For monochrome images.
  force_mkdir(dir_name + '/mono', clean=True)
  # For contours generated from monochrome images.
  force_mkdir(dir_name + '/contour', clean=True)
  # For post-contour debugging information.
  force_mkdir(dir_name + '/scratch', clean=True)
  return dir_name


def force_mkdir(dir_name, clean=False):
  """Creates a directory if it doesn't already exist.

  Args:
    dir_name: Name of the directory to create.
    clean:    (optional) If set to true, cleans image files from the
              directory (if it already exists).
  """
  if os.path.exists(dir_name):
    if clean:
      for image in glob.glob('%s/*.png' % dir_name):
        os.remove(image)
  else:
    os.makedirs(dir_name)


def debug_print(s, *args, **kwargs):
  """Only prints if the test is running in debug mode."""
  if DEBUG:
    print(s, *args, **kwargs)


if __name__ == '__main__':
  main()