"""
Uses model outputs as input to an XGBoost classifier
"""
!pip install git+https://github.com/pabloppp/pytorch-tools -U

# %% Setup
import os
import shutil
import numpy  as np
import pandas as pd

import torch
import torch.nn            as nn
import torch.nn.functional as F

from PIL           import Image
from torchtools.nn import SimpleSelfAttention
from torchvision   import transforms

from xgboost.sklearn import XGBClassifier
from joblib import dump, load


for dirname, _, filenames in os.walk('/kaggle/input'):
    for filename in filenames:
        print(os.path.join(dirname, filename))

#%% Constants
INPUT_DIR = "/kaggle/input/Kannada-MNIST"
DATA_DIR  = "submission"
MODEL_DIR = "/kaggle/input/kannada-mnist-ensemble-v2"


#%% Codebase
# %% Data preparation
# Load the dataset
df = pd.read_csv(INPUT_DIR + "/test.csv")

# Image ID storage
ids = []

# Create a temporary folder if required
if not os.path.exists(DATA_DIR + "/"):
    os.mkdir(DATA_DIR)

if not os.path.exists(DATA_DIR + "/images/"):
    os.mkdir(DATA_DIR + "/images")

# Loop through each row
for i, row in df.iterrows():
    # Get the data components
    img = np.reshape(row[1:].values, newshape=(28, 28)).astype(np.uint8)

    # Convert
    img = Image.fromarray(img)

    # Fetch the label
    label = row[0]

    # Write to storage
    img.save(DATA_DIR + "/images/{}.png".format(row[0]))
    ids.append(row[0])

# Save the IDs to a DataFrame and write to CSV
label_df = pd.DataFrame({"id": ids, "label": 0})
label_df.to_csv("labels.csv", index=False)

# %% Network
class KannadaNet(nn.Module):
    def __init__(self
                 , input_shape
                 , n_class
                 ):
        """
        Builds a model for the Kannada MNIST dataset
        Based on: https://www.kaggle.com/bustam/cnn-in-keras-for-kannada-digits
        :param input_shape:
        :param n_class:
        """
        super(KannadaNet, self).__init__()

        # Build the layers
        self._build_layers(input_shape=input_shape, n_class=n_class)

        # Initialize weights
        self._init_weight()

        # Parameter counts
        print("Model contains {} parameters".format(sum(p.numel() for p in self.parameters() if p.requires_grad)))

    def forward(self, x):
        x = self._conv_forward(x)

        # Classification
        x = self.fc_1(x)
        x = F.relu(x)
        x = self.bn_1(x)

        x = self.fc_2(x)
        x = F.relu(x)
        x = self.bn_2(x)

        confidence = self.confidence(x)
        out = self.classifier(x)

        # return out
        return {"prediction": out.cpu().float(), "confidence": confidence.cpu().float()}

    def _build_layers(self, input_shape, n_class):
        """
        Adds layers to the network
        :param input_shape:
        :param n_class:
        :return:
        """
        # Create the stem of the network
        self.stem = Conv2D_BN(in_channels=input_shape[0]
                              , out_channels=32
                              , padding=1
                              )

        # Inception blocks
        self.inception_res_1 = InceptionBlock(in_shape=[32, 28, 28], residual=True)
        self.attention_1 = SimpleSelfAttention(self.inception_res_1.output_shape[0])
        self.inception_res_2 = InceptionBlock(
            in_shape=[int(x / y) for x, y in zip(self.inception_res_1.output_shape, [1, 2, 2])], residual=True)
        self.attention_2 = SimpleSelfAttention(self.inception_res_2.output_shape[0])
        self.inception_res_3 = InceptionBlock(
            in_shape=[int(x / y) for x, y in zip(self.inception_res_2.output_shape, [1, 2, 2])], residual=True)
        self.attention_3 = SimpleSelfAttention(self.inception_res_3.output_shape[0])

        # Pooling
        self.pool = nn.MaxPool2d(kernel_size=2)

        # Determine the output dimensions
        conv_out_dims = self._conv_forward(torch.zeros([1] + input_shape, dtype=torch.float32))

        # Linear layers
        self.fc_1 = nn.Linear(in_features=conv_out_dims.shape[1], out_features=512)
        self.bn_1 = nn.BatchNorm1d(num_features=512)
        self.fc_2 = nn.Linear(in_features=512, out_features=512)
        self.bn_2 = nn.BatchNorm1d(num_features=512)

        self.confidence = nn.Linear(in_features=512, out_features=1)
        self.classifier = nn.Linear(in_features=512, out_features=n_class)

    # Convolutional forward
    def _conv_forward(self, x):
        """
        Forward function for convolutional layers
        :param x:
        :return:
        """
        x = self.stem(x)
        x = self.inception_res_1(x)
        x = self.attention_1(x)
        x = self.pool(x)
        x = self.inception_res_2(x)
        x = self.attention_2(x)
        x = self.pool(x)
        x = self.inception_res_3(x)
        x = self.attention_3(x)
        x = self.pool(x)

        # Flatten
        x = torch.flatten(x, start_dim=1)

        return x

    # Weight initialization
    def _init_weight(self):
        for m in self.modules():
            print(m)
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight)
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1.0)
                m.bias.data.zero_()

# %% Utility functions
class InceptionBlock(nn.Module):
    def __init__(self, in_shape, residual=True):
        super(InceptionBlock, self).__init__()
        # Class parameters
        self.residual = residual

        # Define the component operations
        self.conv_1x1_1 = nn.Conv2d(kernel_size=1, in_channels=in_shape[0], out_channels=int(in_shape[0] / 2))
        self.bn_1 = nn.BatchNorm2d(num_features=self.conv_1x1_1.out_channels)

        self.conv_1x1_2 = nn.Conv2d(kernel_size=1, in_channels=in_shape[0], out_channels=int(in_shape[0]))
        self.bn_2 = nn.BatchNorm2d(num_features=self.conv_1x1_2.out_channels)

        self.conv_1x1_3 = nn.Conv2d(kernel_size=1, in_channels=in_shape[0], out_channels=int(in_shape[0] / 2))
        self.bn_3 = nn.BatchNorm2d(num_features=self.conv_1x1_3.out_channels)

        self.conv_1x1_4 = nn.Conv2d(kernel_size=1, in_channels=in_shape[0], out_channels=int(in_shape[0] / 2))
        self.bn_4 = nn.BatchNorm2d(num_features=self.conv_1x1_4.out_channels)

        self.conv_3x3_1 = nn.Conv2d(kernel_size=3, in_channels=self.conv_1x1_2.out_channels,
                                    out_channels=self.conv_1x1_1.out_channels * 2, padding=1)
        self.bn_5 = nn.BatchNorm2d(num_features=self.conv_3x3_1.out_channels)

        self.conv_3x3_2 = nn.Conv2d(kernel_size=3, in_channels=self.conv_1x1_3.out_channels,
                                    out_channels=int(in_shape[0] / 4), padding=1)
        self.bn_6 = nn.BatchNorm2d(num_features=self.conv_3x3_2.out_channels)

        self.conv_3x3_3 = nn.Conv2d(kernel_size=3, in_channels=self.conv_3x3_2.out_channels,
                                    out_channels=int(in_shape[0] / 4), padding=1)
        self.bn_7 = nn.BatchNorm2d(num_features=self.conv_3x3_3.out_channels)

        self.pool_3x3 = nn.MaxPool2d(kernel_size=3, stride=1, padding=1)

        # Determine the block output shape
        self.output_shape = self._get_out_shape(in_shape)

    def forward(self, x):
        # First path
        x_1 = self.conv_1x1_1(x)
        x_1 = F.relu(x_1)
        x_1 = self.bn_1(x_1)

        # Second path
        x_2 = self.conv_1x1_2(x)
        x_2 = F.relu(x_2)
        x_2 = self.bn_2(x_2)

        x_2 = self.conv_3x3_1(x_2)
        x_2 = F.relu(x_2)
        x_2 = self.bn_5(x_2)

        # Third path
        x_3 = self.conv_1x1_3(x)
        x_3 = F.relu(x_3)
        x_3 = self.bn_3(x_3)

        x_3 = self.conv_3x3_2(x_3)
        x_3 = F.relu(x_3)
        x_3 = self.bn_6(x_3)

        x_3 = self.conv_3x3_3(x_3)
        x_3 = F.relu(x_3)
        x_3 = self.bn_7(x_3)

        # Fourth path
        x_4 = self.pool_3x3(x)
        x_4 = self.conv_1x1_4(x_4)
        x_4 = F.relu(x_4)
        x_4 = self.bn_4(x_4)

        # Concatenate
        if self.residual:
            x = torch.cat([x_1, x_2, x_3, x_4, x], dim=1)
        else:
            x = torch.cat([x_1, x_2, x_3, x_4], dim=1)

        return x

    def _get_out_shape(self, input_shape):
        return list(self.forward(torch.zeros([1] + input_shape)).shape[1:])


class Conv2D_BN(nn.Module):
    def __init__(self, in_channels, out_channels, kernel=3, stride=1, padding=1, activation=True):
        """
        Convolution followed by Batch Normalization, with optional Mish activation
        :param in_features:
        :param out_features:
        :param kernel_size:
        :param stride:
        :param padding:
        :param groups:
        :param activation:
        """
        super(Conv2D_BN, self).__init__()

        # Class parameters
        self.activation = activation

        # Convolution and BatchNorm class
        self.conv = nn.Conv2d(in_channels=in_channels
                              , out_channels=out_channels
                              , kernel_size=kernel
                              , stride=stride
                              , padding=padding
                              , bias=True
                              )

        self.bn = nn.BatchNorm2d(num_features=out_channels, momentum=0.9)

    def forward(self, x):
        x = self.conv(x)
        if self.activation:
            x = F.relu(x)
        x = self.bn(x)

        return x


def block_forward(layers, x):
    for layer in layers:
        x = layer(x)

    return x


# Mish Activation
def Mish(x):
    r"""
    Mish activation function is proposed in "Mish: A Self
    Regularized Non-Monotonic Neural Activation Function"
    paper, https://arxiv.org/abs/1908.08681.
    """

    return x * torch.tanh(F.softplus(x))


# %% Dataset
# Dataset class
class KannadaMNIST(object):
    def __init__( self
                , image_folder    : str
                , labels          : pd.DataFrame
                , transforms      : list  = None
                ):
        """
        Class for creating a dataset for training the flow chart object detector
        """
        # Class attributes
        self.image_folder = image_folder
        self.labels       = labels
        self.transforms   = transforms

        # Create an index list
        self.index = np.arange(self.labels.shape[0])

    def __len__(self):
        return len(self.index)

    def __getitem__(self, idx):
        # Load the labels
        label = self.labels.iloc[self.index[idx]]

        # Image ID
        id = label["id"]

        # Load the image
        img = Image.open(self.image_folder + "/" + str(label["id"]) + ".png")

        # Populate the labels
        label = label["label"]

        # Apply transformations
        img_ = self.transforms(img)

        return {"image": img_, "label": label, "id": id}

# %% Evaluation
# Set up the dataset
ds = KannadaMNIST( image_folder=DATA_DIR + "/images"
                 , labels      =label_df
                 , transforms  =transforms.ToTensor()
                 )
loader = torch.utils.data.DataLoader(dataset=ds, batch_size=256)

models      = []
for i in range(6):
    # Network
    net = KannadaNet( input_shape=[1, 28, 28]
                    , n_class=10
                    )

    # Load the weights
    net.load_state_dict(torch.load("{}/Submission weights {}.tar".format(MODEL_DIR, i)))

    # Activate evaluation
    net.eval()
    net.cuda()
    models.append(net)

#%%
# Predictions from each model
with torch.no_grad():
    X = torch.stack([torch.cat([net(x["image"].cuda())["prediction"] for x in loader], dim=0)
                     for net in models
                    ]
                   , dim=1
                   )

#%%
# Reshape for XGBoost model
X = X.reshape((X.shape[0], -1))

# Load the XGBoost model
xclas = load("{}/kannda_xgboost.joblib.dat".format(MODEL_DIR))

# Check the class distribution
y_ = xclas.predict(X.reshape((X.shape[0], -1)))

#%% Submission
label_df["label"] = y_
label_df.to_csv("submission.csv", index=False)

#%% Cleanup
shutil.rmtree(DATA_DIR)