"""Top-level model classes.

Author:
    Chris Chute (chute@stanford.edu)
"""

import char_layers
import torch
import torch.nn as nn


class BiDAF(nn.Module):
    """Baseline BiDAF model for SQuAD.

    Based on the paper:
    "Bidirectional Attention Flow for Machine Comprehension"
    by Minjoon Seo, Aniruddha Kembhavi, Ali Farhadi, Hannaneh Hajishirzi
    (https://arxiv.org/abs/1611.01603).

    Follows a high-level structure commonly found in SQuAD models:
        - Embedding layer: Embed word indices to get word vectors.
        - Encoder layer: Encode the embedded sequence.
        - Attention layer: Apply an attention mechanism to the encoded sequence.
        - Model encoder layer: Encode the sequence again.
        - Output layer: Simple layer (e.g., fc + softmax) to get final outputs.

    Args:
        word_vectors (torch.Tensor): Pre-trained word vectors.
        char_vectors (torch.Tensor): Pre-trained character vectors.
        hidden_size (int): Number of features in the hidden state at each layer.
        drop_prob (float): Dropout probability.
    """
    def __init__(self, word_vectors, char_vectors, hidden_size, channel_size, channel_width, drop_prob=0.):
        super(BiDAF, self).__init__()
        self.emb = char_layers.Embedding(word_vectors=word_vectors,
                                    char_vectors=char_vectors,
                                    hidden_size=hidden_size,
                                    channel_size=channel_size,
                                    channel_width=channel_width,
                                    drop_prob=drop_prob)

        self.enc = char_layers.RNNEncoder(input_size=2 * hidden_size,
                                     hidden_size=hidden_size,
                                     num_layers=1,
                                     drop_prob=drop_prob)

        self.att = char_layers.BiDAFAttention(hidden_size=2 * hidden_size,
                                         drop_prob=drop_prob)

        self.mod = char_layers.RNNEncoder(input_size=8 * hidden_size,
                                     hidden_size=hidden_size,
                                     num_layers=2,
                                     drop_prob=drop_prob)

        self.out = char_layers.BiDAFOutput(hidden_size=hidden_size,
                                      drop_prob=drop_prob)

    def forward(self, cw_idxs, cc_idxs, qw_idxs, qc_idxs):
        cw_mask = torch.zeros_like(cw_idxs) != cw_idxs
        qw_mask = torch.zeros_like(qw_idxs) != qw_idxs
        cc_mask = torch.zeros_like(cc_idxs) != cc_idxs
        qc_mask = torch.zeros_like(qc_idxs) != qc_idxs
        
        cw_len, qw_len = cw_mask.sum(-1), qw_mask.sum(-1)
        cc_len, qc_len = cc_mask.sum(-1), qc_mask.sum(-1)

        c_emb = self.emb(cw_idxs,cc_idxs)         # (batch_size, c_len, hidden_size)
        q_emb = self.emb(qw_idxs,qc_idxs)         # (batch_size, q_len, hidden_size)
        #print('c_mask',c_mask)
        #print('q_mask',q_mask)
        #print('c and q len',c_len,q_len)

        #print('c_emb',c_emb)
        #print('q_emb',q_emb)
        self.enc.rnn.flatten_parameters()

        c_enc = self.enc(c_emb, cw_len)    # (batch_size, c_len, 2 * hidden_size)
        q_enc = self.enc(q_emb, qw_len)    # (batch_size, q_len, 2 * hidden_size)
        #print("c_enc: ", c_enc.size(), c_enc)
        #print("q_enc: ", q_enc.size(), q_enc)
        att = self.att(c_enc, q_enc,
                       cw_mask, qw_mask)    # (batch_size, c_len, 8 * hidden_size)

        #print("att: ", att.size(), att)
        self.mod.rnn.flatten_parameters()
        mod = self.mod(att, cw_len)        # (batch_size, c_len, 2 * hidden_size)
        #print("mod: ", mod.size(), mod)

        self.out.rnn.rnn.flatten_parameters()
        out = self.out(att, mod, cw_mask)  # 2 tensors, each (batch_size, c_len)
        #print("out: ", out)
        return out
