r""" PF-WILLOW dataset """

import os

import pandas as pd
import numpy as np
import torch

from .dataset import CorrespondenceDataset


class PFWillowDataset(CorrespondenceDataset):

    def __init__(self, benchmark, datapath, thres, split):
        r"""PF-WILLOW dataset constructor"""
        super(PFWillowDataset, self).__init__(benchmark, datapath, thres, split)

        self.train_data = pd.read_csv(self.spt_path)
        self.src_imnames = np.array(self.train_data.iloc[:, 0])
        self.trg_imnames = np.array(self.train_data.iloc[:, 1])
        self.src_kps = self.train_data.iloc[:, 2:22].values
        self.trg_kps = self.train_data.iloc[:, 22:].values
        self.cls = ['car(G)', 'car(M)', 'car(S)', 'duck(S)',
                    'motorbike(G)', 'motorbike(M)', 'motorbike(S)',
                    'winebottle(M)', 'winebottle(wC)', 'winebottle(woC)']
        self.cls_ids = list(map(lambda names: self.cls.index(names.split('/')[1]), self.src_imnames))
        self.src_imnames = list(map(lambda x: os.path.join(*x.split('/')[1:]), self.src_imnames))
        self.trg_imnames = list(map(lambda x: os.path.join(*x.split('/')[1:]), self.trg_imnames))

    def __getitem__(self, idx):
        r""" Constructs and returns a batch for PF-WILLOW dataset """
        batch = super(PFWillowDataset, self).__getitem__(idx)
        batch['pckthres'] = self.get_pckthres(batch)

        return batch

    def get_pckthres(self, batch):
        r""" Computes PCK threshold """
        if self.thres == 'bbox':
            return max(batch['trg_kps'].max(1)[0] - batch['trg_kps'].min(1)[0]).clone()
        elif self.thres == 'img':
            return torch.tensor(max(batch['trg_img'].size()[1], batch['trg_img'].size()[2]))
        else:
            raise Exception('Invalid pck evaluation level: %s' % self.thres)

    def get_points(self, pts_list, idx, org_imsize):
        r""" Returns key-points of an image """
        point_coords = pts_list[idx, :].reshape(2, 10)
        point_coords = torch.tensor(point_coords.astype(np.float32))
        xy, n_pts = point_coords.size()
        pad_pts = torch.zeros((xy, self.max_pts - n_pts)) - 2
        x_crds = point_coords[0] * (self.img_size / org_imsize[0])
        y_crds = point_coords[1] * (self.img_size / org_imsize[1])
        kps = torch.cat([torch.stack([x_crds, y_crds]), pad_pts], dim=1)

        return kps, n_pts