#!/usr/bin/env python3
# coding=utf-8

import torch
import torch.nn as nn

from model.module.encoder import Encoder

from model.module.transformer import Decoder
from model.head.node_centric_head import NodeCentricHead
from model.head.labeled_edge_head import LabeledEdgeHead
from model.head.sequential_head import SequentialHead
from utility.utils import create_padding_mask


class Model(nn.Module):
    def __init__(self, dataset, args, initialize=True):
        super(Model, self).__init__()
        self.encoder = Encoder(args, dataset)
        if args.n_layers > 0:
            self.decoder = Decoder(args)
        else:
            self.decoder = lambda x, *args: x  # identity function, which ignores all arguments except the first one

        if args.graph_mode == "sequential":
            self.head = SequentialHead(dataset, args, initialize)
        elif args.graph_mode == "node-centric":
            self.head = NodeCentricHead(dataset, args, initialize)
        elif args.graph_mode == "labeled-edge":
            self.head = LabeledEdgeHead(dataset, args, initialize)

        self.query_length = args.query_length
        self.dataset = dataset
        self.args = args

    def forward(self, batch, inference=False, **kwargs):
        every_input, word_lens = batch["every_input"]
        decoder_lens = self.query_length * word_lens
        batch_size, input_len = every_input.size(0), every_input.size(1)
        device = every_input.device

        encoder_mask = create_padding_mask(batch_size, input_len, word_lens, device)
        decoder_mask = create_padding_mask(batch_size, self.query_length * input_len, decoder_lens, device)

        encoder_output, decoder_input = self.encoder(batch["input"], batch["char_form_input"], batch["input_scatter"], input_len)

        decoder_output = self.decoder(decoder_input, encoder_output, decoder_mask, encoder_mask)

        if inference:
            return self.head.predict(encoder_output, decoder_output, encoder_mask, decoder_mask, batch)
        else:
            return self.head(encoder_output, decoder_output, encoder_mask, decoder_mask, batch)

    def get_params_for_optimizer(self, args):
        encoder_decay, encoder_no_decay = self.get_encoder_parameters(args.n_encoder_layers)
        decoder_decay, decoder_no_decay = self.get_decoder_parameters()

        parameters = [{"params": p, "weight_decay": args.encoder_weight_decay} for p in encoder_decay]
        parameters += [{"params": p, "weight_decay": 0.0} for p in encoder_no_decay]
        parameters += [
            {"params": decoder_decay, "weight_decay": args.decoder_weight_decay},
            {"params": decoder_no_decay, "weight_decay": 0.0},
        ]
        return parameters

    def get_decoder_parameters(self):
        no_decay = ["bias", "LayerNorm.weight", "_norm.weight"]
        decay_params = (p for name, p in self.named_parameters() if not any(nd in name for nd in no_decay) and not name.startswith("encoder.bert") and p.requires_grad)
        no_decay_params = (p for name, p in self.named_parameters() if any(nd in name for nd in no_decay) and not name.startswith("encoder.bert") and p.requires_grad)

        return decay_params, no_decay_params

    def get_encoder_parameters(self, n_layers):
        no_decay = ["bias", "LayerNorm.weight", "_norm.weight"]
        decay_params = [
            [p for name, p in self.named_parameters() if not any(nd in name for nd in no_decay) and name.startswith(f"encoder.bert.encoder.layer.{n_layers - 1 - i}.") and p.requires_grad] for i in range(n_layers)
        ]
        no_decay_params = [
            [p for name, p in self.named_parameters() if any(nd in name for nd in no_decay) and name.startswith(f"encoder.bert.encoder.layer.{n_layers - 1 - i}.") and p.requires_grad] for i in range(n_layers)
        ]

        return decay_params, no_decay_params