"""A store for Event objects."""
from DesignationCache import DesignationCache
from Event import Event, EventFileFlags, EventSource
from FileStore import FileStore
from FileFactory import FileFactory
from math import floor
from constants import FD_OPEN, FD_CLOSE
from utils import time2Str, debugEnabled
import sys

class EventStore(object):
    """A store for Event objects.

    EventStore is a store for Events. It sorts all Events by timestamp to
    allow their simulation and the building of a filesystem and information
    flow models.
    """
    __event_store = None

    @staticmethod
    def get():
        """Return the EventStore for the entire application."""
        if not EventStore.__event_store:
            EventStore.__event_store = EventStore()
        return EventStore.__event_store

    @staticmethod
    def reset():
        EventStore.__event_store = None

    def __init__(self):
        """Construct an EventStore."""
        super(EventStore, self).__init__()
        self.clear()

    def clear(self):
        """Empty the EventStore."""
        self.store = list()                   # type: list
        self._sorted = True                   # type: bool
        self.desigcache = DesignationCache()  # type: DesignationCache

    def append(self, event: Event):
        """Append an Event to the store. The store will no longer be sorted."""
        if event.time == 0:
            raise ValueError("Events must have a timestamp.")
        self.store.append(event)
        self._sorted = False

    def insert(self, event: Event):
        """Insert an Event. Maintains the store sorted, if it was sorted."""
        if event.time == 0:
            raise ValueError("Events must have a timestamp.")

        # Binary search, first part: find a good index where to insert
        targetval = event.time
        targetb = 0
        minb = 0
        maxb = len(self.store)
        while maxb > minb + 1:
            currentb = minb + floor((maxb-minb)/2)
            currentval = self.store[currentb].time

            # Equal timestamps: put event after events with the same timestamp
            if currentval == targetval:
                targetb = currentb
                break
                pass
            # Event must go before the current value
            elif currentval > targetval:
                maxb = currentb
            # Event must go after the current value
            else:  # store[nextb] < target:
                minb = currentb
        else:
            # Reached if minb == maxb, aka event goes between minb and maxb
            targetb = minb

        # This serves two purposes: firstly, it inserts the new event after
        # those with a strictly equal timestamp. secondly, it 'finishes' the
        # previous loop by ensuring that we place our targetb cursor on the
        # nearest larger value. This is necessary because insert will insert
        # before the index it's given, rather than after, so we must target
        # the nearest larger value. We could have used ceiling instead of
        # floor in the previous loop to avoid that but we still needed to
        # iterate when timestamps are equal.
        maxb = len(self.store)
        while targetb < maxb and self.store[targetb].time <= targetval:
            targetb += 1
        self.store.insert(targetb, event)

    def sort(self):
        """Sort all the inserted Events by timestamp."""
        self.store = sorted(self.store, key=lambda x: x.time)
        self._sorted = True

    def getAllEvents(self):
        """Return all the events in the store. Very memory-intensive."""
        return self.store

    def getEventCount(self):
        """Return the number of events currently in the store."""
        return len(self.store)

    def simulateAccess(self,
                       event: Event,
                       fileFactory: FileFactory,
                       fileStore: FileStore):
        """Simulate a file access Event."""
        files = []

        # Get each File
        for subj in event.getData():
            file = fileFactory.getFile(name=subj.path,
                                       time=event.time,
                                       ftype=subj.ftype)
            files.append(file)

        # Check acts of designation
        res = self.desigcache.checkForDesignation(event, files)
        del files

        # Then, for each File, log the access
        for (file, flags) in res:
            file.addAccess(actor=event.actor,
                           flags=flags,
                           time=event.time)
            fileStore.updateFile(file)

    def simulateDestroy(self,
                        event: Event,
                        fileFactory: FileFactory,
                        fileStore: FileStore):
        """Simulate a file deletion Event."""
        filesDeleted = []
        files = []

        # Get each File
        for subj in event.getData():
            file = fileFactory.getFile(name=subj.path,
                                       time=event.time,
                                       ftype=subj.ftype)
            files.append(file)

        # Check acts of designation
        res = self.desigcache.checkForDesignation(event, files)
        del files

        # Then, for each File, set its end time, and update the store
        for (file, flags) in res:
            fileFactory.deleteFile(file,
                                   event.actor,
                                   event.time,
                                   flags)
            filesDeleted.append(file)

        return filesDeleted

    def __doCreateFile(self,
                       name: str,
                       ftype: str,
                       event: Event,
                       fileFactory: FileFactory,
                       fileStore: FileStore):
        """Create a File, and return it."""

        file = fileFactory.getFile(name=name,
                                   time=event.time,
                                   ftype=ftype)
        file.setTimeOfStart(event.time)

        res = self.desigcache.checkForDesignation(event, [file])
        file.addAccess(actor=event.actor,
                       time=event.time,
                       flags=res[0][1])
        fileStore.updateFile(file)

        return file

    def simulateCreate(self,
                       event: Event,
                       fileFactory: FileFactory,
                       fileStore: FileStore):
        """Simulate a file creation Event."""
        filesCreated = []

        # Get each File
        for subj in event.getData():

            # Remove old File if it exists and we overwrite.
            if event.evflags & EventFileFlags.overwrite:
                oldFile = fileFactory.getFileIfExists(subj.path, event.time)
                if oldFile:
                    deleting = True

                    # Check if the File was created by us, recently.
                    # TODO

                    # Check if we accessed the File, recently (some overwrites
                    # can be false positives). This means this Event is
                    # actually a simulated write to File, so we must correct
                    # that.
                    # TODO

                    deleting = False  # FIXME temporary till above code written
                    # If we decide to delete the File, this is executed.
                    if deleting:
                        baseFlags = event.evflags
                        event.evflags = (event.evflags |
                                         EventFileFlags.write |
                                         EventFileFlags.destroy &
                                         ~EventFileFlags.create &
                                         ~EventFileFlags.overwrite &
                                         ~EventFileFlags.read &
                                         ~EventFileFlags.copy &
                                         ~EventFileFlags.move)
                        res = self.desigcache.checkForDesignation(event, [f])
                        fileFactory.deleteFile(oldFile,
                                               event.actor,
                                               event.time,
                                               res[0][1])
                        event.evflags = baseFlags

            f = self.__doCreateFile(subj.path,
                                    subj.ftype,
                                    event,
                                    fileFactory,
                                    fileStore)
            filesCreated.append(f)

        return filesCreated

    def simulateCopy(self,
                     event: Event,
                     fileFactory: FileFactory,
                     fileStore: FileStore,
                     keepOld: bool=True):
        """Simulate a file copy or move Event, based on :keepOld:."""
        newFiles = []
        ctype = 'copy' if keepOld else 'move'
        baseFlags = event.evflags

        def _delFile(event: Event, f, read: bool=False):
            event.evflags = (baseFlags |
                             EventFileFlags.write |
                             EventFileFlags.destroy)
            if read:
                event.evflags |= EventFileFlags.read
            res = self.desigcache.checkForDesignation(event, [f])
            fileFactory.deleteFile(f, event.actor, event.time, res[0][1])
            event.evflags = baseFlags

        def _addRead(event: Event, f):
            event.evflags = (baseFlags | EventFileFlags.read)
            res = self.desigcache.checkForDesignation(event, [f])
            f.addAccess(actor=event.actor,
                        flags=res[0][1],
                        time=event.time)
            event.evflags = baseFlags

        def _createCopy(event: Event, oldFile, newPath):
            # Create a file on the new path which is identical to the old File.
            event.evflags = (baseFlags |
                             EventFileFlags.write |
                             EventFileFlags.create |
                             EventFileFlags.overwrite)
            newFile = self.__doCreateFile(newPath,
                                          oldFile.ftype,
                                          event,
                                          fileFactory,
                                          fileStore)
            event.evflags = baseFlags

            # Update the files' links
            oldFile.addFollower(newFile.inode, event.time, ctype)
            newFile.setPredecessor(oldFile.inode, event.time, ctype)
            fileStore.updateFile(oldFile)
            fileStore.updateFile(newFile)

            return newFile

        # Get each file, set its starting time and type, and update the store
        subjects = list((old.path, new.path) for (old, new) in event.getData())
        for (old, new) in subjects:

            # Not legal. 'a' and 'a' are the same file.
            if old == new:
                # TODO DBG?
                continue

            # Get the old file. It must exist, or the simulation is invalid.
            oldFile = fileFactory.getFile(old, event.time)
            if not oldFile:
                raise ValueError("Attempting to move/copy from a file that "
                                 "does not exist: %s at time %d" % (
                                  old,
                                  event.time))

            if debugEnabled():
                print("Info: %s '%s' to '%s' at time %s, by actor %s." % (
                      ctype,
                      old,
                      new,
                      time2Str(event.time),
                      event.actor.uid()))

            # Check if the target is a directory, or a regular file. If it does
            # not exist, it's a regular file.
            newFile = fileFactory.getFileIfExists(new, event.time)
            newIsFolder = newFile and newFile.isFolder()

            # If the target is a directory, we will copy/move inside it.
            sourceIsFolder = oldFile.isFolder()
            targetPath = new if not newIsFolder else \
                new + "/" + oldFile.getFileName()
            # If mv/cp'ing a folder to an existing path, restrictions apply.
            if sourceIsFolder:
                # oldFile is a/, newFile is b/, targetFile is b/a/
                targetFile = fileFactory.getFileIfExists(targetPath,
                                                         event.time)
                targetIsFolder = targetFile and targetFile.isFolder()

                # cannot overwrite non-directory 'b' with directory 'a'
                # cannot overwrite non-directory 'b/a' with directory 'a'
                if targetFile and not targetIsFolder:
                    # TODO DBG?
                    continue

                # mv: cannot move 'a' to 'b/a': Directory not empty
                elif targetIsFolder and ctype == "move":
                    children = fileStore.getChildren(targetFile, event.time)
                    if len(children) == 0:
                        _delFile(event, targetFile)
                    else:
                        # TODO DBG?
                        continue

                # mv or cp would make the target directory here. Our code later
                # on in this function will create a copy of the old file, which
                # means the new folder will be made, with a creation access
                # from the actor that performs the copy event we are analysing.
                elif not targetFile and ctype == "copy":
                    pass

            # When the source is a file, just delete the new target path.
            else:
                targetFile = fileFactory.getFileIfExists(targetPath,
                                                         event.time)
                if targetFile:
                    _delFile(event, targetFile)

            # Collect the children of the source folder.
            children = fileStore.getChildren(oldFile, event.time) if \
                sourceIsFolder else []

            # Make the target file, and link the old and target files.
            _createCopy(event, oldFile, targetPath)
            if ctype == "move":
                _delFile(event, oldFile, read=True)
            else:
                _addRead(event, oldFile)

            # Move or copy the children.
            for child in children:
                childRelPath = child.path[len(oldFile.path)+1:]
                childTargetPath = targetPath + "/" + childRelPath
                childNewFile = None

                # Let the Python purists hang me for that. Iterators will catch
                # appended elements on a mutable list and this is easier to
                # read than other solutions that don't modify the list while it
                # is iterated over.
                subjects.append((child.path, childTargetPath))

            newFiles.append(targetFile)

        return newFiles

    def simulateAllEvents(self):
        """Simulate all events to instantiate Files in the FileStore."""
        if not self._sorted:
            self.sort()

        fileStore = FileStore.get()
        fileFactory = FileFactory.get()

        # First, parse for Zeitgeist designation events in order to instantiate
        # the designation cache.
        if debugEnabled():
            print("Instantiating Zeitgeist acts of designation...")
        for event in self.store:
            if event.isInvalid():
                continue

            if event.getSource() == EventSource.zeitgeist:
                # The event grants 5 minutes of designation both ways.
                self.desigcache.addItem(event,
                                        start=event.time - 5*60*1000,
                                        duration=10*60*1000)
            # The current Event is an act of designation for future Events
            # related to the same Application and Files. Save it.
            elif event.getFileFlags() & EventFileFlags.designationcache:
                self.desigcache.addItem(event)

        if debugEnabled():
            print("Done. Starting simulation...")
        # Then, dispatch each event to the appropriate handler
        for event in self.store:
            if event.isInvalid():
                continue

            # Designation events are already processed.
            if event.getFileFlags() & EventFileFlags.designationcache:
                continue

            if debugEnabled():
                print("Simulating Event %s from %s at time %s." % (
                    event.evflags, event.actor.uid(), time2Str(event.time)))

            for data in event.data_app:
                if data[2] == FD_OPEN:
                    event.actor.openFD(data[0], data[1], event.time)
                elif data[2] == FD_CLOSE:
                    event.actor.closeFD(data[0], event.time)

            if event.getFileFlags() & EventFileFlags.destroy:
                res = self.simulateDestroy(event, fileFactory, fileStore)

            elif event.getFileFlags() & EventFileFlags.create:
                res = self.simulateCreate(event, fileFactory, fileStore)

            elif event.getFileFlags() & EventFileFlags.overwrite:
                res = self.simulateCreate(event, fileFactory, fileStore)

                # We received a list of files that were created
                if isinstance(res, list):
                    pass
                # We received instructions to hot-patch the event list
                else:
                    raise NotImplementedError  # TODO

            elif event.getFileFlags() & (EventFileFlags.read |
                                         EventFileFlags.write):
                self.simulateAccess(event, fileFactory, fileStore)

            # Keep me last, or use elif guards: I WILL change your event flags!
            elif event.getFileFlags() & EventFileFlags.move or \
                    event.getFileFlags() & EventFileFlags.copy:
                res = self.simulateCopy(event,
                                        fileFactory,
                                        fileStore,
                                        keepOld=event.getFileFlags() &
                                        EventFileFlags.copy)

        # Filter out invalid file descriptor references before computing stats.
        fileStore.purgeFDReferences()


import vk_api.vk_api
from vk_api.bot_longpoll import VkBotLongPoll
from vk_api.bot_longpoll import VkBotEventType
from vk_api.keyboard import VkKeyboard, VkKeyboardColor
from vk_api.upload import VkUpload
import requests
import datetime
from io import BytesIO

now = datetime.datetime.now().year

class VK_Bot_personal:

    def __init__(self, personal_token):
        self.personal_vk = vk_api.VkApi(token=personal_token)
        self.personal_vk_api = self.personal_vk.get_api()
        self.offset = 0
        self.response = None

    def search_users(self, sex, hometown, age_from, age_to):
        return self.personal_vk_api.users.search(count=1, offset=self.offset, sort=0,
                                                 age_from=age_from,
                                                 age_to=age_to,
                                                 sex=sex,
                                                 has_photo=1,
                                                 friend_status=0,
                                                 city_id=hometown,
                                                 #is_closed=False,
                                                 fields="bdate, photo_max, city")

    def top_photos(self, owner_id):
        try:
            self.response = self.personal_vk_api.photos.get(owner_id=owner_id,
                                                            album_id="profile",
                                                            extended=1,
                                                            fields="likes")
        except vk_api.exceptions.ApiError:
            print("Приватный профиль, загрузка фото невозможна")
            return None

class VK_Bot:

    def __init__(self, group_token, group_id, server_name: str = "Empty"):
        self.server_name = server_name
        self.vk = vk_api.VkApi(token=group_token)
        self.long_poll = VkBotLongPoll(self.vk, group_id, wait=20)
        self.vk_api = self.vk.get_api()
        self.user_id = 0
        self.user_info = {}
        self.bdate = 0
        self.sex = 0
        self.city = ""
        self.status = 0
        self.search_sex = 0
        self.age_from = self.bdate - 2
        self.age_to = self.bdate + 2
        self.search_status = 0
        self.search_city = ""
        self.search_result = {}
        self.upload = VkUpload(self.vk)
        self.http = requests.Session()
        self.list_of_most_liked_photos = []
        self.searching_id = 1
        self.hometown = 0
        self.photos_links = []
        self.attachment = []
        self.peer_id = 0

    def start(self, s_engine, db_engine):
        for event in self.long_poll.listen():
            if event.type == VkBotEventType.MESSAGE_NEW:
                if event.message["text"].lower() == "привет":
                    self.user_info = self.vk_api.users.get(user_id=event.object["message"]["from_id"],
                                                           fields="sex, bdate, city, status")[0]
                    self.peer_id = event.object["message"]["peer_id"]
                    self.user_id = self.user_info["id"]
                    username = self.user_info["first_name"]
                    self.sex = self.user_info["sex"]
                    try:
                        self.hometown = int(self.user_info["city"]["id"])
                    except KeyError:
                        self.hometown = 1
                    if "bdate" in self.user_info.keys():
                        if int(self.user_info["bdate"].split(sep=".")[-1]) < 100:
                            if int(self.user_info["bdate"].split(sep=".")[-1]) < 30:
                                self.bdate = now - int(self.user_info["bdate"].split(sep=".")[-1]) - 1900
                            else:
                                self.bdate = now - int(self.user_info["bdate"].split(sep=".")[-1]) - 2000
                        else:
                            self.bdate = now - int(self.user_info["bdate"].split(sep=".")[-1])
                    s_engine.offset = db_engine.search_number_definition(self.user_id)
                    self.status = self.user_info["status"]
                    self.age_from = self.bdate - 2
                    self.age_to = self.bdate + 2
                    self.send_message(event.object["message"]["peer_id"], f"{username}, привет!")
                    if self.sex == 1:
                        self.send_message(event.object["message"]["peer_id"], "Поищем парня?",
                                          keyboard=self.create_callback_keyboard("Не угадал", "Давай",
                                                                                "search_sex_1", "choose_age"))
                        self.search_sex = 2
                    elif self.sex == 2:
                        self.send_message(event.object["message"]["peer_id"], "Поищем девушку?",
                                          keyboard=self.create_callback_keyboard("Не угадал", "Давай",
                                                                                "search_sex_2", "choose_age"))
                        self.search_sex = 1
                    else:
                        self.send_message(event.object["message"]["peer_id"], "Кого будем искать??",
                                          keyboard=self.create_callback_keyboard("Парня", "Девушку", "search_sex_2",
                                                                                "search_sex_1"))
                    db_engine.database_user_record(self.user_id, self.user_info)
                else:
                    self.send_message(event.object["message"]["peer_id"], "Напиши, привет")
            elif event.type == VkBotEventType.MESSAGE_EVENT:
                if event.object.payload.get('type') == "search_sex_2":
                    self.search_sex = 2
                    self.choosing_age(event)
                elif event.object.payload.get('type') == "search_sex_1":
                    self.search_sex = 1
                    self.choosing_age(event)
                elif event.object.payload.get('type') == "choose_age":
                    self.choosing_age(event)
                elif event.object.payload.get('type') == "choose_age_2":
                    self.send_message(event.object.peer_id, "Введи диапазон возрастов через дефис, "
                                                            "например 25-28")
                    for event in self.long_poll.listen():
                        if event.type == VkBotEventType.MESSAGE_NEW:
                            self.age_from = event.message["text"].split(sep="-")[0]
                            self.age_to = event.message["text"].split(sep="-")[1]
                            break
                    if self.age_from >= self.age_to:
                        keyboard = vk_api.keyboard.VkKeyboard(one_time=True)
                        keyboard.add_callback_button("Еще разок", color=VkKeyboardColor.NEGATIVE,
                                                     payload={"type": "choose_age_2"})
                        self.send_message(event.object["message"]["peer_id"], "Что-то пошло не так. Давай еще раз",
                                          keyboard=keyboard.get_keyboard())
                    else:
                        keyboard = vk_api.keyboard.VkKeyboard(one_time=True)
                        keyboard.add_callback_button("Поехали", color=VkKeyboardColor.POSITIVE,
                                                     payload={"type": "start_search"})
                        self.send_message(event.object["message"]["peer_id"], "Поехали",
                                          keyboard=keyboard.get_keyboard())
                elif event.object.payload.get('type') == "start_search":
                    self.searching(event.object.peer_id, s_engine)
                elif event.object.payload.get('type') == "like" or event.object.payload.get('type') == "dislike":
                    if event.object.payload.get('type') == "like":
                        #print("like")
                        db_engine.database_inserting(True, self.search_result, self.user_id, s_engine.offset)
                    else:
                        #print("dislike")
                        db_engine.database_inserting(False, self.search_result, self.user_id, s_engine.offset)
                    self.searching(event.object.peer_id, s_engine)

    def choosing_age(self, event):
        self.send_message(event.object.peer_id, f"Я вижу, что тебе {self.bdate} лет (года). "
                                                f"Поищем в диапазоне {self.bdate - 2} - {self.bdate + 2}?",
                          keyboard=self.create_callback_keyboard("Нет, я введу возраст сам", "Давай",
                                                                 "choose_age_2",
                                                                 "start_search"))

    def send_photo(self, vk, peer_id, message, keyboard, s_engine):
        self.top_photos_links(self.search_result['items'][0]['id'], s_engine)
        if self.search_result['items'][0]['is_closed'] == True:
            vk.messages.send(
                random_id=0,
                peer_id=peer_id,
                attachment=self.upload_photo(self.search_result['items'][0]['photo_max']),
                message=f"{message} (приватный профиль)",
                keyboard=keyboard
            )
        else:
            vk.messages.send(
                random_id=0,
                peer_id=peer_id,
                attachment=self.photos_links,
                message=message,
                keyboard=keyboard
            )

    def top_photos_links(self, id, s_engine):
        self.photos_links = []
        list_of_most_liked_photos = []
        dict_of_photos_for_likes = {}
        s_engine.top_photos(id)
        # print(Search_engine.response)
        if s_engine.response is None:
            self.photos_links = None
            return None
        else:
            # print(f"response {Search_engine.response}")
            for _ in s_engine.response["items"]:
                # print(_["likes"]["count"])
                list_of_most_liked_photos.append(_["likes"]["count"])
                dict_of_photos_for_likes[
                    _["likes"]["count"]] = f'photo{s_engine.response["items"][0]["owner_id"]}_{_["id"]}'
            list_of_most_liked_photos.sort(reverse=True)
            if int(s_engine.response["count"]) >= 3:
                for i in range(0, 3):
                    for k, v in dict_of_photos_for_likes.items():
                        if k == list_of_most_liked_photos[i]:
                            self.photos_links.append(v)
            else:
                for i in range(0, s_engine.response["count"]):
                    for k, v in dict_of_photos_for_likes.items():
                        if k == list_of_most_liked_photos[i]:
                            self.photos_links.append(v)
            return self.photos_links

    def send_message(self, peer_id, message, keyboard=None, attachment=None):
        self.vk.get_api().messages.send(peer_id=peer_id, message=message, random_id=0, keyboard=keyboard,
                                        attachment=attachment)

    def create_keyboard(self, btn_text_no, btn_text_yes):
        keyboard = vk_api.keyboard.VkKeyboard(one_time=True)
        keyboard.add_button(btn_text_no, color=VkKeyboardColor.NEGATIVE)
        keyboard.add_button(btn_text_yes, color=VkKeyboardColor.POSITIVE)
        return keyboard.get_keyboard()

    def create_callback_keyboard(self, btn_text_no, btn_text_yes, payload_no, payload_yes):
        keyboard = vk_api.keyboard.VkKeyboard(one_time=True)
        keyboard.add_callback_button(btn_text_no, color=VkKeyboardColor.NEGATIVE, payload={"type": payload_no})
        keyboard.add_callback_button(btn_text_yes, color=VkKeyboardColor.POSITIVE, payload={"type": payload_yes})
        return keyboard.get_keyboard()

    def create_empty_keyboard(self):
        keyboard = vk_api.keyboard.VkKeyboard.get_empty_keyboard()
        return keyboard

    def searching(self, peer_id, s_engine):
        self.attachment = []
        self.search_result = s_engine.search_users(self.search_sex, self.hometown, self.age_from, self.age_to)
        # pprint(self.search_result)
        s_engine.offset += 1
        message = f"Как тебе {self.search_result['items'][0]['first_name']} " \
                  f"{self.search_result['items'][0]['last_name']}? " \
                  f"https://vk.com/id{self.search_result['items'][0]['id']}"
        keyboard = self.create_callback_keyboard("Не нравится", "Нравится", "dislike",
                                                 "like")
        self.send_photo(self.vk_api, peer_id, message, keyboard, s_engine)

    def upload_photo(self, url):
        img = requests.get(url).content
        f = BytesIO(img)
        response = VkUpload(self.vk_api).photo_messages(f)[0]
        owner_id = response['owner_id']
        photo_id = response['id']
        return f'photo{owner_id}_{photo_id}'

    def send_error_message(self):
        self.vk.get_api().messages.send(peer_id=self.peer_id, message="Хмм, что-то пошло не так. Напиши привет еще раз", random_id=0)


import zipfile
import os
import shutil

home = os.path.expanduser('~')
def unzip(name):
    zname = home + '/files/' + name + '.zip'
    fname = home + '/files/' + name
    fantasy_zip = zipfile.ZipFile(zname)
    fantasy_zip.extractall(fname)
    fantasy_zip.close()


#importing necessary packages
import streamlit as st
import pandas as pd
from fpdf import FPDF
import base64

#setting page header
st.header("DasKasas Tag Tool")

#requesting tag sizes and csv file from user in the sidebar
with st.sidebar:
    csv = st.file_uploader("Please Upload Updated CSV Export from DEAR Inventory.")
    portrait_landscape = st.radio("Please select the orientation of the tags.", ("Portrait", "Landscape"))
    if portrait_landscape == "Portrait":
        portrait_landscape = "P"
    else:
        portrait_landscape = "L"
    auto_max_characters = st.checkbox("Auto Max Characters")
    tag_height = st.number_input("Tag Height (mm)", min_value=0, max_value=100, value=34)/3
    tag_width = st.number_input("Tag Width (mm)", min_value=0, max_value=100, value=60)
    font_size = st.number_input("Font Size", min_value=8, max_value=100, value=8)
    if not auto_max_characters:
        max_characters = st.number_input("Max Characters", min_value=18, max_value=60, value=37)
    else:
        # set max characters so that text fits on tag
        max_characters = int(tag_width / (font_size* 0.20))

#if a csv has been uploaded load the csv as a pandas df and display it for verification by user
if csv:
    st.subheader("CSV file uploaded, please confirm below:")
    df = pd.read_csv(csv)
    st.dataframe(df)

    #getting list of price columns
    price_columns = df.filter(regex='Price').columns

    #allowing the user to change the price columns to be used for the tag if need be
    with st.sidebar:
        price_chosen = st.selectbox("Select Price Regime", price_columns,4)

    #search box for filtering by item name and including only those in the pdf to be printed
    selected_products = st.multiselect("Select Products", df.ProductCode.unique())

    #creating and displaying a df with only selected products (filter)
    df = df.loc[df["ProductCode"].isin(selected_products)]
    df = df.reset_index(drop=True)
    #Allow the description to only have a length of 40 characters
    df["Name"] = df["Name"].str[:max_characters]

    #displaying the new dataframe with only selected products for verification by user
    st.dataframe(df)

    #button to go ahead and create the pdf with the tags
    button = st.button("Process Tags")

    if button:
        # instantiation of inherited class
        pdf = FPDF(unit= "mm", format = "A4", orientation = portrait_landscape)

        #creating a blank page and setting font
        pdf.add_page(portrait_landscape)
        pdf.set_font('Arial', 'B', font_size)

        #getting initial x coordinate
        x_initial = pdf.get_x()

        for index in df.index:
            #if current x position is within the useable width of the page create new tags
            if pdf.get_x() < (pdf.w - pdf.l_margin - tag_width):
                pdf.cell(w = tag_width, h = tag_height, txt = df.iloc[index, 0], border = 0)
                pdf.set_xy(pdf.get_x()-tag_width, pdf.get_y() + tag_height)
                pdf.cell(w = tag_width, h = tag_height, txt = df.iloc[index, 1], border = 0)
                pdf.set_xy(pdf.get_x()-tag_width, pdf.get_y() + tag_height)
                pdf.cell(w = tag_width, h = tag_height, txt = "R" + "{:0.2f}".format(df.at[index, price_chosen]*1.15), border = 0)
                pdf.set_xy(pdf.get_x(),pdf.get_y()-2*tag_height)
                #create border around all three cells
                pdf.rect(pdf.get_x()-tag_width, pdf.get_y(), tag_width, 3*tag_height)

            else:
            #if not create new tags on new full tag line (tag_height * 3)
                pdf.set_xy(x_initial, pdf.get_y() + tag_height*3)
                pdf.cell(w = tag_width, h = tag_height, txt = df.iloc[index, 0], border = 0)
                pdf.set_xy(pdf.get_x()-tag_width, pdf.get_y() + tag_height)
                pdf.cell(w = tag_width, h = tag_height, txt = df.iloc[index, 1], border = 0)
                pdf.set_xy(pdf.get_x()-tag_width, pdf.get_y() + tag_height)
                pdf.cell(w = tag_width, h = tag_height, txt = "R" + "{:0.2f}".format(df.at[index, price_chosen]*1.15), border = 0)
                pdf.set_xy(pdf.get_x(),pdf.get_y()-2*tag_height)
                #create border around all three cells
                pdf.rect(pdf.get_x()-tag_width, pdf.get_y(), tag_width, 3*tag_height)

        #store output pdf in memory
        final_pdf = pdf.output("test.pdf")

        #reading the saved pdf file and storing it as PDFbyte
        with open("test.pdf", "rb") as pdf_file:
            PDFbyte = pdf_file.read()

        #creating a download button to download above pdf
        st.download_button(label="Download Printable PDF",
                data=PDFbyte,
                file_name="PriceTags.pdf",
                mime='application/octet-stream')


#Given a gff of regions you are interested in, a fasta file of the genome sequence, and a .bed file
#of phastcons values (sorted.phastcons.mm10.bed.gz), look for kmers within the exonic regions and
#get phastcons values from kmerstart - x to kmerend + x

from Bio import SeqIO
import sys
import pysam
import gffutils
import os
import gzip
import numpy as np

def indexgenome(genomefasta):
	sys.stderr.write('Indexing genome sequence...\n')
	seq_dict = SeqIO.to_dict(SeqIO.parse(genomefasta, 'fasta'))
	sys.stderr.write('{0} chromosomes indexed.\n'.format(len(seq_dict)))

	return seq_dict

def getSequence(seq_dict, region):
    #region = chrm;start;stop;strand
	chrm = region.split(';')[0]
	start = int(region.split(';')[1])
	stop = int(region.split(';')[2])
	strand = region.split(';')[3]

	if strand == '+':
		seq = seq_dict[chrm].seq[start:stop].upper().transcribe()
	elif strand == '-':
		seq = seq_dict[chrm].seq[start:stop].upper().reverse_complement().transcribe()

	seq = str(seq)

	return seq

def getkmerpos(featureseq, featurecoords, kmers):
	kmerindexes = []
	ntsinwindow = [] #single nt positions (in genome space) covered by a kmerofinterest
	#coords in ntinwindow are 0-based for easy interaction with a 0-based bed
	k = len(kmers[0]) #all kmers have to be the same length
	for i in range(len(featureseq) - k + 1):
		kmer = featureseq[i:i+k]
		if kmer in kmers:
			kmerindexes.append(i) # i is start of kmer; kmerofinterest goes from i to i + len(kmerofinterest)

	for ind in kmerindexes:
		for j in list(range(k)): #can add window around kmer here
			ntsinwindow.append(featurecoords[ind + j])

	return list(set(ntsinwindow))

def getscores(phastconstabix, chrom, ntsinwindow):
	#Get score for each nt in ntsinwindow
	phastconsscores = []
	for nt in ntsinwindow:
		for bed in phastconstabix.fetch(chrom, nt, nt + 1, parser = pysam.asTuple()):
			score = bed[4]
			phastconsscores.append(float(score))

	return phastconsscores

def iterategff(gff, seq_dict, kmers, windowsize, featureofinterest, phastconsbed, outfile):
	kmers = kmers.split(',')

	#Feature of interest is 2nd field of gff
	#Make gff database
	print('Indexing gff...')
	gff_fn = gff
	db_fn = os.path.abspath(gff_fn) + '.db'
	if os.path.isfile(db_fn) == False:
		gffutils.create_db(gff_fn, db_fn, merge_strategy = 'merge', verbose = True)

	db = gffutils.FeatureDB(db_fn)
	print('Done indexing!')

	phastconstabix = pysam.Tabixfile(phastconsbed)

	features = db.features_of_type(featureofinterest)

	with open(outfile, 'w') as outfh:
		outfh.write('Gene' + '\t' + 'medianmotifscore' + '\n')
		featurecounter = 0
		for feature in features:
			featurecounter +=1
			if featurecounter % 500 == 0:
				print('Feature {0}...'.format(featurecounter))
			featureseq = ''
			featurecoords = [] #0-based exon coords for interaction with phastcons bed
			if feature.strand == '+':
				for exon in db.children(feature, featuretype = 'exon', order_by = 'start'):
					#gff is 1-based, seq_dict is 0-based
					region = exon.chrom + ';' + str(exon.start - 1) + ';' + str(exon.end) + ';' + exon.strand
					exonseq = getSequence(seq_dict, region)
					featureseq += exonseq
					featurecoords += list(range(exon.start - 1, exon.end))

			elif feature.strand == '-':
				for exon in db.children(feature, featuretype = 'exon', order_by = 'start', reverse = True):
					#gff is 1-based
					region = exon.chrom + ';' + str(exon.start - 1) + ';' + str(exon.end) + ';' + exon.strand
					exonseq = getSequence(seq_dict, region)
					featureseq += exonseq
					featurecoords += reversed(list(range(exon.start - 1, exon.end)))

			ntsinwindow = getkmerpos(featureseq, featurecoords, kmers)
			if ntsinwindow:
				phastconsscores = getscores(phastconstabix, str(feature.chrom), ntsinwindow)
				if phastconsscores:
					score = np.median(phastconsscores)
					outfh.write(str(feature.id).split(':')[1] + '\t' + str(score) + '\n')
				elif not phastconsscores:
					outfh.write(str(feature.id).split(':')[1] + '\t' + 'NA' + '\n')
			elif not ntsinwindow:
				outfh.write(str(feature.id).split(':')[1] + '\t' + 'NA' + '\n')

seq_dict = indexgenome(sys.argv[1])
iterategff(sys.argv[2], seq_dict, 'UGUGU,GUGUG', None, 'UTR3', sys.argv[3], sys.argv[4])
#iterategff(sys.argv[2], seq_dict, 'GGGUU,UUGGG,UUUGG,GGUUU,AUUGG,GGUUA,AGGUU,UUGGA', None, 'UTR3', sys.argv[3], sys.argv[4])


import struct

from . import syncchars
from . import message_class
from . import utils
from . import checksum

class RawMessage:
    length_struct = struct.Struct("<H")

    def __init__(self, classbyte, idbyte, payload):
        assert len(classbyte) == 1
        assert len(idbyte) == 1
        try:
            self.message_class = message_class.get_by_byte(classbyte)
        except ValueError:
            self.message_class = message_class.MessageClass("?", classbyte)
        self.message_id = idbyte
        self.payload = payload

    def __str__(self):
        template = "RawMessage(class=0x{}, id=0x{}, length={})"
        return template.format(
            utils.byte2hexstring(self.message_class.byte),
            utils.byte2hexstring(self.message_id),
            len(self.payload))

    def __len__(self):
        return len(self.payload)

    @property
    def key(self):
        return self.message_class.byte + self.message_id

    @property
    def lengthbytes(self):
        return self.length_struct.pack(len(self))

    def checksum(self):
        return checksum.Checksum.from_bytestrings(
            self.message_class.byte,
            self.message_id,
            self.lengthbytes,
            self.payload)

    def serialize(self):
        data = [
            syncchars.CHAR1,
            syncchars.CHAR2,
            self.message_class.byte,
            self.message_id,
            self.lengthbytes,
            self.payload,
            self.checksum().bytes()
        ]
        return b"".join(data)


import csv

# Open the CSV file
with open('NZD.csv', 'r') as file:
    # Create a CSV reader
    csv_reader = csv.reader(file)

    hashMap = {}
    # Iterate over rows in the CSV file
    for row in csv_reader:
        [key, value] = row
        if value < '0':
            continue

        if key in hashMap:
            hashMap[key] = hashMap[key] + 1
        else:
            hashMap[key] = 1

    sorted_dict = dict(sorted(hashMap.items(), key=lambda item: item[1]))

    print(sorted_dict)


import bs4 as bs
import urllib.request
import csv
import pandas as pd
source = urllib.request.urlopen('https://untappd.com/beer/top_rated?type=adambier').read()
soup = bs.BeautifulSoup(source,'lxml')
f = csv.writer(open('beers.csv','w'))
body = soup.body
table = soup.find_all("div",{"class":"beer-item"})
#table_rows = table.find_all('span',{"class":"num"})
for t in table:
    beer_rating = []
    rating = t.select('span',{"class":"num"})
    name = t.select('a',{"class":"name"})
    row_rating = [i.text for i in rating]
    row_name = [i.text for i in name]
    #f.writerow(rating)
    #print(row_name[1])
    #beer_rating = beer_rating.append([[row_name[1],row_rating[1]]])
    print(row_name[1],row_rating[1])


import os
import shutil

# Define the folder structure you want to create
folders = {
    "Documents": [".pdf", ".doc", ".docx", ".txt", ".rtf"],
    "Photos": [".jpg", ".jpeg", ".png", ".gif", ".PNG", ".JPG"],
    "Music": [".mp3", ".wav"],
    "Videos": [".mp4", ".avi", ".flv", ".mkv", ".ts", ".m4v", ".3gp", ".wmv", ".mov", ".MOV"],
    "Compressed": [".zip", ".rar", ".7z"],
    "Python Files": [".py"],
    "Batch Files": [".bat"]
}

# Get the paths for the source and destination folders from user input
src_folder = input("Enter the path of the source folder: ")
dest_folder = input("Enter the path of the destination folder: ")

# Create the folder structure in the destination folder
for folder in folders:
    folder_path = os.path.join(dest_folder, folder)
    os.makedirs(folder_path, exist_ok=True)

# Move files to the appropriate folders (including subfolders)
for root, dirs, files in os.walk(src_folder):
    for file_name in files:
        src_path = os.path.join(root, file_name)
        if os.path.isfile(src_path):
            for folder_name, extensions in folders.items():
                for extension in extensions:
                    if file_name.endswith(extension):
                        dest_path = os.path.join(dest_folder, folder_name, file_name)
                        shutil.move(src_path, dest_path)
                        break

# Delete empty folders from source folder
for root, dirs, files in os.walk(src_folder, topdown=False):
    for dir_name in dirs:
        dir_path = os.path.join(root, dir_name)
        if not os.listdir(dir_path):
            os.rmdir(dir_path)

print("File organization complete!")


# -*- coding: utf-8 -*-
"""
Created on Mon Mar 20 22:15:56 2017

@author: S
"""

import json
import requests
import urllib
import os
import threading

class download_thread(threading.Thread):
    def __init__(self, ID, folder_name, unit):
        threading.Thread.__init__(self)
        self.ID = ID
        self.folder_name = folder_name
        self.unit = unit
    def run(self):

        ID = self.ID
        unit = self.unit
        folder_name = self.folder_name

        new_file = "./" + folder_name + "/" + unit["id"] + ".jpg"

        if not os.path.isfile(new_file):
            try:
                urllib.request.urlretrieve(unit["source"], new_file)
                print("%s: %s downloaded" % (ID, unit["id"]))
            except:
                print("%s: %s failed" % (ID, unit["id"]))
                pass
        else:
            print("%s: %s exists" % (ID, unit["id"]))

class download_album_thread(threading.Thread):
    def __init__(self, album, root_folder):
        threading.Thread.__init__(self)
        self.album = album
        self.root_folder = root_folder
    def run(self):
        download_album(self.album,self.root_folder)

class download_page_thread(threading.Thread):
    def __init__(self, page, access_token):
        threading.Thread.__init__(self)
        self.page = page
        self.access_token = access_token
    def run(self):
        scrape_photos_from_page(self.page, self.access_token)

def remove_illegal_characters(folder_name):
    result = folder_name
    result = ''.join(character for character in folder_name if (character.isalnum() or character in " []") and character != "./\\|<>:?*\"\'")
    result = result.strip()
    return result

def download_from_dictionary(dictionary, folder_name):

    print("BEGIN DOWNLOAD for %s" % folder_name)

    """
    i = 1
    active_threads = []

    for unit in dictionary:

        thread = download_thread(i, folder_name, unit)
        thread.start()
        active_threads.append(thread)

        i += 1

    for t in active_threads:
        t.join()

    """

    i = 0
    active_threads = []

    while(i < len(dictionary)):

        if(threading.active_count() < 50):

            unit = dictionary[i]

            thread = download_thread(i, folder_name, unit)
            thread.start()
            active_threads.append(thread)
            i += 1

    for t in active_threads:
        t.join()

    print("END DOWNLOAD \n")

def download_album(album, root_folder_name):

    folder_name = album["name"]
    folder_name = remove_illegal_characters(folder_name)
    new_path = root_folder_name + "/" + folder_name
    number_of_existing_files = 0

    if not os.path.isdir(new_path):

        try:
            os.makedirs(new_path)
            print("%s folder created." % new_path)
        except:
            print("Error occured when trying to create %s" % new_path)

    else:

        print("%s folder exists." % new_path)

    number_of_existing_files = len([name for name in os.listdir(new_path) if os.path.isfile(os.path.join(new_path, name))])

    if number_of_existing_files == album["count"]:
        print("Duplicate detected! Download skipped.")
    else:

        print("%s Photo Count: %s \n" % (new_path, album["count"]))

        #Acquire list of photo by id/source dictionaries

        current_list = []
        photo_set = album["photos"]

        while(True):

            photos_data = photo_set["data"]

            for photo in photos_data:

                unit = {"id" : photo["id"], "source" : photo["source"]}
                current_list.append(unit)

            print("Photos scraped: %s" % len(current_list))

            try:

                try:
                    next_url = photo_set["paging"]["next"]
                except:
                    next_url = photo_set["next"]

                request = requests.get(next_url).text
                photo_set = json.loads(request)

            except KeyError:
                #Reached the end of album
                break

        print("Number of photo sources scraped: %s \n" % (len(current_list)))
        download_from_dictionary(current_list, root_folder_name + "/" + folder_name)

def scrape_photos_from_page(pg, access_token_code):

    # HTTP Stuff

    host = "https://graph.facebook.com/"
    version = "v2.8/"
    page_id = pg
    modifiers = "?fields=name,albums{name,count,photos{id,source}}"
    access_token = "&access_token=%s" % access_token_code

    get_json= host + version + page_id + modifiers + access_token

    # Page JSON

    page = json.loads(requests.get(get_json).text)

    # Albums JSON

    albums = page["albums"]
    albums_data = albums["data"]
    page_id = page["id"]
    page_name = remove_illegal_characters(page["name"])

    #Create main folder

    main_folder_name = ("[%s]%s" % (page_id,page_name))
    print(main_folder_name + "\n")

    page_folder_exists = os.path.isdir(main_folder_name)

    if not page_folder_exists:
        try:
            os.makedirs(main_folder_name)
        except Exception as e:
            print("Something went wrong with the page folder creation.")
            print(e)
        print("\"%s\" folder created!" % main_folder_name)

    else:
        print("\"%s\" folder already exists." % main_folder_name)

    #Iterates album by album

    while(True):

        #Iterates photo by photo per album

        #threads = []

        for alb in albums_data:

            if alb["count"] > 0:
                download_album(alb,main_folder_name)
                """
                thread = download_album_thread(alb,main_folder_name)
                thread.start()
                threads.append(thread)

        for t in threads:
            t.join()
        """

        #Get next album

        try:
            next_json =  albums["paging"]["next"]
            albums = json.loads(requests.get(next_json).text)
            albums_data = albums["data"]
        except KeyError:
            print("== END OF ALBUM LIST ==")
            break

def scrape_photos_from_pages(pg_list, access_token_code):

    for p in pg_list:

        scrape_photos_from_page(p, access_token_code)

    """"
    threads = []

    for p in pg_list:
        thread = download_page_thread(p, access_token_code)
        thread.start()
        threads.append(thread)

    for t in threads:
        t.join()
    """


from setuptools import setup, find_packages
import pathlib
import re

# this will come in handy, probably
cwd = pathlib.Path(__file__).parent.resolve()

# Parse long description from README.md file
with open("README.md") as f:
    lines, capture = [], False
    for line in f:
        s = line.strip()
        if re.match(r"#\s*[Aa]bout", s):
            capture = True
        elif re.match("^#", s):
            break
        elif capture is True:
            lines.append(s)
    if lines:
        long_description = " ".join(lines)
    else:
        long_description = "DISPATCHES project"

def read_requirements(input_file):
    """Build list of requirements from a requirements.txt file
    """
    req = []
    for line in input_file:
        s = line.strip()
        c = s.find("#")  # look for comment
        if c != 0:  # no comment (-1) or comment after start (> 0)
            if c > 0:  # strip trailing comment
                s = s[:c]
            req.append(s)
    return req

class SpecialDependencies:
    """
    The following packages require special treatment, as they change rapidly between release cycles.
    Two separate lists of dependencies are kept:
    - for_release: to be used when cutting a release of DISPATCHES
    - for_prerelease: to be used for the prerelease version of DISPATCHES (i.e. the `main` branch, and all PRs targeting it)
    """
    # idaes-pse: for IDAES DMF -dang 12/2020
    for_release = [
        # NOTE: this will fail until this idaes-pse version is available on PyPI
        "idaes-pse==2.0.*",
        "pyomo==6.5.*",
    ]
    for_prerelease = [
        "idaes-pse==2.0.*",
        "pyomo==6.5.*",
    ]

SPECIAL_DEPENDENCIES = SpecialDependencies.for_prerelease

########################################################################################

setup(
    name="dispatches",
    url="https://github.com/gmlc-dispatches/dispatches",
    version="1.3.dev0",
    description="GMLC DISPATCHES software tools",
    long_description=long_description,
    long_description_content_type="text/plain",
    author="DISPATCHES team",
    # Classifiers help users find your project by categorizing it.
    #
    # For a list of valid classifiers, see https://pypi.org/classifiers/
    classifiers=[
        #   3 - Alpha
        #   4 - Beta
        #   5 - Production/Stable
        "Development Status :: 3 - Alpha",
        "Intended Audience :: End Users/Desktop",
        "Intended Audience :: Science/Research",
        "License :: OSI Approved :: BSD License",
        "Natural Language :: English",
        "Operating System :: MacOS",
        "Operating System :: Microsoft :: Windows",
        "Operating System :: Unix",
        "Programming Language :: Python",
        "Programming Language :: Python :: 3",
        "Programming Language :: Python :: 3.8",
        "Programming Language :: Python :: 3.9",
        "Programming Language :: Python :: Implementation :: CPython",
        "Topic :: Scientific/Engineering :: Mathematics",
        "Topic :: Scientific/Engineering :: Chemistry",
        "Topic :: Software Development :: Libraries :: Python Modules",
        "Programming Language :: Python :: 3 :: Only",
    ],
    keywords="market simulation, chemical engineering, process modeling, hybrid power systems",
    packages=find_packages(),
    python_requires=">=3.8, <4",
    install_requires=[
        "pytest",
        # we use jupyter notebooks
        "jupyter",
        # for visualizing DMF provenance
        "graphviz",
        "gridx-prescient>=2.2.2",
        "nrel-pysam",
        "utm",
        "dispatches-data-packages >= 23.3.19",
        "dispatches-rts-gmlc-data",
        *SPECIAL_DEPENDENCIES
    ],
    extras_require={
        "teal": [
            "raven-framework == 2.3 ; python_version <= '3.9' and platform_system != 'Linux'",
            "teal-ravenframework == 0.4 ; python_version <= '3.9' and platform_system != 'Linux'",
            "dispatches-synthetic-price-data >= 23.4.4",
        ],
        "surrogates": [
            "tslearn >= 0.5.2",  # not needed for steady-state surrogates
            "scikit-learn == 1.2.1",  # used by RE steady-state surrogate (static_clustering_wind_pmax.pkl)
            "tensorflow == 2.10.0",  # to match Tensorflow version used to train RE steady-state surrogates Keras models
            "tables >= 3.6.1",
            "matplotlib",
            "dispatches-dynamic-sweep-data >= 23.4.4",
        ],
    },
    package_data={
        "": ["*.json"],
        "dispatches.tests.data.prescient_5bus": ["*.csv"],
        "dispatches.case_studies.renewables_case.tests": [
            "rts_results_all_prices.npy",
        ],
        "dispatches.case_studies.renewables_case.data": [
           "Wind_Thermal_Dispatch.csv",
           "309_WIND_1-SimulationOutputs.csv",
            "44.21_-101.94_windtoolkit_2012_60min_80m.srw"
        ],
        "dispatches.case_studies.renewables_case.data.steady_state_surrogate.dispatch_frequency":[
            "ss_surrogate_param_wind_pmax.json",
            "static_clustering_wind_pmax.pkl"
        ],
        "dispatches.case_studies.renewables_case.data.steady_state_surrogate.dispatch_frequency.ss_surrogate_model_wind_pmax":[
            "keras_metadata.pb",
            "saved_model.pb"
        ],
        "dispatches.case_studies.renewables_case.data.steady_state_surrogate.dispatch_frequency.ss_surrogate_model_wind_pmax.variables":[
            "variables.data-00000-of-00001",
            "variables.index"
        ],
        "dispatches.case_studies.renewables_case.data.steady_state_surrogate.revenue":[
            "RE_revenue_params_2_25.json"
        ],
        "dispatches.case_studies.renewables_case.data.steady_state_surrogate.revenue.RE_revenue_2_25":[
            "keras_metadata.pb",
            "saved_model.pb"
        ],
        "dispatches.case_studies.renewables_case.data.steady_state_surrogate.revenue.RE_revenue_2_25.variables":[
            "variables.data-00000-of-00001",
            "variables.index"
        ],
        "dispatches.case_studies.fossil_case.ultra_supercritical_plant": [
            "pfd_ultra_supercritical_pc.svg",
        ],
        "dispatches.workflow.train_market_surrogates.dynamic.tests.data":[
            "inputdatatest.h5",
            "revdatatest.csv",
            "simdatatest.csv",
            "sample_clustering_model.json"
        ],
    },
)


#you have a roof thats 200m2 and a two different colours. Green colour costs 45 and can cover 0,4 meters of the roof and the Red one costs 53
#and can cover 0,7 meters of the roof. Wich one is the more cost efficient

roof= 200
color1 = 45
colorcoverage1 = 0.4
color2 = 53
colorcoverage2 = 0.7

x= color1*colorcoverage1*roof
y= color2*colorcoverage2*roof

if(x>y):
    print("The green is more cost efficient",x-y)
else: print("The red colour is less expensive",y-x)


from typing import Iterator, List, Tuple

import torch.nn as nn

from .group_rule import GroupRule

def named_parameters(module: nn.Module, prefix: str = '') -> Iterator[Tuple[str, nn.Parameter]]:
    memo = set()
    for name, p in module._parameters.items():
        if p is not None and p not in memo:
            memo.add(p)
            yield f'{prefix}.{name}' if prefix else name, p

def group_parameters(model: nn.Module, rules: List[GroupRule]) -> List[List[nn.Parameter]]:
    results: List[list] = [[] for _ in rules]

    def dfs(module: nn.Module, module_name: str = None, prefix: str = '', match_level: int = None):
        new_match_level = None
        for i, rule in enumerate(rules[:match_level]):
            if rule.match(module, module_name):
                new_match_level = i + 1
                break
        if new_match_level is None:
            new_match_level = match_level

        for name, child in module.named_children():
            dfs(child, name, prefix=f'{prefix}.{name}' if prefix else name, match_level=new_match_level)

        for name, param in named_parameters(module=module):
            new_prefix = f'{prefix}.{name}' if prefix else name
            for i, rule in enumerate(rules[:match_level]):
                if rule.match(module, module_name, param_name=name, prefix=new_prefix):
                    results[i].append(param)
                    break
            else:
                if new_match_level is not None:
                    results[new_match_level - 1].append(param)

    dfs(model)
    return results


from django.contrib import admin

from .models import Donor, Currency, Donation

#repeat for other models.
class DonationAdmin(admin.ModelAdmin):
    list_display = ('donor', 'datetime_of_donation',
                    'currency_type', 'amount','notified','section18a_issued',)
    #list_editable = ('donor', 'datetime_of_donation',
     #               'currency_type', 'amount','notified','section18a_issued',)

admin.site.register(Donor)
admin.site.register(Currency)
admin.site.register(Donation, DonationAdmin)


import pytest

from leetcode344 import Solution

@pytest.mark.parametrize("s,expected", [
    ([], []),
    (["a"], ["a"]),
    (["a", "b", "c"], ["c", "b", "a"]),
    (["a", "b", "c", "d"], ["d", "c", "b", "a"]),
])
def testSolutionReverseString(s, expected):
    assert Solution().reverseString(s) == expected


# Definition for a binary tree node.
from typing import Optional

class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right

class Solution1:
    def traversal(self, node, node_list):
        if node is None:
            return

        node_list.append(node)

        self.traversal(node.left, node_list)
        self.traversal(node.right, node_list)

    def flatten(self, root: Optional[TreeNode]) -> None:
        """
        Time O(n)
        Space O(n) 需要额外的list空间来记录前序遍历顺序下的node
        先前序遍历存储下来所有node顺序，再通过loop list重新给每个node分配right和left指针
        """
        node_list = []
        self.traversal(root, node_list)

        for i in range(len(node_list) - 1):
            cur_node = node_list[i]
            next_node = node_list[i + 1]
            cur_node.right = next_node
            cur_node.left = None

        return root

class Solution2:
    def traversal(self, node):
        if node is None:
            return

        # 后续遍历，需要先左右节点先处理完再处理当前节点
        self.traversal(node.left)
        self.traversal(node.right)

        # 后序遍历位置
        # 1、此时左右子树已经被拉平成一条链表
        left = node.left
        right = node.right

        # 2、将左子树作为右子树
        node.left = None
        node.right = left

        # 3、将原先的右子树接到当前右子树(也就是之前的左子树)的末端，需要while loop一直走到当前右子树的底端先
        p = node
        while p.right:
            p = p.right
        p.right = right

    def flatten(self, root: Optional[TreeNode]) -> None:
        """
        Time O(n)
        Space O(1)  不需要额外list空间存储顺序
        分解问题，分别解决左右子树，则解决了整个树。当前节点左节点变成右节点，然后原先的右节点接入左节点的右节点。
        """

        self.traversal(root)


import argparse
import pandas as pd
import time
import os
import numpy as np
from numpy.linalg import inv
from datetime import datetime
import time
import math
import copy
#import pcl
import open3d as o3d
import cv2

import rospy
from sensor_msgs.msg import Image, PointCloud2, PointField
from sensor_msgs import point_cloud2

from scipy.spatial.transform import Rotation as R
#import tf.transformations as tr
import time
import matplotlib.pyplot as plt

#plt.switch_backend('agg')

import torch
import torch.nn as nn
from model import MotionNet, MotionNetMGDA, FeatEncoder

global init, pcd_old, tf_curr
init = True
tf_curr = np.identity(4)

trained_model_path = '/mnt/Disk2/download/model_MGDA.pth'
img_save_dir = '/home/joinet/MotionNet/logs/itri/tmp'

save_directory = '/mnt/Disk1/itri_global_vel_i2o1_no_mask'

history_scan_num = 2
out_seq_len = 1

voxel_size = (0.25, 0.25, 0.4)
area_extents = np.array([[-32., 32.], [-32., 32.], [-3., 2.]])
use_adj_frame_pred = True
use_motion_state_pred_masking = True
color_map = {0: 'c', 1: 'm', 2: 'k', 3: 'y', 4: 'r'}
cat_names = {0: 'bg', 1: 'bus', 2: 'ped', 3: 'bike', 4: 'other'}

pc_list = []
tf_list = []

lidar_radar_ts_bias = 148.9810913 - 0.19 #1248.675218 # lidar_ts - radar_ts # TODO: this should be calibrated

lidar_buffer = []
lidar_ts_buffer = []
res_buffer = []
pose_buffer = []

radar_buffer = []
radar_ts_buffer = []

def pc_msg_to_np(pc_msg):
    gen = point_cloud2.read_points(pc_msg)
    pointcloud = []
    for p in gen:
        point = [p[0],p[1],p[2],p[3]]
        pointcloud.append(point)
    np_pointcloud = np.array(pointcloud)

    return np_pointcloud

def draw_registration_result(source, target, transformation):
    source_temp = copy.deepcopy(source)
    target_temp = copy.deepcopy(target)
    source.paint_uniform_color([0, 1, 0])
    source_temp.paint_uniform_color([1, 0, 0])
    target_temp.paint_uniform_color([0, 0, 1])
    source_temp.transform(transformation)
    o3d.visualization.draw_geometries([source_temp, target_temp, source])

def trans_pc(pc_list, tf_list):
    len_ = len(tf_list)
    local_tf_list = []
    for i in range(len_):
      local_lidar_tf = np.dot(np.linalg.inv(tf_list[len_-1]), tf_list[i])
      local_tf_list.append(local_lidar_tf)
    trans_pc_list = []
    for i in range(len_):
      trans_pc_list.append( np.dot(local_tf_list[i], pc_list[i].T).T )

    return trans_pc_list
def trans_single_pc(pc, mat):
  hom_pc = np.concatenate((pc, np.ones((pc.shape[0],1))), axis=1)
  hom_pc_ = np.dot(mat, hom_pc.T).T
  #p1 = o3d.geometry.PointCloud()
  #p1.points = o3d.utility.Vector3dVector(hom_pc[:,:-1])
  #p2 = o3d.geometry.PointCloud()
  #p2.points = o3d.utility.Vector3dVector(hom_pc_[:,:-1])
  #p1.paint_uniform_color([0, 0, 1])
  #p2.paint_uniform_color([1, 0, 0])
  #o3d.visualization.draw_geometries([p1, p2])
  return hom_pc_[:,:-1]
def voxelize_occupy(pts, voxel_size, extents=None, return_indices=False):
    """
    Voxelize the input point cloud. We only record if a given voxel is occupied or not, which is just binary indicator.

    The input for the voxelization is expected to be a PointCloud
    with N points in 4 dimension (x,y,z,i). Voxel size is the quantization size for the voxel grid.

    voxel_size: I.e. if voxel size is 1 m, the voxel space will be
    divided up within 1m x 1m x 1m space. This space will be 0 if free/occluded and 1 otherwise.
    min_voxel_coord: coordinates of the minimum on each axis for the voxel grid
    max_voxel_coord: coordinates of the maximum on each axis for the voxel grid
    num_divisions: number of grids in each axis
    leaf_layout: the voxel grid of size (numDivisions) that contain -1 for free, 0 for occupied

    :param pts: Point cloud as N x [x, y, z, i]
    :param voxel_size: Quantization size for the grid, vd, vh, vw
    :param extents: Optional, specifies the full extents of the point cloud.
                    Used for creating same sized voxel grids. Shape (3, 2)
    :param return_indices: Whether to return the non-empty voxel indices.
    """
    # Function Constants
    VOXEL_EMPTY = 0
    VOXEL_FILLED = 1

    # Check if points are 3D, otherwise early exit
    if pts.shape[1] < 3 or pts.shape[1] > 4:
        raise ValueError("Points have the wrong shape: {}".format(pts.shape))

    if extents is not None:
        if extents.shape != (3, 2):
            raise ValueError("Extents are the wrong shape {}".format(extents.shape))

        filter_idx = np.where((extents[0, 0] < pts[:, 0]) & (pts[:, 0] < extents[0, 1]) &
                              (extents[1, 0] < pts[:, 1]) & (pts[:, 1] < extents[1, 1]) &
                              (extents[2, 0] < pts[:, 2]) & (pts[:, 2] < extents[2, 1]))[0]
        pts = pts[filter_idx]

    # Discretize voxel coordinates to given quantization size
    discrete_pts = np.floor(pts[:, :3] / voxel_size).astype(np.int32)

    # Use Lex Sort, sort by x, then y, then z
    x_col = discrete_pts[:, 0]
    y_col = discrete_pts[:, 1]
    z_col = discrete_pts[:, 2]
    sorted_order = np.lexsort((z_col, y_col, x_col))

    # Save original points in sorted order
    discrete_pts = discrete_pts[sorted_order]

    # Format the array to c-contiguous array for unique function
    contiguous_array = np.ascontiguousarray(discrete_pts).view(
        np.dtype((np.void, discrete_pts.dtype.itemsize * discrete_pts.shape[1])))

    # The new coordinates are the discretized array with its unique indexes
    _, unique_indices = np.unique(contiguous_array, return_index=True)

    # Sort unique indices to preserve order
    unique_indices.sort()

    voxel_coords = discrete_pts[unique_indices]

    # Compute the minimum and maximum voxel coordinates
    if extents is not None:
        min_voxel_coord = np.floor(extents.T[0] / voxel_size)
        max_voxel_coord = np.ceil(extents.T[1] / voxel_size) - 1
    else:
        min_voxel_coord = np.amin(voxel_coords, axis=0)
        max_voxel_coord = np.amax(voxel_coords, axis=0)

    # Get the voxel grid dimensions
    num_divisions = ((max_voxel_coord - min_voxel_coord) + 1).astype(np.int32)

    # Bring the min voxel to the origin
    voxel_indices = (voxel_coords - min_voxel_coord).astype(int)

    # Create Voxel Object with -1 as empty/occluded
    leaf_layout = VOXEL_EMPTY * np.ones(num_divisions.astype(int), dtype=np.float32)

    # Fill out the leaf layout
    leaf_layout[voxel_indices[:, 0],
                voxel_indices[:, 1],
                voxel_indices[:, 2]] = VOXEL_FILLED

    if return_indices:
        return leaf_layout, voxel_indices
    else:
        return leaf_layout

def lidar_res_plot(viz_map, cat_pred, motion_pred_numpy, disp_pred, non_empty_map, ax):
    border_meter = 4
    border_pixel = border_meter * 4
    x_lim = [-(32 - border_meter), (32 - border_meter)]
    y_lim = [-(32 - border_meter), (32 - border_meter)]

    cat_pred = np.argmax(cat_pred, axis=0) + 1
    #cat_pred = (cat_pred * non_empty_map * filter_mask).astype(np.int)
    cat_pred = (cat_pred * non_empty_map).astype(np.int)

    idx_x = np.arange(256)
    idx_y = np.arange(256)
    idx_x, idx_y = np.meshgrid(idx_x, idx_y, indexing='ij')
    qk = [None] * len(color_map)  # for quiver key

    ax[0].imshow(np.rot90(viz_map), cmap = 'jet') # non_empty_map binary
    ax[0].axis('off')
    ax[0].set_aspect('equal')
    ax[0].title.set_text('LIDAR data')

    motion_pred_numpy = np.squeeze(motion_pred_numpy) * non_empty_map

    ax[1].imshow(np.rot90(motion_pred_numpy), cmap = 'gray')
    ax[1].axis('off')
    ax[1].set_aspect('equal')
    ax[1].title.set_text('motion segmentation')
    '''
    vel_map = np.sqrt(np.power(disp_pred[19,:,:,0],2)+np.power(disp_pred[19,:,:,1],2)) * non_empty_map
    upper_bound_mask = vel_map >= 20
    vel_map[upper_bound_mask] = 20

    ax[1,1].imshow(np.rot90(vel_map), cmap = 'jet', vmin=0, vmax=10)
    ax[1,1].axis('off')
    ax[1,1].set_aspect('equal')
    ax[1,1].title.set_text('vel_map')

    binary_vel_map = np.ones_like(vel_map)
    thd_mask = vel_map < 0.5
    binary_vel_map[thd_mask] = 0
    ax[1,2].imshow(np.rot90(binary_vel_map), cmap = 'gray')
    ax[1,2].axis('off')
    ax[1,2].set_aspect('equal')
    ax[1,2].title.set_text('vel_map')
    '''

    for k in range(len(color_map)):
      # ------------------------ Prediction ------------------------
      # Show the prediction results. We show the cells corresponding to the non-empty one-hot gt cells.
      mask_pred = cat_pred == (k + 1)
      field_pred = disp_pred[-1]  # Show last prediction, ie., the 20-th frame

      # For cells with very small movements, we threshold them to be static
      field_pred_norm = np.linalg.norm(field_pred, ord=2, axis=-1)  # out: (h, w)
      thd_mask = field_pred_norm <= 0.4
      field_pred[thd_mask, :] = 0

      # We use the same indices as the ground-truth, since we are currently focused on the foreground
      X_pred = idx_x[mask_pred]
      Y_pred = idx_y[mask_pred]
      U_pred = field_pred[:, :, 0][mask_pred] / voxel_size[0]
      V_pred = field_pred[:, :, 1][mask_pred] / voxel_size[1]

      qk[k] = ax[2].quiver(X_pred, Y_pred, U_pred, V_pred, angles='xy', scale_units='xy', scale=1, color=color_map[k])
      ax[2].quiverkey(qk[k], X=0.0 + k/5.0, Y=1.1, U=20, label=cat_names[k], labelpos='E')
      ax[2].quiver(X_pred, Y_pred, U_pred, V_pred, angles='xy', scale_units='xy', scale=1, width = 0.003, color=color_map[k])
      ax[2].set_xlim(border_pixel, field_pred.shape[0] - border_pixel)
      ax[2].set_ylim(border_pixel, field_pred.shape[1] - border_pixel)
      ax[2].set_aspect('equal')
      ax[2].title.set_text('Prediction')
      ax[2].axis('off')

def img_cart_to_polar(img):
  img = cv2.resize(img, (260, 260), interpolation=cv2.INTER_LINEAR)
  img = np.rot90(img,-1) # rot clockwise 90 deg
  value = np.sqrt(((img.shape[0]/2.0)**2.0)+((img.shape[1]/2.0)**2.0))
  polar_image = cv2.linearPolar(img,(img.shape[0]/2, img.shape[1]/2), value, cv2.WARP_FILL_OUTLIERS)
  return polar_image

def img_polar_to_cart(img):
  img = np.concatenate((img[int(img.shape[0]/4):,:], img[:int(img.shape[0]/4),:]), axis=0)
  value = np.sqrt(((img.shape[0]/2.0)**2.0)+((img.shape[1]/2.0)**2.0))
  cart_image = cv2.linearPolar(img,(img.shape[0]/2, img.shape[1]/2), value, cv2.WARP_FILL_OUTLIERS + cv2.WARP_INVERSE_MAP)
  cart_image = cv2.resize(cart_image, (256, 256), interpolation=cv2.INTER_LINEAR)
  return cart_image
def combined_lidar_scans(data_list):
  combined_polar_lidar = np.array([])
  for i in range(len(data_list)):
    polar_lidar = img_cart_to_polar(data_list[i])
    polar_lidar_clip = polar_lidar[i*52:(i+1)*52,:]
    if i == 0:
      combined_polar_lidar = polar_lidar_clip
    else:
      combined_polar_lidar = np.concatenate((combined_polar_lidar, polar_lidar_clip), axis=0)
  return combined_polar_lidar
def gen_gt_vel(disp_list):
  combined_polar_lidar = combined_lidar_scans(disp_list)
  result_cart_radar = img_polar_to_cart(combined_polar_lidar)
  return result_cart_radar
def gen_gt_mask(rmg_data_list, radar_cart):
  MASKING=False
  combined_polar_lidar = combined_lidar_scans(rmg_data_list) # motion_seg_list rmg_data_list
  combined_polar_lidar[combined_polar_lidar<0.5] = 0
  combined_polar_lidar[combined_polar_lidar>0.5] = 1

  radar_polar = img_cart_to_polar(radar_cart)

  if MASKING == True:
    # masking #
    result_polar_radar = mask_radar_by_lidar(radar_polar, combined_polar_lidar)
  else:
    # no masking #
    result_polar_radar = combined_polar_lidar

  result_cart_radar = img_polar_to_cart(result_polar_radar)

  combined_cart_lidar = img_polar_to_cart(combined_polar_lidar)
  combined_cart_lidar[combined_cart_lidar<0.5] = 0
  combined_cart_lidar[combined_cart_lidar>0.5] = 1

  return result_cart_radar, combined_cart_lidar # with mask, without mask

def viz_err_fig(raw_radar, gt, res):
  viz_err = np.zeros((256,256,3))
  viz_err[:,:,0] = gt * np.logical_not(res)
  viz_err[:,:,1] = gt * res
  viz_err[:,:,2] = res * np.logical_not(gt)
  viz_radar = np.stack((raw_radar,raw_radar,raw_radar), axis=2)
  return (viz_err + viz_radar*3)/2.

def get_sync(t, timestamps):
    """get the closest id given the timestamp

    :param t: timestamp in seconds
    :type t: float
    :param all_timestamps: a list with all timestamps
    :type all_timestamps: np.array
    :param time_offset: offset in case there is some unsynchronoised sensor, defaults to 0.0
    :type time_offset: float, optional
    :return: the closest id
    :rtype: int
    """
    idx = np.argmin(np.abs(timestamps - t))
    return idx, timestamps[idx]

def GetMatFromXYYaw(x,y,yaw):
    tf_mat = np.matrix([[math.cos(yaw), -math.sin(yaw), x]
                       ,[math.sin(yaw), math.cos(yaw), y]
                       ,[0, 0, 1]])
    return tf_mat

def gen_static_disp_map(radar_cart_list, tf_mat_list, result_cart_radar_disp_0, result_cart_radar_disp_1, out_seq_len, result_cart_radar_rmg, result_cart_radar_seg):
  N=256
  cart_resolution = 0.25
  static_disp_map_list = list()

  field_gt = np.zeros((N,N,2))
  field_gt[:,:,0] = result_cart_radar_disp_0
  field_gt[:,:,1] = result_cart_radar_disp_1
  mask = result_cart_radar_rmg.astype(np.bool) * result_cart_radar_seg.astype(np.bool)
  #mask = np.zeros((256,256))
  field_gt = field_gt*mask.reshape((N,N,1))

  assert len(tf_mat_list)>out_seq_len
  for idx in range(out_seq_len):
    tf_mat = tf_mat_list[idx+1] # top: x, right: y

    # transform ego-motion to pixel unit and w.r.t disp frame
    print('delta (m) :', tf_mat[0,3], tf_mat[1,3], -np.arctan(tf_mat[1,0] / tf_mat[0,0]))
    delta_x = -tf_mat[1,3]/cart_resolution
    delta_y = -tf_mat[0,3]/cart_resolution
    delta_yaw = -np.arctan(tf_mat[1,0] / tf_mat[0,0])

    delta_tf = GetMatFromXYYaw(delta_x, delta_y, delta_yaw) # w.r.t disp frame
    #print('delta:', delta_x, delta_y, delta_yaw)

    valid_mask = result_cart_radar_rmg.astype(np.bool) # only compute disp for valid pixel
    static_disp_map = np.zeros((N,N,2))
    for i in range(N): # row
      for j in range(N): # col
        if valid_mask[i,j] == 1:
          center=N/2-0.5
          local_tf_curr_t = GetMatFromXYYaw(j-center, center-i, 0) #GetMatFromXYYaw(j-center, center-i, 0)
          disp_wrt_curr_t = field_gt[i,j,:]
          #print(disp_wrt_curr_t)
          disp_tf_wrt_curr_t = GetMatFromXYYaw(-disp_wrt_curr_t[1], -disp_wrt_curr_t[0], 0)
          local_tf_old_t = delta_tf.dot(inv(disp_tf_wrt_curr_t.dot(inv(local_tf_curr_t))))

          #local_tf_old_t = np.dot(delta_tf, local_tf_curr_t)
          x_mo = local_tf_old_t[0,2] - local_tf_curr_t[0,2]
          y_mo = local_tf_old_t[1,2] - local_tf_curr_t[1,2]
          static_disp_map[i,j,0] = x_mo # x-axis of disp frame
          static_disp_map[i,j,1] = y_mo # y-axis of disp frame
    static_disp_map_list.append(static_disp_map*result_cart_radar_rmg.astype(np.bool).reshape((N,N,1)))
  return static_disp_map_list

def callback_lidar(cloud):
  global init, pcd_old, tf_curr
  rospy.loginfo("Lidar %s: %s", cloud.header.seq, cloud.header.stamp.to_sec() - lidar_radar_ts_bias)
  pc = pc_msg_to_np(cloud)
  pc = pc[:,:-1] # rm intensity info
  #print(pc.shape)

  # x y z y p r
  # -0.75 0.1 0 0.06 0 0
  # -0.95 0.0 -1 0.055 0 0
  mat = np.identity(4) # lidar2radar
  r = R.from_euler('z', 0.055, degrees=False)
  mat[:3,:3] = r.as_matrix()
  mat[0,3] = -0.8 #-0.95
  mat[1,3] = 0.0
  mat = np.linalg.inv(mat) # radar2lidar
  pc = trans_single_pc(pc, mat) # trans pc to radar frame

  # rmg
  z_low_idx = np.where(pc[:,2]<-1.3)[0]
  pc_rmg = np.delete(pc, z_low_idx, axis=0)

  pcd = o3d.geometry.PointCloud()
  pcd.points = o3d.utility.Vector3dVector(pc_rmg)

  #hom_pc = np.concatenate((pc, np.ones((pc.shape[0],1))), axis=1)
  #lidar_buffer.append(hom_pc)
  #lidar_ts_buffer.append(cloud.header.stamp.to_sec() - lidar_radar_ts_bias)

  if init:
    pcd_old = copy.deepcopy(pcd)
    init = False

  ### icp ###
  downpcd = pcd.voxel_down_sample(voxel_size=0.3) #0.2
  downpcd_old = pcd_old.voxel_down_sample(voxel_size=0.3)
  threshold = 2
  trans_init = np.identity(4)
  reg_p2p = o3d.registration.registration_icp(
          downpcd, downpcd_old, threshold, trans_init,
          o3d.registration.TransformationEstimationPointToPoint(), o3d.registration.ICPConvergenceCriteria(max_iteration = 300))
  tf_curr = np.dot(tf_curr, reg_p2p.transformation)
  #draw_registration_result(downpcd, downpcd_old, reg_p2p.transformation)
  #print(tf_curr)

  hom_pc = np.concatenate((pc, np.ones((pc.shape[0],1))), axis=1)
  hom_pc_rmg = np.concatenate((pc_rmg, np.ones((pc_rmg.shape[0],1))), axis=1)
  pc_list.append(hom_pc)
  tf_list.append(tf_curr)

  if len(pc_list) > 5 and len(tf_list) > 5:
    pc_list.pop(0)
    tf_list.pop(0)
    trans_pc_list = trans_pc(pc_list, tf_list)

    voxel_indices_list = list()
    padded_voxel_points_list = list()
    for i in range(5):
      res, voxel_indices = voxelize_occupy(trans_pc_list[i], voxel_size=voxel_size, extents=area_extents, return_indices=True)
      voxel_indices_list.append(voxel_indices)
      padded_voxel_points_list.append(res)

    padded_voxel_points = np.stack(padded_voxel_points_list, axis=0).astype(np.bool) # Shape: (5, 256, 256, 13)
    non_empty_map = padded_voxel_points[-1,:,:,:].any(axis=2)

    viz_map = np.zeros(padded_voxel_points[-1,:,:,0].shape)
    occ_map = np.ones(padded_voxel_points[-1,:,:,0].shape)
    for i in range(padded_voxel_points[-1,:,:,:].shape[2]):
      k = padded_voxel_points[-1,:,:,:].shape[2]-1-i
      viz_map = viz_map + padded_voxel_points[-1,:,:,k]*(k+1) * occ_map
      occ_map = np.logical_not(padded_voxel_points[-1,:,:,-i-1:].any(axis=2))

    padded_voxel_points = padded_voxel_points.astype(np.float32)
    padded_voxel_points = torch.from_numpy(padded_voxel_points)
    padded_voxel_points = torch.unsqueeze(padded_voxel_points, 0).to(device)
    model_encoder.eval()
    model_head.eval()
    with torch.no_grad():
      shared_feats = model_encoder(padded_voxel_points)
      disp_pred, cat_pred, motion_pred = model_head(shared_feats)

      disp_pred = disp_pred.cpu().numpy()
      disp_pred = np.transpose(disp_pred, (0, 2, 3, 1))
      cat_pred = np.squeeze(cat_pred.cpu().numpy(), 0)

      if use_adj_frame_pred:  # The prediction are the displacement between adjacent frames
          for c in range(1, disp_pred.shape[0]):
              disp_pred[c, ...] = disp_pred[c, ...] + disp_pred[c - 1, ...]

      if use_motion_state_pred_masking:
          motion_pred_numpy = motion_pred.cpu().numpy()
          motion_pred_numpy = np.argmax(motion_pred_numpy, axis=1)
          motion_mask = motion_pred_numpy == 0 # index of static pixel

          cat_pred_numpy = np.argmax(cat_pred, axis=0)
          cat_mask = np.logical_and(cat_pred_numpy == 0, non_empty_map == 1)
          cat_mask = np.expand_dims(cat_mask, 0)

          cat_weight_map = np.ones_like(motion_pred_numpy, dtype=np.float32)
          cat_weight_map[motion_mask] = 0.0
          cat_weight_map[cat_mask] = 0.0
          cat_weight_map = cat_weight_map[:, :, :, np.newaxis]  # (1, h, w. 1)

          disp_pred = disp_pred * cat_weight_map

    lidar_buffer.append(hom_pc_rmg)
    lidar_ts_buffer.append(cloud.header.stamp.to_sec() - lidar_radar_ts_bias)

    res_dict = dict()
    res_dict['disp_pred'] = disp_pred[-1]
    res_dict['cat_pred'] = cat_pred
    res_dict['motion_pred'] = motion_pred_numpy
    res_dict['non_empty_map'] = non_empty_map
    res_buffer.append(res_dict)

    pose_buffer.append(tf_curr)

    if len(lidar_ts_buffer) > 30:
      lidar_ts_buffer.pop(0)
      lidar_buffer.pop(0)
      res_buffer.pop(0)
      pose_buffer.pop(0)

#    disp_pred = np.rot90(res_dict['disp_pred'],2)
#    non_empty_map = np.rot90(res_dict['non_empty_map'],2)
#    motion_pred = np.rot90(res_dict['motion_pred'],2)

#    vel_map = np.sqrt(np.power(disp_pred[:,:,0],2)+np.power(disp_pred[:,:,1],2)) * non_empty_map
#    binary_vel_map = np.ones_like(vel_map)
#    thd_mask = vel_map < 0.01 # 0.5
#    binary_vel_map[thd_mask] = 0
#    print(binary_vel_map.shape)
#    fig, ax = plt.subplots(1, 1, figsize=(7, 7))
#    ax.imshow(binary_vel_map)
#    field_gt = np.zeros((256,256,2))
#    field_gt[:,:,0] = disp_pred[:,:,1]
#    field_gt[:,:,1] = disp_pred[:,:,0]
#    # viz same as oxford ramnet
#    idx_x = np.arange(field_gt.shape[0])
#    idx_y = np.arange(field_gt.shape[1])
#    idx_x, idx_y = np.meshgrid(idx_x, idx_y, indexing='ij')
#    mask = non_empty_map.astype(np.bool) #  * combined_cart_rmg.astype(np.bool) is better !!
#    X = idx_x[mask]
#    Y = idx_y[mask]
#    U = field_gt[:, :, 0][mask]
#    V = field_gt[:, :, 1][mask]
#    qk = ax.quiver(Y, X, U, V, angles='xy', scale_units='xy', scale=0.2, width=0.0015, headwidth=10, color='r', minlength=3, minshaft=1)
#    plt.show()

    # Plot
    #fig, ax = plt.subplots(1, 3, figsize=(21, 7))
    #lidar_res_plot(viz_map, cat_pred, motion_pred_numpy, disp_pred, non_empty_map, ax)
    #plt.savefig(os.path.join(img_save_dir, 'res_'+str(cloud.header.stamp) + '.png'))
    #plt.show()
    #ax[0].clear()
    #ax[1].clear()
    #ax[2].clear()

  pcd_old = copy.deepcopy(pcd)

def callback_radar(radar_msg):
  global fig2, ax2
  rospy.loginfo("I heard Radar %s: %s", radar_msg.header.seq, radar_msg.header.stamp.to_sec())

  radar_img = np.frombuffer(radar_msg.data, dtype=np.uint8).reshape(radar_msg.height, radar_msg.width, -1)[...,0]
  radar_img = radar_img.astype(np.float32)/255.
#  ax2.imshow(radar_img, cmap='gray')
#  plt.pause(0.1)
#  ax2.clear()

  radar_buffer.append(radar_img)
  radar_ts_buffer.append(radar_msg.header.stamp.to_sec())

  history_scan_num=2

  if len(radar_buffer) > history_scan_num+2:
    curr_radar_img = radar_buffer[-2]
    curr_radar_ts = radar_ts_buffer[-2]
    print('-----RUN-----')
    print(curr_radar_ts)
    print('len(lidar_ts_buffer)',len(lidar_ts_buffer))
    print('len(lidar_buffer)',len(lidar_buffer))
    print('len(res_buffer)',len(res_buffer))
    print('len(pose_buffer)',len(pose_buffer))
    if len(lidar_ts_buffer)<3:
      print('no sync lidar result yet...')
      return
    sync_lidar_idx, sync_lidar_ts = get_sync(curr_radar_ts, np.array(lidar_ts_buffer))
    while(np.abs(sync_lidar_ts-curr_radar_ts)>0.1 or sync_lidar_idx>len(lidar_ts_buffer)-3):
      print('wait for lidar processing...')
      time.sleep(0.55)
      sync_lidar_idx, sync_lidar_ts = get_sync(curr_radar_ts, np.array(lidar_ts_buffer))
    print('len(lidar_ts_buffer)',len(lidar_ts_buffer))
    print(sync_lidar_idx, sync_lidar_ts, 't diff:', np.abs(sync_lidar_ts-curr_radar_ts))

    rmg_data_list = list()
    voxel_size = (0.25, 0.25, 10.0)
    area_extents = np.array([[-32., 32.], [-32., 32.], [-10., 10.]])
    for i in range(5):
      padded_voxel_points, voxel_indices = voxelize_occupy(lidar_buffer[sync_lidar_idx+i], voxel_size=voxel_size, extents=area_extents, return_indices=True)
      non_empty_map = padded_voxel_points[:,:,:].any(axis=2).astype(np.float32)
      non_empty_map = np.rot90(non_empty_map,2)
      rmg_data_list.append(non_empty_map)

    motion_seg_list = list()
    disp_0_list = list()
    disp_1_list = list()
    for i in range(5):
      # Read res dict
      res_data = res_buffer[sync_lidar_idx+i]
      disp_pred = res_data['disp_pred']
      non_empty_map = res_data['non_empty_map']
      motion_pred = res_data['motion_pred']

      # Coordinate transform
      disp_pred = np.rot90(disp_pred,2)
      non_empty_map = np.rot90(non_empty_map,2)
      motion_pred = np.rot90(motion_pred,2)

      # Pre-process
      vel_map = np.sqrt(np.power(disp_pred[:,:,0],2)+np.power(disp_pred[:,:,1],2)) * non_empty_map
      binary_vel_map = np.ones_like(vel_map)
      thd_mask = vel_map < 0.01 # 0.5
      binary_vel_map[thd_mask] = 0
      #print(binary_vel_map.shape)

#      fig, ax = plt.subplots(1, 1, figsize=(7, 7))
#      ax.imshow(binary_vel_map)
#      field_gt = np.zeros((256,256,2))
#      field_gt[:,:,0] = disp_pred[:,:,1]
#      field_gt[:,:,1] = disp_pred[:,:,0]
#      # viz same as oxford ramnet
#      idx_x = np.arange(field_gt.shape[0])
#      idx_y = np.arange(field_gt.shape[1])
#      idx_x, idx_y = np.meshgrid(idx_x, idx_y, indexing='ij')
#      mask = non_empty_map.astype(np.bool) #  * combined_cart_rmg.astype(np.bool) is better !!
#      X = idx_x[mask]
#      Y = idx_y[mask]
#      U = field_gt[:, :, 0][mask]
#      V = field_gt[:, :, 1][mask]
#      qk = ax.quiver(Y, X, U, V, angles='xy', scale_units='xy', scale=0.2, width=0.0015, headwidth=10, color='r', minlength=3, minshaft=1)
#      plt.show()

      motion_seg_list.append(binary_vel_map)
      disp_0_list.append(disp_pred[:,:,0])
      disp_1_list.append(disp_pred[:,:,1])

    tf_mat_list = list()
    curr_pose = pose_buffer[sync_lidar_idx]
    for i in range(history_scan_num):
      old_pose = pose_buffer[sync_lidar_idx-i]
      tf_mat = np.linalg.inv(old_pose).dot(curr_pose)
      tf_mat_list.append(tf_mat)

    result_cart_radar_rmg, combined_cart_rmg = gen_gt_mask(rmg_data_list, radar_img)
    result_cart_radar_seg, combined_cart_seg = gen_gt_mask(motion_seg_list, radar_img)
    result_cart_radar_disp_0 = gen_gt_vel(disp_0_list)
    result_cart_radar_disp_1 = gen_gt_vel(disp_1_list)

    ### Generate disp by motion ###
    #result_cart_radar_disp_0 = np.zeros((256,256))
    #result_cart_radar_disp_1 = np.zeros((256,256))
    #result_disp_map_seq_list = gen_static_disp_map(radar_buffer, tf_mat_list, result_cart_radar_disp_0, result_cart_radar_disp_1, out_seq_len, result_cart_radar_rmg, result_cart_radar_seg)

    save_data_dict = dict()
    for i in range(history_scan_num):
      save_data_dict['raw_radar_' + str(i)] = radar_buffer[-2-i]
    save_data_dict['gt_disp_global'] = np.concatenate((np.expand_dims(result_cart_radar_disp_0,axis=2), np.expand_dims(result_cart_radar_disp_1,axis=2)), axis=2)
    save_data_dict['gt_radar'] = result_cart_radar_rmg
    save_data_dict['gt_moving'] = result_cart_radar_seg
    save_data_dict['gt_lidar'] = combined_cart_rmg

    save_file_name = os.path.join(save_directory, str(curr_radar_ts) + '.npy')
    #np.save(save_file_name, arr=save_data_dict)

    if len(radar_ts_buffer) > 10:
      radar_ts_buffer.pop(0)
      radar_buffer.pop(0)

#    fig, ax = plt.subplots(1, 3, figsize=(14, 7))

#    viz = viz_err_fig(curr_radar_img, result_cart_radar_seg, result_cart_radar_rmg)

#    fig, ax = plt.subplots(1, 1, figsize=(7, 7))
#    ax.imshow(result_cart_radar_seg)
#    field_gt = np.zeros((256,256,2))
#    field_gt[:,:,0] = result_cart_radar_disp_1
#    field_gt[:,:,1] = result_cart_radar_disp_0
#    # viz same as oxford ramnet
#    idx_x = np.arange(field_gt.shape[0])
#    idx_y = np.arange(field_gt.shape[1])
#    idx_x, idx_y = np.meshgrid(idx_x, idx_y, indexing='ij')
#    mask = non_empty_map.astype(np.bool) #  * combined_cart_rmg.astype(np.bool) is better !!
#    X = idx_x[mask]
#    Y = idx_y[mask]
#    U = field_gt[:, :, 0][mask]
#    V = field_gt[:, :, 1][mask]
#    qk = ax.quiver(Y, X, U, V, angles='xy', scale_units='xy', scale=0.2, width=0.0015, headwidth=10, color='r', minlength=3, minshaft=1)
#    plt.show()

    #plt.pause(0.1)
    #plt.savefig(os.path.join(img_save_dir, str(radar_msg.header.stamp.to_sec()) + '.png'))
    #ax2.clear()

def listener():
    global fig, ax, fig2, ax2

    #fig, ax = plt.subplots(1, 3, figsize=(9, 3))
    #fig2, ax2 = plt.subplots(1, 1, figsize=(4, 4))
    #plt.tight_layout()

    rospy.init_node('listener', anonymous=True)
    rospy.Subscriber("/velodyne_points", PointCloud2, callback_lidar, queue_size=100000)
    rospy.Subscriber("/Navtech/Cartesian", Image, callback_radar, queue_size=100000)
    rospy.spin()

if __name__ == '__main__':
    ### Load Model ###
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(device)
    model_encoder = FeatEncoder()
    model_head = MotionNetMGDA(out_seq_len=20, motion_category_num=2)
    model_encoder = nn.DataParallel(model_encoder)
    model_head = nn.DataParallel(model_head)
    checkpoint = torch.load(trained_model_path)
    model_encoder.load_state_dict(checkpoint['encoder_state_dict'])
    model_head.load_state_dict(checkpoint['head_state_dict'])
    model_encoder = model_encoder.to(device)
    model_head = model_head.to(device)
    loaded_models = [model_encoder, model_head]
    print("Loaded pretrained model...")
    model_encoder = loaded_models[0]
    model_head = loaded_models[1]
    listener()


import numpy as np
import time

def count_centroid(base, centroid_l):
    # count the distance for each item and centroid to get the distance_table
    distance_table = None
    for i, vecs in enumerate(base, 0):
        if i == 0:
            distance_table = [np.linalg.norm(base[0] - centroid) for centroid in centroid_l]
            distance_table = np.array([distance_table])
            continue
        tmp_dis = [np.linalg.norm(vecs - centroid) for centroid in centroid_l]
        tmp_dis = np.array([tmp_dis])
        distance_table = np.append(distance_table, tmp_dis, axis=0)
    # print(distance_table.shape)
    # get the nearest centroid and use it as the label
    labels = np.argmin(distance_table, axis=1)
    return labels

if __name__ == '__main__':
    np.random.seed(0)
    data = np.random.normal(size=100000).reshape(-1, 2)
    centroid = np.random.normal(size=100).reshape(-1, 2)
    start_time = time.time()
    labels = count_centroid(data, centroid)
    end_time = time.time()
    np.savetxt('label.txt', labels, fmt='%d')
    print('time consumed %.3f' % (end_time - start_time))