"""DSSM Preprocessor."""

from matchzoo import engine
from matchzoo import preprocessor
from matchzoo import datapack


class DSSMPreprocessor(engine.BasePreprocessor):
    """
    DSSM preprocessor helper.

    Example:
        >>> train_inputs = [
        ...     ("beijing", "Beijing is capital of China", 1, ('id0', 'id1')),
        ...     ("beijing", "China is in east Asia", 0, ('id0', 'id2')),
        ...     ("beijing", "Summer in Beijing is hot.", 1, ('id0', 'id3'))
        ... ]
        >>> dssm_preprocessor = DSSMPreprocessor()
        >>> rv_train = dssm_preprocessor.fit_transform(
        ...     train_inputs,
        ...     stage='train')
        >>> dssm_preprocessor.context['dim_triletter']
        25
        >>> type(rv_train)
        <class 'matchzoo.datapack.DataPack'>
        >>> test_inputs = [("beijing",
        ...                 "I visted beijing yesterday.",
        ...                 ('id0', 'id4'))]
        >>> rv_test = dssm_preprocessor.transform(test_inputs, stage='test')
        >>> type(rv_test)
        <class 'matchzoo.datapack.DataPack'>

    """

    def __init__(self):
        """Initialization."""
        self._context = {}

    @property
    def context(self):
        """Get fitted parameters."""
        return self._context

    def _prepare_stateless_units(self):
        """Prepare."""
        return [
            preprocessor.TokenizeUnit(),
            preprocessor.LowercaseUnit(),
            preprocessor.PuncRemovalUnit(),
            preprocessor.StopRemovalUnit(),
            preprocessor.NgramLetterUnit()
        ]

    def _build_vocab(self, inputs):
        """Build vocabulary before fit transform."""
        vocab = []
        units = self._prepare_stateless_units()
        for input in inputs:
            for unit in units:
                left = unit.transform(input[0])
                right = unit.transform(input[1])
            vocab.extend(left + right)
        return vocab

    def fit(self, inputs):
        """Fit parameters."""
        vocab = self._build_vocab(inputs)
        vocab_unit = preprocessor.VocabularyUnit()
        vocab_unit.fit(vocab)
        self._context['term_index'] = vocab_unit.state['term_index']
        self._context['dim_triletter'] = len(
            vocab_unit.state['term_index']) + 1
        return self

    def transform(self, inputs, stage):
        """Transform."""
        output_left = []
        output_righ = []
        labels = []
        # ids is used to store (qid, did) pairs.
        ids = []
        if stage not in ['train', 'test']:
            msg = f'{stage} is not a valid stage name'
            msg += '`train` or `test` expected.'
            raise ValueError(msg)
        if not self._context.get('term_index'):
            raise ValueError(
                "Please fit term_index before apply transofm function.")
        units = self._prepare_stateless_units()
        units.append(
            preprocessor.WordHashingUnit(self._context['term_index']))
        for input in inputs:
            for unit in units:
                left = unit.transform(input[0])
                righ = unit.transform(input[1])
            output_left.append(left)
            output_righ.append(righ)
            if stage == "train":
                labels.append(input[2])
                ids.append(input[3])
            else:
                ids.append(input[2])
        data = {'text_left': output_left,
                'text_right': output_righ,
                'id': ids}
        if stage == "train":
            data['label'] = labels
        return datapack.DataPack(data=data, context=self._context)

@@ -1278,16 +1278,145 @@
 ontext%0A%0A
+    @context.setter%0A    def context(self, context: dict):%0A        %22%22%22Set pre-fitted context.%22%22%22%0A        self._context = context%0A%0A
     def 
@@ -2086,16 +2086,416 @@
  vocab%0A%0A
+    def _check_transoform_state(self, stage):%0A        %22%22%22check.%22%22%22%0A        if stage not in %5B'train', 'test'%5D:%0A            msg = f'%7Bstage%7D is not a valid stage name'%0A            msg += '%60train%60 or %60test%60 expected.'%0A            raise ValueError(msg)%0A        if not self._context.get('term_index'):%0A            raise ValueError(%0A                %22Please fit term_index before apply transofm function.%22)%0A%0A
     def 
@@ -2927,466 +2927,116 @@
     
-output_left = %5B%5D%0A        output_righ = %5B%5D%0A        labels = %5B%5D%0A        # ids is used to store (qid, did) pairs.%0A        ids = %5B%5D%0A        if stage not in %5B'train', 'test'%5D:%0A            msg = f'%7Bstage%7D is not a valid stage name'%0A            msg += '%60train%60 or %60test%60 expected.'%0A            raise ValueError(msg)%0A        if not self._context.get('term_index'):%0A            raise ValueError(%0A                %22Please fit term_index before apply transofm function.%22
+self._check_transoform_state(stage)%0A        lst_left, lst_right, lst_labels, lst_ids = (%5B%5D for i in range(4)
 )%0A  
@@ -3342,21 +3342,18 @@
         
-outpu
+ls
 t_left.a
@@ -3376,27 +3376,25 @@
         
-outpu
+ls
 t_righ
+t
 .append(
@@ -3448,16 +3448,20 @@
         
+lst_
 labels.a
@@ -3484,32 +3484,36 @@
                 
+lst_
 ids.append(input
@@ -3551,16 +3551,20 @@
         
+lst_
 ids.appe
@@ -3605,21 +3605,18 @@
 _left': 
-outpu
+ls
 t_left,%0A
@@ -3649,19 +3649,17 @@
 t': 
-outpu
+ls
 t_righ
+t
 ,%0A  
@@ -3678,16 +3678,20 @@
   'id': 
+lst_
 ids%7D%0A   
@@ -3744,16 +3744,20 @@
 bel'%5D = 
+lst_
 labels%0A 

add context setter