#
#   Copyright © 2021 Uncharted Software Inc.
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.

import sys
import numpy as np
from tqdm import tqdm

import torch
from torch import nn
from torch.utils.data import DataLoader, Dataset

from torchvision import transforms
from torchvision.datasets.folder import pil_loader
from torchvision.models import resnet18, resnet34, resnet50, resnet101, densenet161

from basenet import BaseNet

from .base import DistilBaseModel
from .forest import ForestCV
from .helpers import tiebreaking_vote
from .metrics import metrics, classification_metrics, regression_metrics

# --
# IO helper


class PathDataset(Dataset):
    def __init__(self, paths, transform=None, preload=True):
        self.paths = paths

        self.preload = preload
        if self.preload:
            print("PathDataset: preloading", file=sys.stderr)
            self._samples = []
            for p in tqdm(self.paths):
                self._samples.append(pil_loader(p))

        self.transform = transform

    def __getitem__(self, idx):

        if not self.preload:
            sample = pil_loader(self.paths[idx])
        else:
            sample = self._samples[idx]

        if self.transform is not None:
            sample = self.transform(sample)

        return sample, -1

    def __len__(self):
        return self.paths.shape[0]


# --
# Model


class FixedCNNFeatureExtractor(BaseNet):
    def __init__(self, base_model, drop_last=1):
        super().__init__()
        self._model = nn.Sequential(*list(base_model.children())[:-drop_last])

    def forward(self, x):
        x = self._model(x)
        while len(x.shape) > 2:
            x = x.mean(dim=-1).squeeze()

        return x


class FixedCNNForest(DistilBaseModel):
    def __init__(self, target_metric):

        self.target_metric = target_metric
        self.is_classification = target_metric in classification_metrics

        self._feature_extractors = [
            resnet18,
            resnet34,
            resnet50,
            resnet101,
            densenet161,
        ]
        self._y_train = None
        self._models = None

        self.transform = transforms.Compose(
            [
                transforms.Resize((224, 224)),
                transforms.ToTensor(),
                transforms.Normalize(
                    mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
                ),
            ]
        )

    def extract_features(self, fe, dataloaders, mode):
        model = FixedCNNFeatureExtractor(fe(pretrained=True)).to("cuda")
        model.verbose = True
        _ = model.eval()

        feats, _ = model.predict(dataloaders, mode=mode)
        del model

        return feats

    def fit(self, X_train, y_train, U_train=None):
        self._y_train = y_train

        dataloaders = {
            "train": DataLoader(
                PathDataset(paths=X_train, transform=self.transform),
                batch_size=32,
                shuffle=False,
            ),
        }

        self._models = []
        for fe in self._feature_extractors:
            train_feats = self.extract_features(fe, dataloaders, mode="train")
            model = ForestCV(target_metric=self.target_metric)
            model = model.fit(train_feats, y_train)
            self._models.append(model)

        return self

    def predict(self, X):
        dataloaders = {
            "test": DataLoader(
                PathDataset(paths=X, transform=self.transform),
                batch_size=32,
                shuffle=False,
            ),
        }

        all_preds = []
        for fe, model in zip(self._feature_extractors, self._models):
            test_feats = self.extract_features(fe, dataloaders, mode="test")
            all_preds.append(model.predict(test_feats))

        if self.is_classification:
            ens_pred = tiebreaking_vote(np.vstack(all_preds), self._y_train)
        else:
            ens_pred = np.stack(all_preds).mean(axis=0)

        return ens_pred
