# Copyright 2023 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.
# ==============================================================================

"""GeoMatch++ model definition.

Extends GeoMatch (geomatch.py) with:
  1. A morphology encoder: a third GCN that encodes the robot's kinematic-tree
     graph (links as nodes, joints as edges, 9-dimensional node features).
  2. A DCP-style cross-attention transformer that learns correspondence between
     the object embedding and the morphology embedding.  Its output is a
     residual that is added to the object GCN encoding before the contact-map
     prediction head and the autoregressive modules.

The object GCN (obj_encoder) and robot surface GCN (robot_encoder) are loaded
from a pretrained GeoMatch checkpoint and kept frozen throughout training.
Only the morphology encoder, the transformer, the projection heads, and the
autoregressive modules are updated.

Reference:
    GeoMatch++: Morphology-Aware Grasping via Correspondence Learning
    https://arxiv.org/abs/2412.18998
"""

import torch
from torch import nn

from models.gnn import GCN
from models.geomatch import GeoMatchARModule


# ---------------------------------------------------------------------------
# DCP Transformer
# ---------------------------------------------------------------------------

class _DCPLayer(nn.Module):
  """Single Deep Closest Point cross-attention block.

  Applies self-attention on each sequence, then bidirectional cross-attention.
  All attention operations use pre-norm (LayerNorm before attention) for
  training stability.

  Args:
      embed_dim: Dimensionality of input/output tokens.
      n_heads:   Number of attention heads.  embed_dim must be divisible by n_heads.
      dropout:   Dropout rate applied inside MultiheadAttention.
  """

  def __init__(self, embed_dim: int, n_heads: int, dropout: float = 0.1):
    super().__init__()
    kw = dict(embed_dim=embed_dim, num_heads=n_heads,
              dropout=dropout, batch_first=True)

    self.obj_self_attn   = nn.MultiheadAttention(**kw)
    self.morph_self_attn = nn.MultiheadAttention(**kw)
    self.obj_cross_attn  = nn.MultiheadAttention(**kw)
    self.morph_cross_attn = nn.MultiheadAttention(**kw)

    self.norm_obj_sa    = nn.LayerNorm(embed_dim)
    self.norm_morph_sa  = nn.LayerNorm(embed_dim)
    self.norm_obj_ca    = nn.LayerNorm(embed_dim)
    self.norm_morph_ca  = nn.LayerNorm(embed_dim)

  def forward(
      self,
      obj_embed: torch.Tensor,    # [B, N_obj, D]
      morph_embed: torch.Tensor,  # [B, N_morph, D]
  ):
    # Self-attention on object tokens
    obj_sa, _   = self.obj_self_attn(
        self.norm_obj_sa(obj_embed),
        self.norm_obj_sa(obj_embed),
        self.norm_obj_sa(obj_embed),
    )
    obj_sa = obj_embed + obj_sa             # residual [B, N_obj, D]

    # Self-attention on morphology tokens
    morph_sa, _ = self.morph_self_attn(
        self.norm_morph_sa(morph_embed),
        self.norm_morph_sa(morph_embed),
        self.norm_morph_sa(morph_embed),
    )
    morph_sa = morph_embed + morph_sa       # residual [B, N_morph, D]

    # Cross-attention: object queries morphology
    obj_ca, _   = self.obj_cross_attn(
        self.norm_obj_ca(obj_sa),           # query  [B, N_obj, D]
        self.norm_morph_ca(morph_sa),        # key    [B, N_morph, D]
        self.norm_morph_ca(morph_sa),        # value  [B, N_morph, D]
    )                                        # output [B, N_obj, D]

    # Cross-attention: morphology queries object
    morph_ca, _ = self.morph_cross_attn(
        self.norm_morph_ca(morph_sa),        # query  [B, N_morph, D]
        self.norm_obj_ca(obj_sa),            # key    [B, N_obj, D]
        self.norm_obj_ca(obj_sa),            # value  [B, N_obj, D]
    )                                        # output [B, N_morph, D]

    return obj_ca, morph_ca


class DCPTransformer(nn.Module):
  """Multi-layer DCP cross-attention transformer.

  Stacks n_layers of _DCPLayer with accumulated residuals.

  Args:
      embed_dim: Token dimensionality (must match GCN output dim = 512).
      n_heads:   Attention heads per layer.
      n_layers:  Number of stacked DCP blocks (paper uses 1).
      dropout:   Attention dropout probability.
  """

  def __init__(
      self,
      embed_dim: int = 512,
      n_heads: int = 4,
      n_layers: int = 1,
      dropout: float = 0.1,
  ):
    super().__init__()
    self.layers = nn.ModuleList(
        [_DCPLayer(embed_dim, n_heads, dropout) for _ in range(n_layers)]
    )

  def forward(
      self,
      obj_embed: torch.Tensor,    # [B, N_obj, D]
      morph_embed: torch.Tensor,  # [B, N_morph, D]
  ) -> torch.Tensor:
    """Returns the net residual for the object embedding: [B, N_obj, D].

    The caller adds this residual to the frozen object GCN output:
        obj_embed_final = obj_embed_raw + dcp_transformer(obj_embed_raw, morph_embed)
    """
    obj_out, morph_out = obj_embed, morph_embed
    for layer in self.layers:
      obj_ca, morph_ca = layer(obj_out, morph_out)
      obj_out   = obj_out   + obj_ca
      morph_out = morph_out + morph_ca

    # Net residual relative to the original obj_embed input
    return obj_out - obj_embed   # [B, N_obj, D]


# ---------------------------------------------------------------------------
# GeoMatch++ model
# ---------------------------------------------------------------------------

class GeoMatchPP(nn.Module):
  """GeoMatch++ model.

  Architecture:
      obj_encoder      (frozen) : GCN(3  -> 256x3 -> 512) encodes object PC
      robot_encoder    (frozen) : GCN(3  -> 256x3 -> 512) encodes robot surface PC
      morphology_encoder        : GCN(9  -> 256x3 -> 512) encodes kinematic graph
      dcp_transformer           : DCPTransformer(512, n_heads=4, n_layers=1)
      obj_proj                  : Linear(512 -> 64, no bias)
      robot_proj                : Linear(512 -> 64, no bias)
      kp_ar_model_1..5          : five GeoMatchARModule instances

  Forward pass (see forward() for tensor shapes):
      obj_embed_raw  = L2-norm(obj_encoder(obj_pc, obj_adj))
      robot_embed    = L2-norm(robot_encoder(robot_pc, robot_adj))
      morph_embed    = L2-norm(morphology_encoder(morph_features, morph_adj))
      obj_embed      = obj_embed_raw + DCPTransformer(obj_embed_raw, morph_embed)
      contact_map    = obj_embed @ keypoint_feat.T        [B, 2048, 6, 1]
      (AR modules unchanged from GeoMatch)
  """

  def __init__(self, config) -> None:
    super().__init__()

    self.config            = config
    self.n_kp              = config.keypoint_n
    self.robot_weighting   = config.robot_weighting
    self.match_weighting   = config.matchnet_weighting
    self.dist_loss_weight  = config.dist_loss_weight
    self.match_loss_weight = config.match_loss_weight

    # Frozen encoders (same architecture as GeoMatch — weights loaded separately)
    self.obj_encoder = GCN(
        nfeat=config.obj_in_feats,
        nhid=config.hidden_n,
        nout=config.obj_out_feats,
        dropout=0.5,
        num_hidden=config.num_hidden,
    )
    self.robot_encoder = GCN(
        nfeat=config.robot_in_feats,
        nhid=config.hidden_n,
        nout=config.robot_out_feats,
        dropout=0.5,
        num_hidden=config.num_hidden,
    )

    # Morphology encoder — trained from scratch
    self.morphology_encoder = GCN(
        nfeat=config.morph_in_feats,    # 9
        nhid=config.hidden_n,           # 256
        nout=config.morph_out_feats,    # 512  (must == obj_out_feats)
        dropout=0.5,
        num_hidden=config.num_hidden,   # 3
    )

    # DCP cross-attention transformer — trained from scratch
    self.dcp_transformer = DCPTransformer(
        embed_dim=config.transformer_embed_dim,   # 512
        n_heads=config.transformer_n_heads,        # 4
        n_layers=config.transformer_n_layers,      # 1
    )

    # Projection heads — re-initialised (trainable)
    self.obj_proj   = nn.Linear(config.obj_out_feats,   64, bias=False)
    self.robot_proj = nn.Linear(config.robot_out_feats, 64, bias=False)

    # Autoregressive keypoint modules — re-initialised (trainable)
    self.kp_ar_model_1 = GeoMatchARModule(config, 1)
    self.kp_ar_model_2 = GeoMatchARModule(config, 2)
    self.kp_ar_model_3 = GeoMatchARModule(config, 3)
    self.kp_ar_model_4 = GeoMatchARModule(config, 4)
    self.kp_ar_model_5 = GeoMatchARModule(config, 5)

  # ── Weight loading and freezing ───────────────────────────────────────────

  def load_geomatch_weights(self, pretrained_path: str, device: str = 'cpu'):
    """Load obj_encoder and robot_encoder weights from a GeoMatch checkpoint.

    Only state-dict keys beginning with 'obj_encoder.' or 'robot_encoder.'
    are copied.  All other parameters remain at their random initialisation.
    """
    pretrained = torch.load(
        pretrained_path, map_location=device, weights_only=True
    )
    own_state = self.state_dict()
    loaded, skipped = [], []
    for k, v in pretrained.items():
      if k.startswith(('obj_encoder.', 'robot_encoder.')):
        own_state[k] = v
        loaded.append(k)
      else:
        skipped.append(k)
    self.load_state_dict(own_state)
    print(f'Loaded {len(loaded)} pretrained parameter tensors '
          f'(obj_encoder + robot_encoder).')
    print(f'Skipped {len(skipped)} keys (morphology encoder, transformer, '
          f'projections, AR modules — trained from scratch).')

  def freeze_pretrained_encoders(self):
    """Freeze obj_encoder and robot_encoder.  Call after load_geomatch_weights."""
    for param in self.obj_encoder.parameters():
      param.requires_grad = False
    for param in self.robot_encoder.parameters():
      param.requires_grad = False
    n_frozen = sum(
        p.numel() for p in self.parameters() if not p.requires_grad
    )
    n_train  = sum(
        p.numel() for p in self.parameters() if p.requires_grad
    )
    print(f'Frozen {n_frozen:,} params  |  Trainable {n_train:,} params')

  # ── Override train() to keep frozen encoders in eval mode ─────────────────

  def train(self, mode: bool = True):
    """Keep frozen encoders in eval mode regardless of overall training mode.

    This prevents the Dropout layers inside obj_encoder and robot_encoder from
    randomly dropping activations during GeoMatch++ training.
    """
    super().train(mode)
    self.obj_encoder.eval()
    self.robot_encoder.eval()
    return self

  # ── Shared embedding helper ───────────────────────────────────────────────

  def encode_embed(self, encoder, feature, adj_mat, normalize_emb=True):
    x = encoder(feature, adj_mat)
    if normalize_emb:
      x = x.clone() / (torch.norm(x, dim=-1, keepdim=True) + 1e-6)
    return x

  # ── Forward pass ─────────────────────────────────────────────────────────

  def forward(
      self,
      obj_pc,               # [B, 2048, 3]       object point cloud
      robot_pc,             # [B, N_rob, 3]      robot surface points (~1000)
      robot_key_point_idx,  # [B, 6]             keypoint indices into robot_pc
      obj_adj,              # [B, 2048, 2048]    object graph adjacency
      robot_adj,            # [B, N_rob, N_rob]  robot graph adjacency
      xyz_prev,             # [B, 6, 3]          previous keypoint positions
      morph_features,       # [B, 32, 9]         morphology node features
      morph_adj,            # [B, 32, 32]        morphology graph adjacency
  ):
    # 1. Frozen object encoder
    obj_embed_raw = self.encode_embed(
        self.obj_encoder, obj_pc, obj_adj
    )                                             # [B, 2048, 512]

    # 2. Frozen robot surface encoder
    robot_embed = self.encode_embed(
        self.robot_encoder, robot_pc, robot_adj
    )                                             # [B, N_rob, 512]

    # 3. Morphology encoder (trainable)
    morph_embed = self.encode_embed(
        self.morphology_encoder, morph_features, morph_adj
    )                                             # [B, 32, 512]

    # 4. DCP cross-attention: object attends to morphology → residual
    obj_residual = self.dcp_transformer(
        obj_embed_raw, morph_embed
    )                                             # [B, 2048, 512]

    # 5. Enhanced object embedding
    obj_embed = obj_embed_raw + obj_residual      # [B, 2048, 512]

    # 6. Contact map prediction (dot product, identical to GeoMatch)
    robot_feat_size = robot_embed.shape[2]        # 512
    keypoint_feat = torch.gather(
        robot_embed,
        1,
        robot_key_point_idx[..., None].long().repeat(1, 1, robot_feat_size),
    )                                             # [B, 6, 512]
    contact_map_pred = torch.matmul(
        obj_embed, keypoint_feat.transpose(2, 1)
    )[..., None]                                  # [B, 2048, 6, 1]

    # 7. Projection heads
    obj_proj_embed   = self.obj_proj(obj_embed)    # [B, 2048, 64]
    robot_proj_embed = self.robot_proj(robot_embed) # [B, N_rob, 64]

    # 8. Autoregressive keypoint modules
    output_1 = self.kp_ar_model_1(
        obj_proj_embed, obj_pc, robot_proj_embed, xyz_prev
    )   # [B, 2048, 1]
    output_2 = self.kp_ar_model_2(
        obj_proj_embed, obj_pc, robot_proj_embed, xyz_prev
    )
    output_3 = self.kp_ar_model_3(
        obj_proj_embed, obj_pc, robot_proj_embed, xyz_prev
    )
    output_4 = self.kp_ar_model_4(
        obj_proj_embed, obj_pc, robot_proj_embed, xyz_prev
    )
    output_5 = self.kp_ar_model_5(
        obj_proj_embed, obj_pc, robot_proj_embed, xyz_prev
    )

    output = torch.cat(
        (output_1, output_2, output_3, output_4, output_5), dim=-1
    )[..., None]                                  # [B, 2048, 5, 1]

    return contact_map_pred, output

  # ── Loss ─────────────────────────────────────────────────────────────────

  def calc_loss(self, gt_contact_map, contact_map_pred, pred, label):
    """Identical loss to GeoMatch: 0.5 * l_dist + 0.5 * l_match."""
    flat_pred_map = contact_map_pred.view(
        contact_map_pred.shape[0]
        * contact_map_pred.shape[1]
        * contact_map_pred.shape[2],
        1,
    )
    flat_gt_map = gt_contact_map.view(
        gt_contact_map.shape[0]
        * gt_contact_map.shape[1]
        * gt_contact_map.shape[2],
        1,
    )

    pos_weight = torch.tensor([self.robot_weighting]).to(flat_pred_map.device)
    loss = nn.BCEWithLogitsLoss(pos_weight=pos_weight)(
        flat_pred_map, flat_gt_map
    )
    l_dist = torch.mean(loss)

    pos_weight = torch.tensor([self.match_weighting]).to(pred.device)
    match_losses = []
    for i in range(self.n_kp - 1):
      pred_i  = pred[:, :, i].reshape(-1, 1)
      label_i = label[:, :, i].reshape(-1, 1)
      match_losses.append(
          nn.BCEWithLogitsLoss(pos_weight=pos_weight)(pred_i, label_i)
      )
    l_match = torch.mean(torch.stack(match_losses))

    return self.dist_loss_weight * l_dist + self.match_loss_weight * l_match