"""
Collection of different neural network architectures together with
architecture-specific helper functions.
"""

from collections import Counter

from torch import nn
import torch
import numpy as np

import hannds_data as hd


class Network88(nn.Module):
    def __init__(self, hidden_size, n_layers, bidirectional, n_features, n_categories, rnn_type):
        super(Network88, self).__init__()
        self.rnn = _rnn_with_type(rnn_type, n_features, hidden_size, n_layers, bidirectional)

        self.n_directions = 2 if bidirectional else 1
        self.out_linear = nn.Linear(hidden_size * self.n_directions, n_features * n_categories)
        self.n_layers = n_layers
        self.n_features = n_features
        self.n_categories = n_categories
        self.criterion = nn.CrossEntropyLoss()

    def forward(self, input, hidden_in):
        rnn_output, hidden_out = self.rnn.forward(input, hidden_in)
        output = self.out_linear(rnn_output)
        output = output.view(-1, output.shape[1], self.n_features, self.n_categories)
        # Residual connection
        # output[:, :, :, 0] += 1.0 - input
        # output[:, :, :, 1] += input
        # output[:, :, :, 2] += input
        return output, hidden_out

    def compute_loss(self, output, labels):
        return self.criterion(output.view((-1, self.n_categories)), labels.view(-1))

    def compute_accuracy(self, X_batch, Y_batch, prediction):
        predicted_classes = torch.argmax(prediction, dim=3)
        filter_func = causal_filter
        X_numpy = X_batch.squeeze().cpu().numpy()
        Y_numpy = Y_batch.squeeze().cpu().numpy()
        classes_numpy = predicted_classes.squeeze().cpu().numpy()
        return compute_accuracy88(X_numpy, Y_numpy, classes_numpy, filter_func)
        # va1 = valid_accuracy_plain[0] + compute_accuracy88(X_numpy, Y_numpy, classes_numpy)
        # va2 = valid_accuracy_plain[1] + 1
        # valid_accuracy_plain = (va1, va2)


class Network88Tanh(nn.Module):
    def __init__(self, hidden_size, n_layers, bidirectional, n_features, rnn_type):
        super(Network88Tanh, self).__init__()
        self.rnn = _rnn_with_type(rnn_type, n_features, hidden_size, n_layers, bidirectional)
        self.n_directions = 2 if bidirectional else 1
        self.out_linear = nn.Linear(hidden_size * self.n_directions, n_features)
        self.n_layers = n_layers
        self.criterion = nn.MSELoss()

    def forward(self, input, hidden_in):
        rnn_output, hidden_out = self.rnn.forward(input, hidden_in)
        output = self.out_linear(rnn_output)
        output = torch.tanh(output)
        output = output * input
        return output, hidden_out

    def compute_loss(self, output, labels):
        return self.criterion(output, labels)

    def compute_accuracy(self, X_batch, Y_batch, prediction):
        num_notes = torch.sum(X_batch)
        left_hand = (prediction < 0.0).float()
        right_hand = (prediction > 0.0).float()
        prediction = left_hand * -1.0 + right_hand * 1.0
        prediction = causal_filter(prediction)
        diff = (prediction != Y_batch).float()
        errors_percent = torch.sum(diff) / num_notes * 100.0
        return 100.0 - errors_percent


class NetworkMidi(nn.Module):
    def __init__(self, hidden_size, n_layers, rnn_type, bidirectional):
        super(NetworkMidi, self).__init__()
        n_features = 4
        self.rnn = _rnn_with_type(rnn_type, n_features, hidden_size, n_layers, bidirectional=bidirectional)
        self.n_directions = 2 if bidirectional else 1
        self.linear = nn.Linear(hidden_size * self.n_directions, out_features=1)
        self.criterion = nn.BCEWithLogitsLoss()

    def forward(self, input, hidden_in):
        out, hidden_out = self.rnn(input, hidden_in)
        out = self.linear(out)
        return out, hidden_out

    def compute_loss(self, output, labels):
        return self.criterion(output.view(-1), labels.view(-1))

    def compute_accuracy(self, X_batch, Y_batch, prediction):
        labels_flat = Y_batch.view(-1)
        prediction_flat = (torch.sigmoid(prediction.view(-1)) >= 0.5).float()
        num_events = float(labels_flat.shape[0])
        num_correct = float(torch.sum(labels_flat == prediction_flat))