# coding=utf-8
# Copyright 2022 The Google Research 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.

# Lint as: python3
"""NTS-Net adapted for perturbed top-k.

Based on the original PyTorch code
https://github.com/yangze0930/NTS-Net/blob/master/core/model.py
"""

import enum
import functools
import math
from typing import List, Tuple

from absl import app
from absl import flags
from absl import logging
import chex
from clu import platform
import einops
from flax.deprecated import nn
import jax
import jax.numpy as jnp
import ml_collections
import ml_collections.config_flags as config_flags
from off_the_grid.lib import data
from off_the_grid.lib import models
from off_the_grid.lib import utils
import off_the_grid.lib.classification_utils as classification_lib
from off_the_grid.lib.layers import sample_patches
from off_the_grid.lib.layers import transformer
import optax
import tensorflow as tf


FLAGS = flags.FLAGS

config_flags.DEFINE_config_file(
    "config", None, "Training configuration.", lock_config=True)
flags.DEFINE_string("workdir", None, "Work unit directory.")
NUM_CLASSES = 200

ANCHORS_SETTINGS = (
    dict(
        layer="p3",
        stride=32,
        size=48,
        scale=[2**(1. / 3.), 2**(2. / 3.)],
        aspect_ratio=[0.667, 1, 1.5]),  # Anchors 0-5
    dict(
        layer="p4",
        stride=64,
        size=96,
        scale=[2**(1. / 3.), 2**(2. / 3.)],
        aspect_ratio=[0.667, 1, 1.5]),  # Anchors 6-11
    dict(
        layer="p5",
        stride=128,
        size=192,
        scale=[1, 2**(1. / 3.), 2**(2. / 3.)],
        aspect_ratio=[0.667, 1, 1.5]),  # Anchors 12-20
)


class Communication(str, enum.Enum):
  NONE = "none"
  SQUEEZE_EXCITE_D = "squeeze_excite_d"
  SQUEEZE_EXCITE_X = "squeeze_excite_x"
  TRANSFORMER = "transformer"


def zeroone(scores, x_min, x_max):
  """Normalize values to lie between [0, 1]."""
  return [(x - x_min) / (x_max - x_min + 1e-5) for x in scores]


class ProposalNet(nn.Module):
  """FPN inspired scorer module."""

  def apply(self, x,
            communication = Communication.NONE,
            train = True):
    """Forward pass."""
    batch_size = x.shape[0]

    if communication is Communication.SQUEEZE_EXCITE_X:
      x = sample_patches.SqueezeExciteLayer(x)
    # end if squeeze excite x

    d1 = nn.relu(nn.Conv(
        x, 128, kernel_size=(3, 3), strides=(1, 1), bias=True, name="down1"))
    d2 = nn.relu(nn.Conv(
        d1, 128, kernel_size=(3, 3), strides=(2, 2), bias=True, name="down2"))
    d3 = nn.relu(nn.Conv(
        d2, 128, kernel_size=(3, 3), strides=(2, 2), bias=True, name="down3"))

    if communication is Communication.SQUEEZE_EXCITE_D:
      d1_flatten = einops.rearrange(d1, "b h w c -> b (h w) c")
      d2_flatten = einops.rearrange(d2, "b h w c -> b (h w) c")
      d3_flatten = einops.rearrange(d3, "b h w c -> b (h w) c")

      nd1 = d1_flatten.shape[1]
      nd2 = d2_flatten.shape[1]

      d_together = jnp.concatenate([d1_flatten, d2_flatten, d3_flatten], axis=1)

      num_channels = d_together.shape[-1]
      y = d_together.mean(axis=1)
      y = nn.Dense(y, features=num_channels // 4, bias=False)
      y = nn.relu(y)
      y = nn.Dense(y, features=num_channels, bias=False)
      y = nn.sigmoid(y)

      d_together = d_together * y[:, None, :]

      # split and reshape
      d1 = d_together[:, :nd1].reshape(d1.shape)
      d2 = d_together[:, nd1:nd1+nd2].reshape(d2.shape)
      d3 = d_together[:, nd1+nd2:].reshape(d3.shape)

    elif communication is Communication.TRANSFORMER:
      d1_flatten = einops.rearrange(d1, "b h w c -> b (h w) c")
      d2_flatten = einops.rearrange(d2, "b h w c -> b (h w) c")
      d3_flatten = einops.rearrange(d3, "b h w c -> b (h w) c")

      nd1 = d1_flatten.shape[1]
      nd2 = d2_flatten.shape[1]

      d_together = jnp.concatenate([d1_flatten, d2_flatten, d3_flatten], axis=1)

      positional_encodings = self.param(
          "scale_ratio_position_encodings",
          shape=(1,) + d_together.shape[1:],
          initializer=jax.nn.initializers.normal(1. / d_together.shape[-1]))
      d_together = transformer.Transformer(
          d_together + positional_encodings,
          num_layers=2,
          num_heads=8,
          is_training=train)

      # split and reshape
      d1 = d_together[:, :nd1].reshape(d1.shape)
      d2 = d_together[:, nd1:nd1+nd2].reshape(d2.shape)
      d3 = d_together[:, nd1+nd2:].reshape(d3.shape)

    t1 = nn.Conv(
        d1, 6, kernel_size=(1, 1), strides=(1, 1), bias=True, name="tidy1")
    t2 = nn.Conv(
        d2, 6, kernel_size=(1, 1), strides=(1, 1), bias=True, name="tidy2")
    t3 = nn.Conv(
        d3, 9, kernel_size=(1, 1), strides=(1, 1), bias=True, name="tidy3")

    raw_scores = (jnp.split(t1, 6, axis=-1) +
                  jnp.split(t2, 6, axis=-1) +
                  jnp.split(t3, 9, axis=-1))

    # The following is for normalization.
    t = jnp.concatenate((jnp.reshape(t1, [batch_size, -1]),
                         jnp.reshape(t2, [batch_size, -1]),
                         jnp.reshape(t3, [batch_size, -1])), axis=1)
    t_min = jnp.reshape(jnp.min(t, axis=-1), [batch_size, 1, 1, 1])
    t_max = jnp.reshape(jnp.max(t, axis=-1), [batch_size, 1, 1, 1])
    normalized_scores = zeroone(raw_scores, t_min, t_max)

    stats = {
        "scores": normalized_scores,
        "raw_scores": t,
    }
    # removes the split dimension. scores are now b x h' x w' shaped
    normalized_scores = [s.squeeze(-1) for s in normalized_scores]

    return normalized_scores, stats


def extract_weighted_patches(x,
                             weights,
                             kernel,
                             stride,
                             padding):
  """Weighted average of patches using jax.lax.scan."""
  logging.info("recompiling for kernel=%s and stride=%s and padding=%s", kernel,
               stride, padding)
  x = jnp.pad(x, ((0, 0),
                  (padding[0], padding[0] + kernel[0]),
                  (padding[1], padding[1] + kernel[1]),
                  (0, 0)))
  batch_size, _, _, channels = x.shape
  _, k, weights_h, weights_w = weights.shape

  def accumulate_patches(acc, index_i_j):
    i, j = index_i_j
    patch = jax.lax.dynamic_slice(
        x,
        (0, i * stride[0], j * stride[1], 0),
        (batch_size, kernel[0], kernel[1], channels))
    weight = weights[:, :, i, j]

    weighted_patch = jnp.einsum("bk, bijc -> bkijc", weight, patch)
    acc += weighted_patch
    return acc, None

  indices = jnp.stack(
      jnp.meshgrid(jnp.arange(weights_h), jnp.arange(weights_w), indexing="ij"),
      axis=-1)
  indices = indices.reshape((-1, 2))

  init_patches = jnp.zeros((batch_size, k, kernel[0], kernel[1], channels))
  patches, _ = jax.lax.scan(accumulate_patches, init_patches, indices)

  return patches


def weighted_anchor_aggregator(x, weights):
  """Given a tensor of weights per anchor computes the weighted average."""
  counter = 0
  all_sub_aggregates = []

  for anchor_info in ANCHORS_SETTINGS:
    stride = anchor_info["stride"]
    size = anchor_info["size"]
    for scale in anchor_info["scale"]:
      for aspect_ratio in anchor_info["aspect_ratio"]:
        kernel_size = (
            int(size * scale / float(aspect_ratio) ** 0.5),
            int(size * scale * float(aspect_ratio) ** 0.5))
        padding = (
            math.ceil((kernel_size[0] - stride) / 2.),
            math.ceil((kernel_size[1] - stride) / 2.))
        aggregate = extract_weighted_patches(
            x, weights[counter], kernel_size, (stride, stride), padding)
        aggregate = jnp.reshape(aggregate,
                                [-1, kernel_size[0], kernel_size[1], 3])
        aggregate_224 = jax.image.resize(aggregate,
                                         [aggregate.shape[0], 224, 224, 3],
                                         "bilinear")
        all_sub_aggregates.append(aggregate_224)
        counter += 1

  return jnp.sum(jnp.stack(all_sub_aggregates, axis=0), axis=0)


class AttentionNet(nn.Module):
  """The complete NTS-Net model using perturbed top-k."""

  def apply(self,
            x,
            config,
            num_classes,
            train = True):
    """Creates a model definition."""
    b, c = x.shape[0], x.shape[3]
    k = config.k
    sigma = config.ptopk_sigma
    num_samples = config.ptopk_num_samples

    sigma *= self.state("sigma_mutiplier", shape=(),
                        initializer=nn.initializers.ones).value

    stats = {"x": x, "sigma": sigma}

    feature_extractor = models.ResNet50.shared(train=train, name="ResNet_0")

    rpn_feature = feature_extractor(x)