# Copyright 2020 DeepMind Technologies Limited.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Define some common object shortcuts for RGB objects.

We can configure generative objects somewhat freely in parameters, but in
current experiments, we use a discrete set of nicknamed objects. These objects
are defined with a per-object constant set of parameters. For easier use, these
are specified here.

The initial version of RGB-objects (named RGB30) was created manually and is not
a part of the current parametric object pipeline.

There is a visualization of the objects and more information can be found here:
https://sites.google.com/corp/google.com/rgb--stacking#h.p_Hbvm_ijsde_K
"""
import collections
import copy
import itertools

from typing import Dict, Tuple
from dm_robotics.manipulation.props.parametric_object import parametric_object
from dm_robotics.manipulation.props.parametric_object.rgb_objects import parametric_rgb_object


# RGB-objects v1.0 are created with 3 deformations of a seed object (a cube with
# a 50mm side): G minor deformation, B average deformation, R major
# deformation. Deformations are chosen by sampling independently 5 parameters of
# the RGB-shapes. We have chosen: 1 - hollowness; 2 - number of sides;
# 3 - shrinking; 4 - not used; 5 - shear; 6 - form factor.

ParametersDict = parametric_object.ParametersDict
RgbVersion = parametric_rgb_object.RgbVersion

# Breaking these over > 100 lines does not help visibility so:
# pylint: disable=line-too-long, bad-whitespace
# pyformat: disable

_OBJECTS_V1_0 = collections.OrderedDict({
    "g1": ParametersDict({ "sds":  4, "shr":  0, "drf": 0, "hlw": 15, "shx": 0, "shy":  0, "scx": 50, "scy": 50, "scz":  50 }),
    "b1": ParametersDict({ "sds":  4, "shr":  0, "drf": 0, "hlw": 20, "shx": 0, "shy":  0, "scx": 50, "scy": 50, "scz":  50 }),
    "r1": ParametersDict({ "sds":  4, "shr":  0, "drf": 0, "hlw": 35, "shx": 0, "shy":  0, "scx": 50, "scy": 50, "scz":  50 }),
    "g2": ParametersDict({ "sds":  6, "shr":  0, "drf": 0, "hlw":  0, "shx": 0, "shy":  0, "scx": 46, "scy": 46, "scz":  50 }),
    "b2": ParametersDict({ "sds":  8, "shr":  0, "drf": 0, "hlw":  0, "shx": 0, "shy":  0, "scx": 45, "scy": 45, "scz":  50 }),
    "r2": ParametersDict({ "sds": 10, "shr":  0, "drf": 0, "hlw":  0, "shx": 0, "shy":  0, "scx": 45, "scy": 45, "scz":  50 }),
    "g3": ParametersDict({ "sds":  4, "shr": 25, "drf": 0, "hlw":  0, "shx": 0, "shy":  0, "scx": 51, "scy": 51, "scz":  60 }),
    "b3": ParametersDict({ "sds":  4, "shr": 48, "drf": 0, "hlw":  0, "shx": 0, "shy":  0, "scx": 46, "scy": 49, "scz":  63 }),
    "r3": ParametersDict({ "sds":  4, "shr": 75, "drf": 0, "hlw":  0, "shx": 0, "shy":  0, "scx": 41, "scy": 49, "scz":  71 }),
    "g5": ParametersDict({ "sds":  4, "shr":  0, "drf": 0, "hlw":  0, "shx": 0, "shy": 20, "scx": 50, "scy": 50, "scz":  50 }),
    "b5": ParametersDict({ "sds":  4, "shr":  0, "drf": 0, "hlw":  0, "shx": 0, "shy": 31, "scx": 50, "scy": 50, "scz":  50 }),
    "r5": ParametersDict({ "sds":  4, "shr":  0, "drf": 0, "hlw":  0, "shx": 0, "shy": 42, "scx": 50, "scy": 50, "scz":  50 }),
    "g6": ParametersDict({ "sds":  4, "shr":  0, "drf": 0, "hlw":  0, "shx": 0, "shy":  0, "scx": 40, "scy": 56, "scz":  80 }),
    "b6": ParametersDict({ "sds":  4, "shr":  0, "drf": 0, "hlw":  0, "shx": 0, "shy":  0, "scx": 32, "scy": 48, "scz":  96 }),
    "r6": ParametersDict({ "sds":  4, "shr":  0, "drf": 0, "hlw":  0, "shx": 0, "shy":  0, "scx": 29, "scy": 29, "scz": 150 })})
# pylint: enable=line-too-long, bad-whitespace
# pyformat: enable


# RGB-objects v1.3 adds two deformation axes to v1.1 (axis 7 and axis 8). The
# axis 6 in v1.1 is the deformation of the form factor by increasing the size
# along the z-axis. With v1.3 we introduce deformations which correspond to
# scaling along the x-axis (axis 7) and along the y-axis (axis 8). Moreover,
# we generate all objects as interpolations from a seed object s and 7 objects
# r2, r3, r5, r6, r7, r8 which are the maximum deformations of the seed
# object along the available axes 2, 3, 5, 6, 7, 8. Finally also, v1.3
# defines the objects Omn with Omn = (Om + On)//2.

# pylint: disable=line-too-long, bad-whitespace
# pyformat: disable
_OBJECTS_V1_3 = collections.OrderedDict({
    "s":  ParametersDict({ "sds": 4, "shr": 0, "drf": 0, "hlw": 0, "shx": 0, "shy": 0, "scx": 50, "scy": 50, "scz": 50}),
    "r2": ParametersDict({ "sds": 10, "shr": 0, "drf": 0, "hlw": 0, "shx": 0, "shy": 0, "scx": 45, "scy": 45, "scz": 50}),
    "r3": ParametersDict({ "sds": 4, "shr": 75, "drf": 0, "hlw": 0, "shx": 0, "shy": 0, "scx": 41, "scy": 49, "scz": 71}),
    "r5": ParametersDict({ "sds": 4, "shr": 0, "drf": 0, "hlw": 0, "shx": 0, "shy": 42, "scx": 50, "scy": 50, "scz": 50}),
    "r6": ParametersDict({ "sds": 4, "shr": 0, "drf": 0, "hlw": 0, "shx": 0, "shy": 0, "scx": 29, "scy": 29, "scz": 150}),
    "r7": ParametersDict({ "sds": 4, "shr": 0, "drf": 0, "hlw": 0, "shx": 0, "shy": 0, "scx": 29, "scy": 150, "scz": 29}),
    "r8": ParametersDict({ "sds": 4, "shr": 0, "drf": 0, "hlw": 0, "shx": 0, "shy": 0, "scx": 150, "scy": 29, "scz": 29})})
_OBJECTS_V1_3_NON_GRASPABLE = ("l36", "m36", "y36", "r36", "r67", "v36")
_OBJECTS_V1_3_NON_UNIQUE = ("f25", "h27", "l27", "m27", "r27", "r68", "u27", "v27", "x27", "y27", "h28", "l28", "m28", "r28", "u28", "v28", "x28", "y28", "r78")

# pylint: enable=line-too-long, bad-whitespace
# pyformat: enable


def parameters_interpolations(
    params_dict1: ParametersDict,
    params_dict2: ParametersDict,
    interpolation_length: int = 1,
    interpolation_keys: Tuple[str, ...] = ()) -> collections.OrderedDict:
  """Function to interpolate in between two parametersDict.

  This function can be used to interpolate in between parametersDicts. The
  function takes as input two parameterDicts. The function interpolates in
  between the parametersDicts generating a given number of equally-spaced
  samples. By default, only one sample is added, corresponding to an element
  in between the two provided parameterDicts (e.g. m = (s+e)/2). Generated
  parameterDicts are returned in a collection. By default, associated labels are
  m1, m2, ... or otherwise specified by the user through a tuple.

  Args:
    params_dict1: the first parameterDict.
    params_dict2: the second parameterDict.
    interpolation_length: the numnber of interpolation samples.
    interpolation_keys: the keys used in the resulting collection.
  Returns:
    the collection of combinations
  """