"""Matchzoo point generator."""

from matchzoo import engine
from matchzoo import tasks
from matchzoo.datapack import DataPack
from matchzoo import utils

import numpy as np
import typing


class PointGenerator(engine.BaseGenerator):
    """PointGenerator for Matchzoo.

    Ponit generator can be used for classification as well as ranking.

    Examples:
        >>> import pandas as pd
        >>> relation = [['qid0', 'did0', 1]]
        >>> left = [['qid0', [1, 2]]]
        >>> right = [['did0', [2, 3]]]
        >>> relation = pd.DataFrame(relation,
        ...                         columns=['id_left', 'id_right', 'label'])
        >>> left = pd.DataFrame(left, columns=['id_left', 'text_left'])
        >>> left.set_index('id_left', inplace=True)
        >>> left['length_left'] = left.apply(lambda x: len(x['text_left']),
        ...                                  axis=1)
        >>> right = pd.DataFrame(right, columns=['id_right', 'text_right'])
        >>> right.set_index('id_right', inplace=True)
        >>> right['length_right'] = right.apply(lambda x: len(x['text_right']),
        ...                                     axis=1)
        >>> input = datapack.DataPack(relation=relation,
        ...                           left=left,
        ...                           right=right
        ... )
        >>> task = tasks.Classification()
        >>> generator = PointGenerator(input, task, 1, 'train', False)
        >>> x, y = generator[0]
        >>> x['text_left'].tolist()
        [[1, 2]]
        >>> x['text_right'].tolist()
        [[2, 3]]
        >>> x['id_left'].tolist()
        ['qid0']
        >>> x['id_right'].tolist()
        ['did0']
        >>> x['length_left'].tolist()
        [2]
        >>> x['length_right'].tolist()
        [2]
        >>> y.tolist()
        [[0.0, 1.0]]

    """

    def __init__(
        self,
        datapack,
        task: engine.BaseTask = tasks.Classification(2),
        batch_size: int = 32,
        shuffle: bool = True,
        post_process=None
    ):
        """Construct the point generator.

        :param datapack: the output generated by :class:`DataPack`.
        :param task: the task is a instance of :class:`engine.BaseTask`.
        :param batch_size: number of instances in a batch.
            `evaluate`, or `predict` expected.
        :param shuffle: whether to shuffle the instances while generating a
            batch.
        """
        self._datapack = datapack
        self._task = task
        self._post_process = post_process
        super().__init__(batch_size, len(datapack.relation), shuffle)

    def _get_batch_of_transformed_samples(
        self,
        index_array: np.array
    ) -> typing.Tuple[dict, typing.Any]:
        """Get a batch of samples based on their ids.

        :param index_array: a list of instance ids.
        :return: A batch of transformed samples.
        """
        sliced_data = self._datapack[index_array]
        if self._datapack.stage == 'train':
            columns = list(sliced_data.columns)
            columns.remove('label')
            x = sliced_data[columns].to_dict(orient='list')
            for key, val in x.items():
                x[key] = np.array(val)
            y = np.array(sliced_data['label'])
            return x, y
        else:
            x = sliced_data.to_dict(orient='list')
            for key, val in x.items():
                x[key] = np.array(val)
            return x, None

@@ -2955,264 +2955,123 @@
     
-if self._datapack.stage == 'train':%0A            columns = list(sliced_data.columns)%0A            columns.remove('label')%0A            x = sliced_data%5Bcolumns%5D.to_dict(orient='list')%0A            for key, val in x.items():%0A                x%5Bkey%5D = np.array(val
+columns = list(sliced_data.columns)%0A%0A        if self._datapack.stage == 'train':%0A            columns.remove('label'
 )%0A  
@@ -3127,42 +3127,36 @@
     
-    return x, y%0A        else:%0A    
+else:%0A            y = None%0A%0A
     
@@ -3174,16 +3174,25 @@
 ced_data
+%5Bcolumns%5D
 .to_dict
@@ -3215,20 +3215,16 @@
         
-  
-  
 for key,
@@ -3254,20 +3254,16 @@
         
-    
 x%5Bkey%5D =
@@ -3277,20 +3277,17 @@
 ay(val)%0A
-    
+%0A
         
@@ -3296,13 +3296,10 @@
 turn x, 
-None
+y
 %0A

Refactor point generator.