from collections import defaultdict

import nltk


class Place:
    title = None
    words = None

    def __init__(self, title: str):
        title = title.lower()
        self.title = title
        self.words = nltk.word_tokenize(title)

    def __repr__(self) -> str:
        return 'Place(%s)' % self.title

    def contains_word(self, word: str):
        return word in self.words


place_titles = [
    'Moscow', 'New York', 'New Castle', 'Atlanta'
]

all_places = []

places_by_word = defaultdict(list)

text1 = '''
The streets of New York City are so desolate now that you half expect tumbleweed to blow along the pavement where cars and cabs once clustered. There is barely a plane in the sky. You hear the wheeze of an empty bus rounding a corner, the flutter of pigeons on a fire escape, the wail of an ambulance. The sirens are unnervingly frequent. But even on these sunny, early-spring days there are few people in sight. For weeks, as the distancing rules of the pandemic took hold, a
'''

text2 = '''
Lines of cars stretched for blocks around five Atlanta Public Schools facilities on Saturday as families sheltering in place from the coronavirus pandemic scrambled for food before the system’s spring break holidays.
Atlanta Public Schools’ emergency meal distribution sites were swamped Saturday, with all five locations running out of food after less than two hours.The school system, which is closed through April 10, will have no food distribution over the break, offered families a week’s worth of food at Saturday’s pickup sites. APS Superintendent Meria Carstarphen said long lines were waiting when they opened at 9 a.m.“It’s been a massive demand,” she said. “We knew people would come and we’d distribute all the food. We didn’t know how fast it would happen.”
'''

def init_system():
    for title in place_titles:
        place = Place(title)
        all_places.append(place)
        for word in place.words:
            places_by_word[word].append(place)


def break_into_words(text: str):
    words = nltk.word_tokenize(text)
    return words


class Searcher:
    def __init__(self, places: set, words):
        self.places = places
        self.words = words
        self.i = 0

    def search(self):
        self.i = 0
        while self.i < len(self.words):
            word = self.words[self.i]
            places1 = places_by_word.get(word)
            if places1:
                places2 = list(places1)
                while self.i < len(self.words):
                    self.i += 1
                    word_next = self.words[self.i]
                    places3 = list(filter(lambda x: x.contains_word(word_next), places2))
                    if places3:
                        places2 = places3
                    else:
                        self.places.update(places2)
                        break
                # self.places.update(places1)
            self.i += 1


def find_place(text: str) -> list:
    text = text.lower()
    places_result = set()
    words = break_into_words(text)
    searcher = Searcher(places_result, words)
    searcher.search()
    return list(places_result)


def main():
    global all_places, places_by_word
    # nltk.download()
    init_system()
    for text in [text1, text2]:
        pp = find_place(text)
        print('text: %s' % text)
        print('places: %s' % pp)


if __name__ == '__main__':
    main()