#!/usr/bin/env python
# coding: utf-8

# pylint: disable-all
from __future__ import absolute_import
from sklearn.preprocessing import LabelEncoder
from pathlib import Path

import torch
from torch.autograd import Variable
import torch.nn as nn
import torch.optim as optim


class BinModel(nn.Module):
    expected_target_type = torch.FloatTensor

    def __init__(self, input_size):
        super(BinModel, self).__init__()
        self.fc1 = nn.Linear(input_size, 50)
        self.relu1 = nn.ReLU()
        self.dout = nn.Dropout(0.2)
        self.fc2 = nn.Linear(50, 100)
        self.prelu = nn.PReLU(1)
        self.out = nn.Linear(100, 1)
        self.out_act = nn.Sigmoid()

    def forward(self, input_):
        a1 = self.fc1(input_)
        h1 = self.relu1(a1)
        dout = self.dout(h1)
        a2 = self.fc2(dout)
        h2 = self.prelu(a2)
        a3 = self.out(h2)
        y = self.out_act(a3)
        return y


class RegModel(nn.Module):
    def __init__(self, input_size):
        super(RegModel, self).__init__()
        self.fc1 = nn.Linear(input_size, 50)
        self.relu1 = nn.ReLU()
        self.dout = nn.Dropout(0.2)
        self.fc2 = nn.Linear(50, 100)
        self.prelu = nn.PReLU(1)
        self.out = nn.Linear(100, 1)

    def forward(self, input_):
        a1 = self.fc1(input_)
        h1 = self.relu1(a1)
        dout = self.dout(h1)
        a2 = self.fc2(dout)
        h2 = self.prelu(a2)
        y = self.out(h2)
        return y


class MultiModel(nn.Module):
    expected_target_type = torch.LongTensor

    def __init__(self, input_size, output_size):
        super(MultiModel, self).__init__()
        self.layer1 = nn.Linear(input_size, 8)
        self.relu = nn.ReLU()
        self.layer2 = nn.Linear(8, output_size)
        self.out = nn.Softmax()

    def forward(self, input_):
        out = self.layer1(input_)
        out = self.relu(out)
        out = self.layer2(out)
        out = self.out(out)
        return out


def train_epoch(model, opt, criterion, X, y, batch_size=50):
    model.train()
    losses = []
    for beg_i in range(0, X.size(0), batch_size):
        x_batch = X[beg_i : beg_i + batch_size, :]
        # y_hat will be (batch_size, 1) dim, so coerce target to look the same
        y_batch = y[beg_i : beg_i + batch_size].reshape(-1, 1)
        x_batch = Variable(x_batch)
        y_batch = Variable(y_batch)

        opt.zero_grad()
        # (1) Forward
        y_hat = model(x_batch)
        # (2) Compute diff
        loss = criterion(y_hat, y_batch)
        # (3) Compute gradients
        loss.backward()
        # (4) update weights
        opt.step()
        losses.append(loss.data.numpy())
    return losses


def build_classifier(X, num_labels):
    class_model = BinModel(X.shape[1]) if num_labels == 2 else MultiModel(X.shape[1], num_labels)
    class_opt = optim.Adam(class_model.parameters(), lr=0.001)
    class_criterion = nn.BCELoss() if num_labels == 2 else nn.CrossEntropyLoss()

    return class_model, class_opt, class_criterion


def build_regressor(X):
    reg_model = RegModel(X.shape[1])
    reg_opt = optim.Adam(reg_model.parameters(), lr=0.001)
    reg_criterion = nn.MSELoss()

    return reg_model, reg_opt, reg_criterion


def train_classifier(X, y, class_model, class_opt, class_criterion, n_epochs=5):
    target_encoder = LabelEncoder()
    target_encoder.fit(y)
    transformed_y = target_encoder.transform(y)
    bin_t_X = torch.from_numpy(X.values).type(torch.FloatTensor)
    bin_t_y = torch.from_numpy(transformed_y).type(class_model.expected_target_type)

    for e in range(n_epochs):
        train_epoch(class_model, class_opt, class_criterion, bin_t_X, bin_t_y)


def train_regressor(X, y, reg_model, reg_opt, reg_criterion, n_epochs=5):
    reg_t_X = torch.from_numpy(X.values).type(torch.FloatTensor)
    reg_t_y = torch.from_numpy(y.values).type(torch.FloatTensor)

    for e in range(n_epochs):
        train_epoch(reg_model, reg_opt, reg_criterion, reg_t_X, reg_t_y)


def save_torch_model(model, output_dir_path, filename="torch_bin.pth"):
    output_file_path = Path(output_dir_path) / filename
    torch.save(model, output_file_path)


def subset_data(X):
    numerics = ["int16", "int32", "int64", "float16", "float32", "float64"]
    # exclude any completely-missing columns when checking for numerics
    num_features = list(X.dropna(axis=1, how="all").select_dtypes(include=numerics).columns)
    # keep numeric features, zero-impute any missing values
    # obviously this is a very rudimentary approach to handling missing values
    # a more sophisticated imputer can be implemented by making use of custom transform, load, and predict hooks
    return X[num_features].fillna(0)

[((111, 20, 111, 32), 'torch.nn.MSELoss', 'nn.MSELoss', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((117, 21, 117, 35), 'sklearn.preprocessing.LabelEncoder', 'LabelEncoder', ({}, {}), '()', False, 'from sklearn.preprocessing import LabelEncoder\n'), ((137, 4, 137, 39), 'torch.save', 'torch.save', ({(137, 15, 137, 20): 'model', (137, 22, 137, 38): 'output_file_path'}, {}), '(model, output_file_path)', False, 'import torch\n'), ((20, 19, 20, 44), 'torch.nn.Linear', 'nn.Linear', ({(20, 29, 20, 39): 'input_size', (20, 41, 20, 43): '50'}, {}), '(input_size, 50)', True, 'import torch.nn as nn\n'), ((21, 21, 21, 30), 'torch.nn.ReLU', 'nn.ReLU', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((22, 20, 22, 35), 'torch.nn.Dropout', 'nn.Dropout', ({(22, 31, 22, 34): '0.2'}, {}), '(0.2)', True, 'import torch.nn as nn\n'), ((23, 19, 23, 37), 'torch.nn.Linear', 'nn.Linear', ({(23, 29, 23, 31): '50', (23, 33, 23, 36): '100'}, {}), '(50, 100)', True, 'import torch.nn as nn\n'), ((24, 21, 24, 32), 'torch.nn.PReLU', 'nn.PReLU', ({(24, 30, 24, 31): '1'}, {}), '(1)', True, 'import torch.nn as nn\n'), ((25, 19, 25, 36), 'torch.nn.Linear', 'nn.Linear', ({(25, 29, 25, 32): '100', (25, 34, 25, 35): '1'}, {}), '(100, 1)', True, 'import torch.nn as nn\n'), ((26, 23, 26, 35), 'torch.nn.Sigmoid', 'nn.Sigmoid', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((42, 19, 42, 44), 'torch.nn.Linear', 'nn.Linear', ({(42, 29, 42, 39): 'input_size', (42, 41, 42, 43): '50'}, {}), '(input_size, 50)', True, 'import torch.nn as nn\n'), ((43, 21, 43, 30), 'torch.nn.ReLU', 'nn.ReLU', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((44, 20, 44, 35), 'torch.nn.Dropout', 'nn.Dropout', ({(44, 31, 44, 34): '0.2'}, {}), '(0.2)', True, 'import torch.nn as nn\n'), ((45, 19, 45, 37), 'torch.nn.Linear', 'nn.Linear', ({(45, 29, 45, 31): '50', (45, 33, 45, 36): '100'}, {}), '(50, 100)', True, 'import torch.nn as nn\n'), ((46, 21, 46, 32), 'torch.nn.PReLU', 'nn.PReLU', ({(46, 30, 46, 31): '1'}, {}), '(1)', True, 'import torch.nn as nn\n'), ((47, 19, 47, 36), 'torch.nn.Linear', 'nn.Linear', ({(47, 29, 47, 32): '100', (47, 34, 47, 35): '1'}, {}), '(100, 1)', True, 'import torch.nn as nn\n'), ((64, 22, 64, 46), 'torch.nn.Linear', 'nn.Linear', ({(64, 32, 64, 42): 'input_size', (64, 44, 64, 45): '8'}, {}), '(input_size, 8)', True, 'import torch.nn as nn\n'), ((65, 20, 65, 29), 'torch.nn.ReLU', 'nn.ReLU', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((66, 22, 66, 47), 'torch.nn.Linear', 'nn.Linear', ({(66, 32, 66, 33): '8', (66, 35, 66, 46): 'output_size'}, {}), '(8, output_size)', True, 'import torch.nn as nn\n'), ((67, 19, 67, 31), 'torch.nn.Softmax', 'nn.Softmax', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((84, 18, 84, 35), 'torch.autograd.Variable', 'Variable', ({(84, 27, 84, 34): 'x_batch'}, {}), '(x_batch)', False, 'from torch.autograd import Variable\n'), ((85, 18, 85, 35), 'torch.autograd.Variable', 'Variable', ({(85, 27, 85, 34): 'y_batch'}, {}), '(y_batch)', False, 'from torch.autograd import Variable\n'), ((103, 22, 103, 34), 'torch.nn.BCELoss', 'nn.BCELoss', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((103, 59, 103, 80), 'torch.nn.CrossEntropyLoss', 'nn.CrossEntropyLoss', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((136, 23, 136, 44), 'pathlib.Path', 'Path', ({(136, 28, 136, 43): 'output_dir_path'}, {}), '(output_dir_path)', False, 'from pathlib import Path\n'), ((120, 14, 120, 40), 'torch.from_numpy', 'torch.from_numpy', ({(120, 31, 120, 39): 'X.values'}, {}), '(X.values)', False, 'import torch\n'), ((121, 14, 121, 45), 'torch.from_numpy', 'torch.from_numpy', ({(121, 31, 121, 44): 'transformed_y'}, {}), '(transformed_y)', False, 'import torch\n'), ((128, 14, 128, 40), 'torch.from_numpy', 'torch.from_numpy', ({(128, 31, 128, 39): 'X.values'}, {}), '(X.values)', False, 'import torch\n'), ((129, 14, 129, 40), 'torch.from_numpy', 'torch.from_numpy', ({(129, 31, 129, 39): 'y.values'}, {}), '(y.values)', False, 'import torch\n')]