# coding=utf-8
# Copyright 2019 The Hal Authors.
#
# 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.

"""Utilities that mutate a given sentences."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import copy
import re
import random

import numpy as np

relation_synonyms = {'on the left side': ['left', 'on the left'],
                     'on the right side': ['right', 'on the right'],
                     'in front of': ['front of']}
material_synonyms = {'matte': ['rubber', ''], 'rubber': ['matte', ''], 'shiny': ['metallic', ''],
                     'metallic': ['shiny', '']}
object_synonyms = {'object': ['sphere', 'object', 'thing'],
                   'sphere': ['object', 'ball', 'thing'],
                   'ball': ['sphere', 'object', 'thing'],
                   'objects': ['spheres', 'objects', 'things'],
                   'spheres': ['objects', 'balls', 'things'],
                   'balls': ['spheres', 'objects', 'things']}
adjective_synonyms = {'any': ['']}
misc_synonyms = {'are': ['is']}

synonym_tables = [relation_synonyms, material_synonyms, object_synonyms,
                  adjective_synonyms, misc_synonyms]

colors = [{'red': ['']}, {'blue': ['']}, {'cyan': ['']}, {'purple': ['']}, {'green': ['']}]

other_colors = {
    'red': ['blue', 'cyan', 'purple', 'green'],
    'blue': ['red', 'cyan', 'purple', 'green'],
    'cyan': ['blue', 'red', 'purple', 'green'],
    'purple': ['blue', 'cyan', 'red', 'green'],
    'green': ['blue', 'cyan', 'purple', 'red'],
}

other_directions = {
    'left': ['right'],
    'right': ['left'],
    'front': ['behind'],
    'behind': ['front'],
}


def mutate_sentence(text, k=2, delete_color=False):
  tables = random.sample(synonym_tables, k)
  if delete_color and random.uniform(0, 1) < 0.5:
    tables = random.sample(colors, 5)
    subed = False
    for t in tables:
      if subed: break
      for w in t:
        if w in text:
          text = re.sub(w, random.choice(t[w]), text)
          subed = True
  else:
    for t in tables:
      for w in t:
        if w in text:
          text = re.sub(w, random.choice(t[w]), text)
  return text


def mutate_sentence_and_program(text, program):
  copied_program = copy.deepcopy(program)
  delete_options = [(0, 0), (0, 1), (1, 0), (1, 1)]
  return None


def negate_unary_sentence(text):
  words = text.split(' ')
  mutate_candiate = {}
  for i, w in enumerate(words):
    if w in other_colors:
      mutate_candiate['color'] = (i, w)
    elif w in other_directions:
      mutate_candiate['direction'] = (i, w)
  toss = random.random()
  if toss < 0.33 and 'color' in mutate_candiate:
    i, color = mutate_candiate['color']
    new_color = random.choice(other_colors[color])
    words[i] = new_color
  elif 0.33 < random.random() < 0.66 and 'direction' in mutate_candiate:
    i, direction = mutate_candiate['direction']
    new_direction = random.choice(other_directions[direction])
    words[i] = new_direction
  elif 'direction' in mutate_candiate and 'color' in mutate_candiate:
    i, color = mutate_candiate['color']
    new_color = random.choice(other_colors[color])
    words[i] = new_color
    i, direction = mutate_candiate['direction']
    new_direction = random.choice(other_directions[direction])
    words[i] = new_direction
  else:
    return None

  mutated_text = ' '.join(words)
  return mutated_text
