import os.path

from .cityscapes import remap_labels_to_train_ids
import numpy as np
import scipy.io
import torch
import torch.utils.data as data
from PIL import Image

from .data_loader import register_data_params, register_dataset_obj
from .data_loader import DatasetParams
from .cityscapes import id2label as LABEL2TRAIN
import pandas as pd
import random

@register_data_params('cape_dataset_src')
class CapeParamsSrc(DatasetParams):
    num_channels = 3
    image_size = 200#TODO (jm) (to use resize this can't be equal to net.imagesize )
    mean = 0.5
    std = 0.5
    num_cls = 8
    target_transform = None

@register_dataset_obj('cape_dataset_src')
class CapeSrc(data.Dataset):

    def __init__(self, root,train=True, download=True,
     num_cls=19, split='train', remap_labels=True, 
            transform=None, target_transform=None):
            
        print('ROOT', root)
        pass
        # self.root = root
        # self.split = split
        # self.remap_labels = remap_labels
        # #self.ids = self.collect_ids()
        # self.transform = transform
        # self.target_transform = target_transform
        # m = scipy.io.loadmat(os.path.join(self.root, 'mapping.mat'))
        # full_classes = [x[0] for x in m['classes'][0]]
        # self.classes = []
        # for old_id, new_id in LABEL2TRAIN.items():
        #     if not new_id == 255 and old_id > 0:
        #         self.classes.append(full_classes[old_id])
        # self.num_cls = 19
        self.train = train
        self.opt = 1#opt
        self.root = './cyclegan/cape_analytics/' #1 opt.dataroot

        if self.train:
            self.path_A_df = './cyclegan/cape_analytics/trainvaltrain.csv'
        else:
            self.path_A_df = './cyclegan/cape_analytics/test_a.csv'

        self.df_A = pd.read_csv(self.path_A_df)
        self.category_columns = sorted([i for i in self.df_A if i.startswith('nvotes')])

        self.A_size = len(self.df_A)

        print('transform', transform)
        # new code 
        self.transform = transform
    def img_path(self, id):
        filename = '{:05d}.png'.format(id)
        return os.path.join(self.root, 'images', filename)

    def label_path(self, id):
        filename = '{:05d}.png'.format(id)
        return os.path.join(self.root, 'labels', filename)

    def __getitem__(self, index):
        A_row = self.df_A.iloc[index % self.A_size, :]

        A_path = os.path.join(self.root, A_row['filename'])

        A_img = Image.open(A_path).convert('RGB')

        A = self.transform(A_img)

        A_label = A_row[self.category_columns].values.argmax()

        return A, A_label

    def __len__(self):
        return self.A_size

