import unittest

from convokit.tests.politeness_strategies.politeness_corpus_util import parsed_politeness_test_corpus, parsed_politeness_test_zh_corpus
from convokit.politenessStrategies import PolitenessStrategies

class TestPolitenessStrategies(unittest.TestCase):

    def test_politeness_api(self):
        ps = PolitenessStrategies(strategy_collection='politeness_api')
        corpus = parsed_politeness_test_corpus()
    
        corpus = ps.transform(corpus)
        expected_strategy_uses = [
                {'feature_politeness_==1st_person==': 0,
                'feature_politeness_==1st_person_pl.==': 0,
                'feature_politeness_==1st_person_start==': 1,
                'feature_politeness_==2nd_person==': 1,
                'feature_politeness_==2nd_person_start==': 0,
                'feature_politeness_==Apologizing==': 0,
                'feature_politeness_==Deference==': 0,
                'feature_politeness_==Direct_question==': 0,
                'feature_politeness_==Direct_start==': 0,
                'feature_politeness_==Factuality==': 1,
                'feature_politeness_==Gratitude==': 1,
                'feature_politeness_==HASHEDGE==': 0,
                'feature_politeness_==HASNEGATIVE==': 0,
                'feature_politeness_==HASPOSITIVE==': 1,
                'feature_politeness_==Hedges==': 0,
                'feature_politeness_==INDICATIVE==': 0,
                'feature_politeness_==Indirect_(btw)==': 0,
                'feature_politeness_==Indirect_(greeting)==': 0,
                'feature_politeness_==Please==': 0,
                'feature_politeness_==Please_start==': 0,
                'feature_politeness_==SUBJUNCTIVE==': 0},
                {'feature_politeness_==1st_person==': 0,
                'feature_politeness_==1st_person_pl.==': 0,
                'feature_politeness_==1st_person_start==': 0,
                'feature_politeness_==2nd_person==': 1,
                'feature_politeness_==2nd_person_start==': 0,
                'feature_politeness_==Apologizing==': 0,
                'feature_politeness_==Deference==': 1,
                'feature_politeness_==Direct_question==': 0,
                'feature_politeness_==Direct_start==': 0,
                'feature_politeness_==Factuality==': 0,
                'feature_politeness_==Gratitude==': 0,
                'feature_politeness_==HASHEDGE==': 0,
                'feature_politeness_==HASNEGATIVE==': 0,
                'feature_politeness_==HASPOSITIVE==': 1,
                'feature_politeness_==Hedges==': 0,
                'feature_politeness_==INDICATIVE==': 0,
                'feature_politeness_==Indirect_(btw)==': 0,
                'feature_politeness_==Indirect_(greeting)==': 0,
                'feature_politeness_==Please==': 0,
                'feature_politeness_==Please_start==': 0,
                'feature_politeness_==SUBJUNCTIVE==': 0},
                {'feature_politeness_==1st_person==': 0,
                'feature_politeness_==1st_person_pl.==': 0,
                'feature_politeness_==1st_person_start==': 0,
                'feature_politeness_==2nd_person==': 1,
                'feature_politeness_==2nd_person_start==': 0,
                'feature_politeness_==Apologizing==': 0,
                'feature_politeness_==Deference==': 0,
                'feature_politeness_==Direct_question==': 0,
                'feature_politeness_==Direct_start==': 0,
                'feature_politeness_==Factuality==': 0,
                'feature_politeness_==Gratitude==': 0,
                'feature_politeness_==HASHEDGE==': 0,
                'feature_politeness_==HASNEGATIVE==': 0,
                'feature_politeness_==HASPOSITIVE==': 0,
                'feature_politeness_==Hedges==': 0,
                'feature_politeness_==INDICATIVE==': 0,
                'feature_politeness_==Indirect_(btw)==': 0,
                'feature_politeness_==Indirect_(greeting)==': 1,
                'feature_politeness_==Please==': 0,
                'feature_politeness_==Please_start==': 0,
                'feature_politeness_==SUBJUNCTIVE==': 0},
                {'feature_politeness_==1st_person==': 1,
                'feature_politeness_==1st_person_pl.==': 0,
                'feature_politeness_==1st_person_start==': 0,
                'feature_politeness_==2nd_person==': 1,
                'feature_politeness_==2nd_person_start==': 0,
                'feature_politeness_==Apologizing==': 1,
                'feature_politeness_==Deference==': 0,
                'feature_politeness_==Direct_question==': 0,
                'feature_politeness_==Direct_start==': 0,
                'feature_politeness_==Factuality==': 0,
                'feature_politeness_==Gratitude==': 0,
                'feature_politeness_==HASHEDGE==': 0,
                'feature_politeness_==HASNEGATIVE==': 1,
                'feature_politeness_==HASPOSITIVE==': 1,
                'feature_politeness_==Hedges==': 0,
                'feature_politeness_==INDICATIVE==': 0,
                'feature_politeness_==Indirect_(btw)==': 0,
                'feature_politeness_==Indirect_(greeting)==': 0,
                'feature_politeness_==Please==': 0,
                'feature_politeness_==Please_start==': 0,
                'feature_politeness_==SUBJUNCTIVE==': 0},
                {'feature_politeness_==1st_person==': 0,
                'feature_politeness_==1st_person_pl.==': 0,
                'feature_politeness_==1st_person_start==': 0,
                'feature_politeness_==2nd_person==': 1,
                'feature_politeness_==2nd_person_start==': 0,
                'feature_politeness_==Apologizing==': 0,
                'feature_politeness_==Deference==': 0,
                'feature_politeness_==Direct_question==': 0,
                'feature_politeness_==Direct_start==': 0,
                'feature_politeness_==Factuality==': 0,
                'feature_politeness_==Gratitude==': 0,
                'feature_politeness_==HASHEDGE==': 0,
                'feature_politeness_==HASNEGATIVE==': 0,
                'feature_politeness_==HASPOSITIVE==': 0,
                'feature_politeness_==Hedges==': 0,
                'feature_politeness_==INDICATIVE==': 0,
                'feature_politeness_==Indirect_(btw)==': 0,
                'feature_politeness_==Indirect_(greeting)==': 0,
                'feature_politeness_==Please==': 1,
                'feature_politeness_==Please_start==': 0,
                'feature_politeness_==SUBJUNCTIVE==': 1},
                {'feature_politeness_==1st_person==': 0,
                'feature_politeness_==1st_person_pl.==': 0,
                'feature_politeness_==1st_person_start==': 0,
                'feature_politeness_==2nd_person==': 0,
                'feature_politeness_==2nd_person_start==': 0,
                'feature_politeness_==Apologizing==': 0,
                'feature_politeness_==Deference==': 0,
                'feature_politeness_==Direct_question==': 0,
                'feature_politeness_==Direct_start==': 0,
                'feature_politeness_==Factuality==': 0,
                'feature_politeness_==Gratitude==': 0,
                'feature_politeness_==HASHEDGE==': 0,
                'feature_politeness_==HASNEGATIVE==': 0,
                'feature_politeness_==HASPOSITIVE==': 0,
                'feature_politeness_==Hedges==': 0,
                'feature_politeness_==INDICATIVE==': 0,
                'feature_politeness_==Indirect_(btw)==': 0,
                'feature_politeness_==Indirect_(greeting)==': 0,
                'feature_politeness_==Please==': 0,
                'feature_politeness_==Please_start==': 1,
                'feature_politeness_==SUBJUNCTIVE==': 0},
                {'feature_politeness_==1st_person==': 0,
                'feature_politeness_==1st_person_pl.==': 0,
                'feature_politeness_==1st_person_start==': 0,
                'feature_politeness_==2nd_person==': 1,
                'feature_politeness_==2nd_person_start==': 0,
                'feature_politeness_==Apologizing==': 0,
                'feature_politeness_==Deference==': 0,
                'feature_politeness_==Direct_question==': 0,
                'feature_politeness_==Direct_start==': 0,
                'feature_politeness_==Factuality==': 0,
                'feature_politeness_==Gratitude==': 0,
                'feature_politeness_==HASHEDGE==': 0,
                'feature_politeness_==HASNEGATIVE==': 0,
                'feature_politeness_==HASPOSITIVE==': 0,
                'feature_politeness_==Hedges==': 0,
                'feature_politeness_==INDICATIVE==': 0,
                'feature_politeness_==Indirect_(btw)==': 1,
                'feature_politeness_==Indirect_(greeting)==': 0,
                'feature_politeness_==Please==': 0,
                'feature_politeness_==Please_start==': 0,
                'feature_politeness_==SUBJUNCTIVE==': 0},
                {'feature_politeness_==1st_person==': 0,
                'feature_politeness_==1st_person_pl.==': 0,
                'feature_politeness_==1st_person_start==': 0,
                'feature_politeness_==2nd_person==': 1,
                'feature_politeness_==2nd_person_start==': 0,
                'feature_politeness_==Apologizing==': 0,
                'feature_politeness_==Deference==': 0,
                'feature_politeness_==Direct_question==': 1,
                'feature_politeness_==Direct_start==': 0,
                'feature_politeness_==Factuality==': 0,
                'feature_politeness_==Gratitude==': 0,
                'feature_politeness_==HASHEDGE==': 0,
                'feature_politeness_==HASNEGATIVE==': 0,
                'feature_politeness_==HASPOSITIVE==': 0,
                'feature_politeness_==Hedges==': 0,
                'feature_politeness_==INDICATIVE==': 0,
                'feature_politeness_==Indirect_(btw)==': 0,
                'feature_politeness_==Indirect_(greeting)==': 0,
                'feature_politeness_==Please==': 0,
                'feature_politeness_==Please_start==': 0,
                'feature_politeness_==SUBJUNCTIVE==': 0},
                {'feature_politeness_==1st_person==': 0,
                'feature_politeness_==1st_person_pl.==': 0,
                'feature_politeness_==1st_person_start==': 0,
                'feature_politeness_==2nd_person==': 1,
                'feature_politeness_==2nd_person_start==': 0,
                'feature_politeness_==Apologizing==': 0,
                'feature_politeness_==Deference==': 0,
                'feature_politeness_==Direct_question==': 0,
                'feature_politeness_==Direct_start==': 1,
                'feature_politeness_==Factuality==': 0,
                'feature_politeness_==Gratitude==': 0,
                'feature_politeness_==HASHEDGE==': 0,
                'feature_politeness_==HASNEGATIVE==': 0,
                'feature_politeness_==HASPOSITIVE==': 0,
                'feature_politeness_==Hedges==': 0,
                'feature_politeness_==INDICATIVE==': 1,
                'feature_politeness_==Indirect_(btw)==': 0,
                'feature_politeness_==Indirect_(greeting)==': 0,
                'feature_politeness_==Please==': 0,
                'feature_politeness_==Please_start==': 0,
                'feature_politeness_==SUBJUNCTIVE==': 0},
                {'feature_politeness_==1st_person==': 0,
                'feature_politeness_==1st_person_pl.==': 0,
                'feature_politeness_==1st_person_start==': 0,
                'feature_politeness_==2nd_person==': 1,
                'feature_politeness_==2nd_person_start==': 0,
                'feature_politeness_==Apologizing==': 0,
                'feature_politeness_==Deference==': 0,
                'feature_politeness_==Direct_question==': 0,
                'feature_politeness_==Direct_start==': 0,
                'feature_politeness_==Factuality==': 0,
                'feature_politeness_==Gratitude==': 0,
                'feature_politeness_==HASHEDGE==': 0,
                'feature_politeness_==HASNEGATIVE==': 0,
                'feature_politeness_==HASPOSITIVE==': 0,
                'feature_politeness_==Hedges==': 0,
                'feature_politeness_==INDICATIVE==': 0,
                'feature_politeness_==Indirect_(btw)==': 0,
                'feature_politeness_==Indirect_(greeting)==': 0,
                'feature_politeness_==Please==': 1,
                'feature_politeness_==Please_start==': 0,
                'feature_politeness_==SUBJUNCTIVE==': 1},
                {'feature_politeness_==1st_person==': 0,
                'feature_politeness_==1st_person_pl.==': 0,
                'feature_politeness_==1st_person_start==': 0,
                'feature_politeness_==2nd_person==': 1,
                'feature_politeness_==2nd_person_start==': 0,
                'feature_politeness_==Apologizing==': 0,
                'feature_politeness_==Deference==': 0,
                'feature_politeness_==Direct_question==': 0,
                'feature_politeness_==Direct_start==': 0,
                'feature_politeness_==Factuality==': 0,
                'feature_politeness_==Gratitude==': 0,
                'feature_politeness_==HASHEDGE==': 0,
                'feature_politeness_==HASNEGATIVE==': 0,
                'feature_politeness_==HASPOSITIVE==': 0,
                'feature_politeness_==Hedges==': 0,
                'feature_politeness_==INDICATIVE==': 1,
                'feature_politeness_==Indirect_(btw)==': 0,
                'feature_politeness_==Indirect_(greeting)==': 0,
                'feature_politeness_==Please==': 0,
                'feature_politeness_==Please_start==': 0,
                'feature_politeness_==SUBJUNCTIVE==': 0},
                {'feature_politeness_==1st_person==': 0,
                'feature_politeness_==1st_person_pl.==': 1,
                'feature_politeness_==1st_person_start==': 1,
                'feature_politeness_==2nd_person==': 0,
                'feature_politeness_==2nd_person_start==': 0,
                'feature_politeness_==Apologizing==': 0,
                'feature_politeness_==Deference==': 0,
                'feature_politeness_==Direct_question==': 0,
                'feature_politeness_==Direct_start==': 0,
                'feature_politeness_==Factuality==': 0,
                'feature_politeness_==Gratitude==': 0,
                'feature_politeness_==HASHEDGE==': 1,
                'feature_politeness_==HASNEGATIVE==': 0,
                'feature_politeness_==HASPOSITIVE==': 1,
                'feature_politeness_==Hedges==': 1,
                'feature_politeness_==INDICATIVE==': 0,
                'feature_politeness_==Indirect_(btw)==': 0,
                'feature_politeness_==Indirect_(greeting)==': 0,
                'feature_politeness_==Please==': 0,
                'feature_politeness_==Please_start==': 0,
                'feature_politeness_==SUBJUNCTIVE==': 0},
                {'feature_politeness_==1st_person==': 0,
                'feature_politeness_==1st_person_pl.==': 1,
                'feature_politeness_==1st_person_start==': 0,
                'feature_politeness_==2nd_person==': 0,
                'feature_politeness_==2nd_person_start==': 0,
                'feature_politeness_==Apologizing==': 0,
                'feature_politeness_==Deference==': 0,
                'feature_politeness_==Direct_question==': 0,
                'feature_politeness_==Direct_start==': 0,
                'feature_politeness_==Factuality==': 1,
                'feature_politeness_==Gratitude==': 0,
                'feature_politeness_==HASHEDGE==': 1,
                'feature_politeness_==HASNEGATIVE==': 0,
                'feature_politeness_==HASPOSITIVE==': 1,
                'feature_politeness_==Hedges==': 0,
                'feature_politeness_==INDICATIVE==': 0,
                'feature_politeness_==Indirect_(btw)==': 0,
                'feature_politeness_==Indirect_(greeting)==': 0,
                'feature_politeness_==Please==': 0,
                'feature_politeness_==Please_start==': 0,
                'feature_politeness_==SUBJUNCTIVE==': 0}
        ]
        
        for expected_use, utterance in zip(expected_strategy_uses, corpus.iter_utterances()):
            extracted_use = utterance.retrieve_meta('politeness_strategies')
            for k, v in expected_use.items():
                self.assertEqual(v, extracted_use[k])

    
    def test_politeness_cscw_zh(self):
        ps = PolitenessStrategies(strategy_collection='politeness_cscw_zh')
        corpus = parsed_politeness_test_zh_corpus()
   
        corpus = ps.transform(corpus)
        expected_strategy_uses = [
               {'apologetic': 0,
                'best_wishes': 0,
                'can_you': 0,
                'could_you': 0,
                'emergency': 0,
                'factuality': 0,
                'first_person_plural': 0,
                'first_person_singular': 0,
                'gratitude': 1,
                'greeting': 0,
                'hedge': 0,
                'honorifics': 0,
                'indirect_btw': 0,
                'ingroup_iden': 0,
                'please': 0,
                'praise': 0,
                'promise': 0,
                'start_i': 0,
                'start_please': 0,
                'start_question': 0,
                'start_so': 0,
                'start_you': 0,
                'taboo': 0,
                'together': 0,
                'you_direct': 0,
                'you_honorific': 1},
               {'apologetic': 0,
                'best_wishes': 0,
                'can_you': 0,
                'could_you': 0,
                'emergency': 0,
                'factuality': 0,
                'first_person_plural': 0,
                'first_person_singular': 0,
                'gratitude': 0,
                'greeting': 0,
                'hedge': 0,
                'honorifics': 0,
                'indirect_btw': 0,
                'ingroup_iden': 0,
                'please': 0,
                'praise': 1,
                'promise': 0,
                'start_i': 0,
                'start_please': 0,
                'start_question': 0,
                'start_so': 0,
                'start_you': 0,
                'taboo': 1,
                'together': 0,
                'you_direct': 0,
                'you_honorific': 0},
               {'apologetic': 0,
                'best_wishes': 0,
                'can_you': 0,
                'could_you': 0,
                'emergency': 0,
                'factuality': 0,
                'first_person_plural': 0,
                'first_person_singular': 0,
                'gratitude': 0,
                'greeting': 1,
                'hedge': 0,
                'honorifics': 0,
                'indirect_btw': 0,
                'ingroup_iden': 0,
                'please': 0,
                'praise': 0,
                'promise': 0,
                'start_i': 0,
                'start_please': 0,
                'start_question': 0,
                'start_so': 0,
                'start_you': 0,
                'taboo': 0,
                'together': 0,
                'you_direct': 1,
                'you_honorific': 0},
               {'apologetic': 1,
                'best_wishes': 0,
                'can_you': 0,
                'could_you': 0,
                'emergency': 0,
                'factuality': 0,
                'first_person_plural': 0,
                'first_person_singular': 1,
                'gratitude': 0,
                'greeting': 0,
                'hedge': 0,
                'honorifics': 0,
                'indirect_btw': 0,
                'ingroup_iden': 0,
                'please': 0,
                'praise': 0,
                'promise': 0,
                'start_i': 0,
                'start_please': 0,
                'start_question': 0,
                'start_so': 0,
                'start_you': 1,
                'taboo': 0,
                'together': 0,
                'you_direct': 1,
                'you_honorific': 0},
               {'apologetic': 0,
                'best_wishes': 0,
                'can_you': 0,
                'could_you': 0,
                'emergency': 0,
                'factuality': 0,
                'first_person_plural': 0,
                'first_person_singular': 0,
                'gratitude': 0,
                'greeting': 0,
                'hedge': 1,
                'honorifics': 0,
                'indirect_btw': 0,
                'ingroup_iden': 0,
                'please': 1,
                'praise': 0,
                'promise': 0,
                'start_i': 0,
                'start_please': 0,
                'start_question': 0,
                'start_so': 0,
                'start_you': 0,
                'taboo': 0,
                'together': 0,
                'you_direct': 0,
                'you_honorific': 1},
               {'apologetic': 0,
                'best_wishes': 0,
                'can_you': 0,
                'could_you': 0,
                'emergency': 0,
                'factuality': 0,
                'first_person_plural': 1,
                'first_person_singular': 0,
                'gratitude': 0,
                'greeting': 0,
                'hedge': 0,
                'honorifics': 0,
                'indirect_btw': 0,
                'ingroup_iden': 0,
                'please': 0,
                'praise': 0,
                'promise': 0,
                'start_i': 0,
                'start_please': 1,
                'start_question': 0,
                'start_so': 0,
                'start_you': 0,
                'taboo': 0,
                'together': 0,
                'you_direct': 0,
                'you_honorific': 0},
               {'apologetic': 0,
                'best_wishes': 0,
                'can_you': 1,
                'could_you': 0,
                'emergency': 0,
                'factuality': 0,
                'first_person_plural': 0,
                'first_person_singular': 0,
                'gratitude': 0,
                'greeting': 0,
                'hedge': 0,
                'honorifics': 0,
                'indirect_btw': 0,
                'ingroup_iden': 0,
                'please': 0,
                'praise': 0,
                'promise': 0,
                'start_i': 0,
                'start_please': 0,
                'start_question': 0,
                'start_so': 0,
                'start_you': 0,
                'taboo': 0,
                'together': 0,
                'you_direct': 1,
                'you_honorific': 0},
               {'apologetic': 0,
                'best_wishes': 0,
                'can_you': 0,
                'could_you': 0,
                'emergency': 0,
                'factuality': 0,
                'first_person_plural': 0,
                'first_person_singular': 0,
                'gratitude': 0,
                'greeting': 0,
                'hedge': 0,
                'honorifics': 0,
                'indirect_btw': 0,
                'ingroup_iden': 0,
                'please': 0,
                'praise': 0,
                'promise': 0,
                'start_i': 0,
                'start_please': 0,
                'start_question': 1,
                'start_so': 0,
                'start_you': 0,
                'taboo': 0,
                'together': 0,
                'you_direct': 0,
                'you_honorific': 0},
               {'apologetic': 0,
                'best_wishes': 0,
                'can_you': 0,
                'could_you': 0,
                'emergency': 0,
                'factuality': 0,
                'first_person_plural': 0,
                'first_person_singular': 0,
                'gratitude': 0,
                'greeting': 0,
                'hedge': 1,
                'honorifics': 0,
                'indirect_btw': 0,
                'ingroup_iden': 0,
                'please': 0,
                'praise': 0,
                'promise': 0,
                'start_i': 0,
                'start_please': 0,
                'start_question': 0,
                'start_so': 0,
                'start_you': 0,
                'taboo': 0,
                'together': 0,
                'you_direct': 0,
                'you_honorific': 0},
               {'apologetic': 0,
                'best_wishes': 0,
                'can_you': 0,
                'could_you': 0,
                'emergency': 0,
                'factuality': 1,
                'first_person_plural': 0,
                'first_person_singular': 1,
                'gratitude': 0,
                'greeting': 0,
                'hedge': 0,
                'honorifics': 0,
                'indirect_btw': 0,
                'ingroup_iden': 0,
                'please': 0,
                'praise': 0,
                'promise': 0,
                'start_i': 0,
                'start_please': 0,
                'start_question': 0,
                'start_so': 0,
                'start_you': 0,
                'taboo': 0,
                'together': 0,
                'you_direct': 0,
                'you_honorific': 0}
        ]

        for expected_use, utterance in zip(expected_strategy_uses, corpus.iter_utterances()):
            extracted_use = utterance.retrieve_meta('politeness_strategies')
            for k, v in expected_use.items():
                self.assertEqual(v, extracted_use[k])
