# -*- coding: utf-8 -*-

import datetime
import json
import sys
from threading import Thread

from resources.lib.common import tools
from resources.lib.indexers.trakt import TraktAPI
from resources.lib.modules import database
from resources.lib.modules.trakt_sync.shows import TraktSyncDatabase
from resources.lib.modules.trakt_sync.hidden import TraktSyncDatabase as HiddenDatabase

try:
    from Queue import Queue
except:
    from queue import Queue

sysaddon = sys.argv[0]
try:
    syshandle = int(sys.argv[1])
except:
    syshandle = ''
trakt = TraktAPI()

language_code = tools.get_language_code()

trakt_database = TraktSyncDatabase()
hidden_database = HiddenDatabase()

class Menus:
    def __init__(self):
        self.itemList = []
        self.threadList = []
        self.direct_episode_threads = []
        self.title_appends = tools.getSetting('general.appendtitles')
        self.task_queue = Queue(40)

    ######################################################
    # MENUS
    ######################################################

    def onDeckShows(self):
        hidden_shows = hidden_database.get_hidden_items('progress_watched', 'shows')
        trakt_list = trakt.json_response('sync/playback/episodes', limit=True)

        if trakt_list is None:
            return
        trakt_list = [i for i in trakt_list if i['show']['ids']['trakt'] not in hidden_shows]
        trakt_list = sorted(trakt_list, key=lambda i: tools.datetime_workaround(i['paused_at'][:19],
                                                                                format="%Y-%m-%dT%H:%M:%S",
                                                                                date_only=False), reverse=True)
        filter_list = []
        showList = []
        sort_list = []
        for i in trakt_list:
            if i['show']['ids']['trakt'] not in filter_list:
                if int(i['progress']) != 0:
                    showList.append(i)
                    filter_list.append(i['show']['ids']['trakt'])
                    sort_list.append(i['show']['ids']['trakt'])

        sort = {'type': 'showInfo', 'id_list': sort_list}
        self.mixedEpisodeBuilder(showList, sort=sort)
        tools.closeDirectory('tvshows')

    def discoverShows(self):

        tools.addDirectoryItem(tools.lang(32007), 'showsPopular&page=1', '', '')
        if tools.getSetting('trakt.auth') is not '':
            tools.addDirectoryItem(tools.lang(32008), 'showsRecommended', '', '')
        # tools.addDirectoryItem('This Years Most Popular', '', '', '')
        tools.addDirectoryItem(tools.lang(32009), 'showsTrending&page=1', '', '')
        tools.addDirectoryItem(tools.lang(32067), 'showsNew', '', '')
        tools.addDirectoryItem(tools.lang(32010), 'showsPlayed&page=1', '', '')
        tools.addDirectoryItem(tools.lang(32011), 'showsWatched&page=1', '', '')
        tools.addDirectoryItem(tools.lang(32012), 'showsCollected&page=1', '', '')
        tools.addDirectoryItem(tools.lang(32013), 'showsAnticipated&page=1', '', '')
        tools.addDirectoryItem(tools.lang(32014), 'showsUpdated&page=1', '', '')
        tools.addDirectoryItem(tools.lang(40121), 'showsNetworks', '', '')
        tools.addDirectoryItem(tools.lang(40123), 'showYears', '', '')
        tools.addDirectoryItem(tools.lang(32062), 'tvGenres', '', '')
        tools.addDirectoryItem(tools.lang(40151), 'showsByActor', '', '')
        # show genres is now labeled as tvGenres to support genre icons in skins
        if tools.getSetting('searchHistory') == 'false':
            tools.addDirectoryItem(tools.lang(32016), 'showsSearch', '', '')
        else:
            tools.addDirectoryItem(tools.lang(32016), 'showsSearchHistory', '', '')
        tools.closeDirectory('addons')

    def myShows(self):
        tools.addDirectoryItem(tools.lang(32063), 'onDeckShows', None, None)
        tools.addDirectoryItem(tools.lang(32017), 'showsMyCollection', '', '')
        tools.addDirectoryItem(tools.lang(32018), 'showsMyWatchlist', '', '')
        tools.addDirectoryItem('Next Up', 'showsNextUp', '', '')
        tools.addDirectoryItem('Upcoming Episodes', 'myUpcomingEpisodes', '', '')
        tools.addDirectoryItem('Unfinished Shows in Collection', 'showsMyProgress', '', '')
        tools.addDirectoryItem('Recent Episodes', 'showsMyRecentEpisodes', '', '')
        tools.addDirectoryItem('My Show Lists', 'myTraktLists&actionArgs=shows', '', '')
        tools.closeDirectory('addons')

    def myShowCollection(self):
        trakt_list = trakt_database.get_collected_episodes()
        trakt_list = [i for i in trakt_list if i is not None]
        trakt_list = list(set([i['show_id'] for i in trakt_list]))
        trakt_list = [{'ids': {'trakt': i}} for i in trakt_list]
        trakt_list = [i for i in trakt_list if i is not None]
        if trakt_list is None:
            return
        self.showListBuilder(trakt_list)
        tools.closeDirectory('tvshows', sort='title')

    def myShowWatchlist(self):
        trakt_list = trakt.json_response('users/me/watchlist/shows', limit=False)
        if trakt_list is None:
            return
        try:
            sort_by = trakt.response_headers['X-Sort-By']
            sort_how = trakt.response_headers['X-Sort-How']
            trakt_list = trakt.sort_list(sort_by, sort_how, trakt_list, 'show')
        except:
            tools.log('Failed to sort trakt list by response headers', 'error')
            pass
        self.showListBuilder(trakt_list)
        tools.closeDirectory('tvshows')

    def myProgress(self):

        collected_episodes = trakt_database.get_collected_episodes()
        collection = list(set([i['show_id'] for i in collected_episodes]))
        if len(collection) == 0:
            return

        show_dicts = []
        for i in collection:
            show_dicts.append({'show': {'ids': {'trakt': i}}})

        show_meta_list = trakt_database.get_show_list(show_dicts)
        unfinished = []

        for show in show_meta_list:
            if show['info']['playcount'] == 0:
                unfinished.append(show)

        self.showListBuilder(unfinished)
        tools.closeDirectory('tvshows', sort='title')

    def newShows(self):

        hidden = hidden_database.get_hidden_items('recommendations', 'shows')
        datestring = datetime.datetime.today() - datetime.timedelta(days=29)
        trakt_list = database.get(trakt.json_response, 12, 'calendars/all/shows/new/%s/30?languages=%s' %
                                  (datestring.strftime('%d-%m-%Y'), language_code))

        if trakt_list is None:
            return
        # For some reason trakt messes up their list and spits out tons of duplicates so we filter it
        duplicate_filter = []
        temp_list = []
        for i in trakt_list:
            if not i['show']['ids']['tvdb'] in duplicate_filter:
                duplicate_filter.append(i['show']['ids']['tvdb'])
                temp_list.append(i)

        trakt_list = temp_list

        trakt_list = [i for i in trakt_list if i['show']['ids']['trakt'] not in hidden]

        if len(trakt_list) > 40:
            trakt_list = trakt_list[:40]
        self.showListBuilder(trakt_list)
        tools.closeDirectory('tvshows')

    def myNextUp(self, ):

        watched_shows = trakt_database.get_watched_shows()
        hidden_shows = hidden_database.get_hidden_items('progress_watched', 'shows')

        watched_shows = [i for i in watched_shows if i['trakt_id'] not in hidden_shows]
        watched_episodes = trakt_database.get_watched_episodes()

        self._start_queue_workers()

        for show in watched_shows:
            self.task_queue.put((self._get_next_episode_to_watch, (show, watched_episodes)), block=True)

        self._finish_queue_workers()

        if tools.getSetting('nextup.sort') == '1':
            watched_list = trakt.json_response('users/me/watched/shows')
            watched_list = sorted(watched_list, key=lambda i: i['last_watched_at'], reverse=True)
            watched_list = [i['show']['ids']['trakt'] for i in watched_list]
            sort = {'type': 'showInfo', 'id_list': watched_list}
        else:
            sort = None

        episodes = self.itemList
        self.itemList = []

        self.mixedEpisodeBuilder(episodes, sort=sort, hide_watched=True)

        tools.closeDirectory('tvshows')

    def _get_next_episode_to_watch(self, show_db_dict, watched_episodes):
        try:
            show_id = show_db_dict['trakt_id']

            if show_db_dict['kodi_meta'] == {}:
                show_db_dict['kodi_meta'] = trakt_database.get_single_show(show_id)

            watched_episodes = [i for i in watched_episodes if i['show_id'] == show_id]
            watched_episodes = sorted(watched_episodes, key=lambda episode: episode['season'], reverse=True)

            season = watched_episodes[0]['season']
            season_meta = trakt_database.get_single_season(show_id, season)

            watched_episodes = [i for i in watched_episodes if i['season'] == season]
            watched_episodes = sorted(watched_episodes, key=lambda episode: episode['number'], reverse=True)
            last_watched_episode = watched_episodes[0]['number']
            next_episode = int(watched_episodes[0]['number']) + 1

            if season_meta is None:
                tools.log('Could not acquire season meta information for %s Season %s' % (show_id, season), 'error')
                return

            if season_meta['info']['episode_count'] == len(watched_episodes) \
                    or season_meta['info']['episode_count'] == last_watched_episode:

                if int(show_db_dict['kodi_meta']['info']['season_count']) > season:
                    season += 1
                    next_episode = 1

            episode_dict = {'show': {'ids': {'trakt': show_id}},
                            'episode': {'season': season, 'number': next_episode}}

            self.itemList.append(episode_dict)
        except KeyError:
            import traceback
            traceback.print_exc()
            pass
        except:
            import traceback
            traceback.print_exc()

    def myRecentEpisodes(self):
        hidden_shows = hidden_database.get_hidden_items('calendar', 'shows')
        datestring = datetime.datetime.today() - datetime.timedelta(days=13)
        trakt_list = database.get(trakt.json_response, 12, 'calendars/my/shows/%s/14' %
                                  datestring.strftime('%d-%m-%Y'))

        if trakt_list is None:
            return
        trakt_list = [i for i in trakt_list if i['show']['ids']['trakt'] not in hidden_shows]
        self.mixedEpisodeBuilder(trakt_list)
        tools.closeDirectory('episodes')

    def myUpcomingEpisodes(self):
        tomorrow = (datetime.date.today() + datetime.timedelta(days=1)).strftime('%Y-%m-%d')
        upcoming_episodes = database.get(trakt.json_response, 24, 'calendars/my/shows/%s/30' % tomorrow)

        sort = sorted(upcoming_episodes, key=lambda i: i['first_aired'])
        sort = [i['episode']['ids']['trakt'] for i in sort]
        sort = {'type': None, 'id_list': sort}
        self.mixedEpisodeBuilder(upcoming_episodes, sort=sort, hide_watched=False, hide_unaired=False,
                                 prepend_date=True)
        tools.closeDirectory('episodes')

    def showsNetworks(self):
        trakt_list = database.get(trakt.json_response, 24, 'networks')

        if trakt_list is None:
            return
        list_items = []
        for i in trakt_list:
            list_items.append(tools.addDirectoryItem(i['name'], 'showsNetworkShows&actionArgs=%s&page=1' % i['name'],
                                                     '', '', bulk_add=True))
        tools.addMenuItems(syshandle, list_items, len(list_items))
        tools.closeDirectory('addons')

    def showsNetworkShows(self, network, page):

        trakt_list = database.get(trakt.json_response, 24, 'shows/popular?networks=%s&page=%s' % (network, page))

        if trakt_list is None:
            return

        self.showListBuilder(trakt_list)

        if len(trakt_list) == int(tools.getSetting('item.limit')):
            tools.addDirectoryItem(tools.lang(32019), 'showsNetworkShows&actionArgs=%s&page=%s' %
                                   (network, int(page) + 1), '', '')

        tools.closeDirectory('tvshows')

    def showsPopular(self, page):
        trakt_list = database.get(trakt.json_response, 12, 'shows/popular?page=%s' % page)

        if trakt_list is None:
            return

        self.showListBuilder(trakt_list)
        tools.addDirectoryItem(tools.lang(32019), 'showsPopular&page=%s' % (int(page) + 1), '', '')
        tools.closeDirectory('tvshows')

    def showsRecommended(self):
        trakt_list = database.get(trakt.json_response, 12, 'recommendations/shows?ignore_collected=true',
                                  limit=True, limitOverride=100)
        if trakt_list is None:
            return
        self.showListBuilder(trakt_list)
        tools.closeDirectory('tvshows')

    def showsTrending(self, page):
        trakt_list = database.get(trakt.json_response, 12, 'shows/trending?page=%s' % page)
        if trakt_list is None:
            return
        self.showListBuilder(trakt_list)
        tools.addDirectoryItem(tools.lang(32019), 'showsTrending&page=%s' % (int(page) + 1), '', '')
        tools.closeDirectory('tvshows')

    def showsPlayed(self, page):
        trakt_list = database.get(trakt.json_response, 12, 'shows/played?page=%s' % page)
        if trakt_list is None:
            return
        self.showListBuilder(trakt_list)
        tools.addDirectoryItem(tools.lang(32019), 'showsPlayed&page=%s' % (int(page) + 1), '', '')
        tools.closeDirectory('tvshows')

    def showsWatched(self, page):
        trakt_list = database.get(trakt.json_response, 12, 'shows/watched?page=%s' % page)
        if trakt_list is None:
            return
        self.showListBuilder(trakt_list)
        tools.addDirectoryItem(tools.lang(32019), 'showsWatched&page=%s' % (int(page) + 1), '', '')
        tools.closeDirectory('tvshows')

    def showsCollected(self, page):
        trakt_list = database.get(trakt.json_response, 12, 'shows/collected?page=%s' % page)
        if trakt_list is None:
            return
        self.showListBuilder(trakt_list)
        tools.addDirectoryItem(tools.lang(32019), 'showsCollected&page=%s' % (int(page) + 1), '', '')
        tools.closeDirectory('tvshows')

    def showsAnticipated(self, page):
        trakt_list = database.get(trakt.json_response, 12, 'shows/anticipated?page=%s&language=%s'
                                  % (page, language_code))
        if trakt_list is None:
            return
        self.showListBuilder(trakt_list)
        tools.addDirectoryItem(tools.lang(32019), 'showsAnticipated&page=%s' % (int(page) + 1), '', '')
        tools.closeDirectory('tvshows')

    def showsUpdated(self, page):
        import datetime
        date = datetime.date.today() - datetime.timedelta(days=31)
        date = date.strftime('%Y-%m-%d')
        trakt_list = database.get(trakt.json_response, 12, 'shows/updates/%s?page=%s' % (date, page))
        if trakt_list is None:
            return
        self.showListBuilder(trakt_list)
        tools.addDirectoryItem(tools.lang(32019), 'showsUpdated&page=%s' % (int(page) + 1), '', '')
        tools.closeDirectory('tvshows')

    def showSearchHistory(self):
        history = database.getSearchHistory('show')
        tools.addDirectoryItem(tools.lang(40142), 'showsSearch', '', '')
        tools.addDirectoryItem(tools.lang(40140), 'clearSearchHistory', '', '', isFolder=False)

        for i in history:
            tools.addDirectoryItem(i, 'showsSearch&actionArgs=%s' % tools.quote(i), '', '')
        tools.closeDirectory('addon')

    def showsSearch(self, actionArgs=None):

        if actionArgs == None:
            k = tools.showKeyboard('', tools.lang(32016))
            k.doModal()
            query = (k.getText() if k.isConfirmed() else None)
            if query == None or query == '':
                return
        else:
            query = actionArgs
        database.addSearchHistory(query, 'show')
        query = tools.deaccentString(tools.display_string(query))
        tools.quote(query)
        tools.closeAllDialogs()
        tools.closeDirectory('tvshows')
        tools.execute("Container.Update(%s?action=showsSearchResults&actionArgs=%s, replace)'" % (sysaddon, query))

    def showsSearchResults(self, query):
        query = tools.quote_plus(tools.unquote(query))
        trakt_list = trakt.json_response('search/show?query=%s&extended=full&type=show&field=title' % query)
        if trakt_list is None:
            return
        self.showListBuilder(trakt_list)
        tools.closeDirectory('tvshows')

    def showsByActor(self, actionArgs):

        if actionArgs == None:
            k = tools.showKeyboard('', tools.lang(32016))
            k.doModal()
            query = (k.getText() if k.isConfirmed() else None)
            if query == None or query == '':
                return
        else:
            query = tools.unquote(actionArgs)

        database.addSearchHistory(query, 'showActor')
        query = tools.deaccentString(query)
        query = query.replace(' ', '-')
        query = tools.quote_plus(query)

        trakt_list = trakt.json_response('people/%s/shows' % query, limit=True)

        try:
            trakt_list = trakt_list['cast']
        except:
            import traceback
            traceback.print_exc()
            trakt_list = []

        trakt_list = [i['show'] for i in trakt_list]

        self.showListBuilder(trakt_list)

        tools.closeDirectory('tvshows')

    def showSeasons(self, args):

        args = tools.get_item_information(args)

        self.seasonListBuilder(args['ids']['trakt'])

        tools.closeDirectory('seasons')

    def seasonEpisodes(self, args):

        args = tools.get_item_information(args)

        show_id = args['showInfo']['ids']['trakt']

        if 'seasonInfo' in args:
            season_number = args['seasonInfo']['info']['season']
        else:
            season_number = args['info']['season']

        self.episodeListBuilder(show_id, season_number)
        tools.closeDirectory('episodes', sort='episode')

    def showGenres(self):
        tools.addDirectoryItem(tools.lang(32065), 'showGenresGet', '', '', isFolder=True)
        genres = database.get(trakt.json_response, 24, 'genres/shows')

        if genres is None:
            return

        for i in genres:
            tools.addDirectoryItem(i['name'], 'showGenresGet&actionArgs=%s' % i['slug'], '', '', isFolder=True)
        tools.closeDirectory('addons')

    def showGenreList(self, args, page):
        if page is None:
            page = 1

        if args is None:
            genre_display_list = []
            genre_string = ''
            genres = database.get(trakt.json_response, 24, 'genres/shows')

            for genre in genres:
                genre_display_list.append(genre['name'])
            genre_multiselect = tools.showDialog.multiselect(tools.addonName + ": Genre Selection", genre_display_list)

            if genre_multiselect is None: return
            for selection in genre_multiselect:
                genre_string += ', %s' % genres[selection]['slug']
            genre_string = genre_string[2:]

        else:
            genre_string = args

        page = int(page)
        trakt_list = database.get(trakt.json_response, 12,
                                  'shows/popular?genres=%s&page=%s' % (genre_string, page))
        if trakt_list is None:
            return

        self.showListBuilder(trakt_list)
        tools.addDirectoryItem(tools.lang(32019),
                               'showGenresGet&actionArgs=%s&page=%s' % (genre_string, page + 1), None, None)
        tools.closeDirectory('tvshows')

    def showsRelated(self, args):
        trakt_list = database.get(trakt.json_response, 12, 'shows/%s/related' % args)
        if trakt_list is None:
            return

        self.showListBuilder(trakt_list)
        tools.closeDirectory('tvshows')

    def showYears(self, year=None, page=None):
        if year is None:
            current_year = int(tools.datetime_workaround(datetime.datetime.today().strftime('%Y-%m-%d')).year)
            all_years = reversed([year for year in range(1900, current_year+1)])
            menu_items = []
            for year in all_years:
                menu_items.append(tools.addDirectoryItem(str(year), 'showYears&actionArgs=%s' % year, '', '',
                                  bulk_add=True))
            tools.addMenuItems(syshandle, menu_items, len(menu_items))
            tools.closeDirectory('tvshows')
        else:
            if page is None:
                page = 1

            trakt_list = trakt.json_response('shows/popular?years=%s&page=%s' % (year, page))
            self.showListBuilder(trakt_list)
            tools.addDirectoryItem(tools.lang(32019),
                                   'showYears&actionArgs=%s&page=%s' % (year, int(page) + 1), None, None)
            tools.closeDirectory('tvshows')

    ######################################################
    # MENU TOOLS
    ######################################################

    def seasonListBuilder(self, show_id, smartPlay=False):

        self.itemList = trakt_database.get_season_list(show_id)

        self.itemList = [x for x in self.itemList if x is not None and 'info' in x]

        self.itemList = sorted(self.itemList, key=lambda k: k['info']['season'])

        if len(self.itemList) == 0:
            tools.log('We received no titles to build a list', 'error')
            return

        hide_specials = False

        if tools.getSetting('general.hideSpecials') == 'true':
            hide_specials = True

        item_list = []

        for item in self.itemList:

            try:
                if hide_specials and int(item['info']['season']) == 0:
                    continue

                action = 'seasonEpisodes'

                args = {'trakt_id': item['showInfo']['ids']['trakt'],
                        'season': item['info']['season'],
                        'item_type': 'season'}

                args = tools.quote(json.dumps(args, sort_keys=True))

                item['trakt_object']['show_id'] = item['showInfo']['ids']['trakt']
                name = item['info']['season_title']

                if not self.is_aired(item['info']) or 'aired' not in item['info']:
                    if tools.getSetting('general.hideUnAired') == 'true':
                        continue
                    name = tools.colorString(name, 'red')
                    name = tools.italic_string(name)
                    item['info']['title'] = name

                item['info'] = tools.clean_air_dates(item['info'])

            except:
                import traceback
                traceback.print_exc()
                continue

            if smartPlay is True:
                return args

            cm = []
            if tools.getSetting('trakt.auth') != '':
                cm.append(('Trakt Manager', 'RunPlugin(%s?action=traktManager&actionArgs=%s)' % (sysaddon, args)))

            if tools.context_addon():
                cm = []

            item_list.append(tools.addDirectoryItem(name, action, item['info'], item['art'], cm=cm, isFolder=True,
                                                    isPlayable=False, actionArgs=args, set_ids=item['ids'],
                                                    bulk_add=True))

        tools.addMenuItems(syshandle, item_list, len(item_list))

    def episodeListBuilder(self, show_id, season_number, smartPlay=False, hide_unaired=False):

        try:
            item_list = []

            self.itemList = trakt_database.get_season_episodes(show_id, season_number)
            self.itemList = [x for x in self.itemList if x is not None and 'info' in x]

            if len(self.itemList) == 0:
                tools.log('We received no titles to build a list', 'error')
                return

            try:
                self.itemList = sorted(self.itemList, key=lambda k: k['info']['episode'])
            except:
                pass

            for item in self.itemList:

                cm = []

                try:

                    if tools.getSetting('smartplay.playlistcreate') == 'true' and smartPlay is False:
                        action = 'smartPlay'
                        playable = False
                    else:
                        playable = True
                        action = 'getSources'

                    args = {'trakt_id': item['showInfo']['ids']['trakt'],
                            'season': item['info']['season'],
                            'episode': item['info']['episode'],
                            'item_type': 'episode'}

                    args = tools.quote(json.dumps(args, sort_keys=True))

                    name = item['info']['title']

                    if not self.is_aired(item['info']):
                        if tools.getSetting('general.hideUnAired') == 'true' or hide_unaired:
                            continue
                        else:
                            name = tools.colorString(name, 'red')
                            name = tools.italic_string(name)
                            item['info']['title'] = name

                    item['info'] = tools.clean_air_dates(item['info'])

                except:
                    import traceback
                    traceback.print_exc()
                    continue
                cm.append((tools.lang(32070),
                           'XBMC.PlayMedia(%s?action=shufflePlay&actionArgs=%s)' % (sysaddon, args)))

                cm.append(('Browse Season',
                           'XBMC.Container.Update(%s?action=seasonEpisodes&actionArgs=%s)' %
                           (sysaddon,
                            tools.quote(json.dumps({'trakt_id': item['showInfo']['ids']['trakt'],
                                                    'season': item['info']['season'],
                                                    'item_type': 'season'})))))

                cm.append((tools.lang(33022),
                           'PlayMedia(%s?action=getSources&seren_reload=true&actionArgs=%s)' % (sysaddon, args)))

                cm.append((tools.lang(32066),
                           'PlayMedia(%s?action=getSources&source_select=true&actionArgs=%s)' % (sysaddon, args)))

                if tools.getSetting('trakt.auth') != '':
                    cm.append(('Trakt Manager', 'RunPlugin(%s?action=traktManager&actionArgs=%s)' % (sysaddon, args)))

                if tools.context_addon():
                    cm = []

                if tools.getSetting('premiumize.enabled') == 'true' and tools.getSetting('premiumize.pin') != '':
                    cm.append((tools.lang(32068),
                               'XBMC.RunPlugin(%s?action=filePicker&actionArgs=%s)' % (sysaddon, args)))

                item_list.append(tools.addDirectoryItem(name, action, item['info'], item['art'], isFolder=False,
                                                        isPlayable=playable, actionArgs=args, bulk_add=True,
                                                        set_ids=item['ids'], cm=cm))

            if smartPlay is True:
                return item_list
            else:
                tools.addMenuItems(syshandle, item_list, len(item_list))

        except:
            import traceback
            traceback.print_exc()

    def mixedEpisodeBuilder(self, trakt_list, sort=None, hide_watched=False, smartPlay=False, hide_unaired=True,
                            prepend_date=False):

        self.threadList = []

        try:
            if len(trakt_list) == 0:
                tools.log('We received no titles to build a list', 'error')
                return

            self.itemList = trakt_database.get_episode_list(trakt_list)

            self.itemList = [x for x in self.itemList if x is not None and 'info' in x]
            self.itemList = [i for i in self.itemList if 'info' in i and i['info'].get('premiered', None) is not None]
            if sort is None:
                self.itemList = sorted(self.itemList,
                                       key=lambda i: tools.datetime_workaround(i['info']['premiered'],
                                                                               tools.trakt_gmt_format, False),
                                       reverse=True)
            elif sort is not False:
                sort_list = []
                for trakt_id in sort['id_list']:
                    try:
                        if not sort['type']:
                            item = [i for i in self.itemList if i['ids']['trakt'] == trakt_id][0]
                        else:
                            item = [i for i in self.itemList if i[sort['type']]['ids']['trakt'] == trakt_id][0]
                        sort_list.append(item)
                    except IndexError:
                        continue
                    except:
                        import traceback
                        traceback.print_exc()
                self.itemList = sort_list

            item_list = []

            for item in self.itemList:
                if item is None:
                    continue

                if item['info'].get('title', '') == '':
                    continue

                if hide_watched and item['info']['playcount'] != 0:
                    continue

                cm = []

                try:
                    name = tools.display_string(item['info']['title'])

                    if not self.is_aired(item['info']) and hide_unaired is True:
                        continue
                    elif not self.is_aired(item['info']):
                        name = tools.colorString(name, 'red')
                        name = tools.italic_string(name)
                        item['info']['title'] = name

                    item['info'] = tools.clean_air_dates(item['info'])

                    args = {'trakt_id': item['showInfo']['ids']['trakt'],
                            'season': item['info']['season'],
                            'episode': item['info']['episode'],
                            'item_type': 'episode'}

                    args = tools.quote(json.dumps(args, sort_keys=True))

                    if tools.getSetting('smartplay.playlistcreate') == 'true' and smartPlay is False:
                        action = 'smartPlay'
                        playable = False
                    else:
                        playable = True
                        action = 'getSources'

                    if self.title_appends == 'true':
                        name = "%s: %sx%s %s" % (tools.colorString(item['showInfo']['info']['tvshowtitle']),
                                                 tools.display_string(item['info']['season']).zfill(2),
                                                 tools.display_string(item['info']['episode']).zfill(2),
                                                 tools.display_string(item['info']['title']))
                    if prepend_date:
                        release_day = tools.datetime_workaround(item['info']['aired'])
                        release_day = release_day.strftime('%d %b')
                        name = '[%s] %s' % (release_day, name)

                    cm.append((tools.lang(32069),
                               'XBMC.Container.Update(%s?action=showSeasons&actionArgs=%s)' %
                               (sysaddon, tools.quote(json.dumps({'trakt_id': item['showInfo']['ids']['trakt'],
                                                     'item_type': 'show'})))))

                    cm.append(('Browse Season',
                               'XBMC.Container.Update(%s?action=seasonEpisodes&actionArgs=%s)' %
                               (sysaddon,
                                tools.quote(json.dumps({'trakt_id': item['showInfo']['ids']['trakt'],
                                                        'season': item['info']['season'],
                                                        'item_type': 'season'})))))

                    cm.append((tools.lang(32070),
                               'XBMC.PlayMedia(%s?action=shufflePlay&actionArgs=%s)' % (sysaddon, args)))

                    cm.append((tools.lang(32066),
                               'PlayMedia(%s?action=getSources&source_select=true&actionArgs=%s)' % (sysaddon, args)))

                    cm.append((tools.lang(33022),
                               'PlayMedia(%s?action=getSources&seren_reload=true&actionArgs=%s)' % (sysaddon, args)))

                    if tools.getSetting('trakt.auth') != '':
                        cm.append(('Trakt Manager', 'RunPlugin(%s?action=traktManager&actionArgs=%s)'
                                   % (sysaddon, tools.quote(json.dumps(item['trakt_object'])))))

                    if tools.context_addon():
                        cm = []

                    if tools.getSetting('premiumize.enabled') == 'true' and tools.getSetting('premiumize.pin') != '':
                        cm.append((tools.lang(32068),
                                   'XBMC.RunPlugin(%s?action=filePicker&actionArgs=%s)' % (sysaddon, args)))

                    item['info']['title'] = item['info']['originaltitle'] = name

                    item_list.append(tools.addDirectoryItem(name, action, item['info'], item['art'], isFolder=False,
                                                            isPlayable=playable, actionArgs=args, bulk_add=True,
                                                            set_ids=item['ids'], cm=cm))

                except:
                    import traceback
                    traceback.print_exc()
                    continue

            if smartPlay is True:
                return item_list
            else:
                tools.addMenuItems(syshandle, item_list, len(item_list))

        except:
            import traceback
            traceback.print_exc()

    def showListBuilder(self, trakt_list, forceResume=False, info_only=False):

        try:
            if len(trakt_list) == 0:
                tools.log('We received no titles to build a list', 'error')
                return
        except:
            import traceback
            traceback.print_exc()
            return

        if 'show' in trakt_list[0]:
            trakt_list = [i['show'] for i in trakt_list]

        show_ids = [i['ids']['trakt'] for i in trakt_list]

        self.itemList = trakt_database.get_show_list(show_ids)
        self.itemList = [x for x in self.itemList if x is not None and 'info' in x]
        self.itemList = tools.sort_list_items(self.itemList, trakt_list)

        item_list = []

        for item in self.itemList:
            try:
                # Add Arguments to pass with items
                args = {'trakt_id': item['ids']['trakt'], 'item_type': 'show'}
                args = tools.quote(json.dumps(args, sort_keys=True))

                cm = []

                name = tools.display_string(item['info']['tvshowtitle'])

                if info_only == True:
                    return args

                if not self.is_aired(item['info']):
                    if tools.getSetting('general.hideUnAired') == 'true':
                        continue
                    name = tools.colorString(name, 'red')
                    name = tools.italic_string(name)

                item['info'] = tools.clean_air_dates(item['info'])

                if 'setCast' in item:
                    set_cast = item['setCast']
                else:
                    set_cast = False

                if tools.getSetting('smartplay.clickresume') == 'true' or forceResume is True:
                    action = 'playbackResume'
                else:
                    action = 'showSeasons'

                # Context Menu Items

                cm.append((tools.lang(32070),
                           'XBMC.PlayMedia(%s?action=shufflePlay&actionArgs=%s)' % (sysaddon, args)))

                cm.append((tools.lang(32020),
                           'Container.Update(%s?action=showsRelated&actionArgs=%s)' % (sysaddon, item['ids']['trakt'])))

                cm.append((tools.lang(32069),
                           'XBMC.Container.Update(%s?action=showSeasons&actionArgs=%s)' % (sysaddon, args)))

                if tools.getSetting('trakt.auth') != '':
                    cm.append(('Trakt Manager', 'RunPlugin(%s?action=traktManager&actionArgs=%s)' % (sysaddon, args)))

                cm.append((tools.lang(40153),
                           'XBMC.PlayMedia(%s?action=playFromRandomPoint&actionArgs=%s' % (sysaddon, args)))

                if tools.context_addon():
                    cm = []

            except:
                import traceback
                traceback.print_exc()
                continue

            item_list.append(tools.addDirectoryItem(name, action, item['info'], item['art'], cm=cm, isFolder=True,
                                                    isPlayable=False, actionArgs=args, bulk_add=True, set_cast=set_cast,
                                                    set_ids=item['ids']))

        tools.addMenuItems(syshandle, item_list, len(item_list))

    def runThreads(self, join=True):
        for thread in self.threadList:
            thread.start()

        if join == True:
            for thread in self.threadList:
                thread.join()

    def _start_queue_workers(self):

        self.queue_finished = False

        for i in range(40):
            self.threadList.append(Thread(target=self._queue_worker))

        for i in self.threadList:
            i.start()

    def _finish_queue_workers(self):

        self.queue_finished = True

        for i in self.threadList:
            i.join()

        self.threadList = []

    def _queue_worker(self):
        while not self.task_queue.empty() or not self.queue_finished:
            try:
                target = self.task_queue.get(timeout=3)
            except:
                continue
            try:
                target[0](*target[1])
            except:
                import traceback
                traceback.print_exc()
                pass

    def is_aired(self, info):
        try:
            try:air_date = info['aired']
            except: air_date = info.get('premiered')
            if air_date == '' or air_date is None:
                return False
            if int(air_date[:4]) < 1970:
                return True

            time_format = tools.trakt_gmt_format
            if len(air_date) == 10:
                time_format = '%Y-%m-%d'

            air_date = tools.gmt_to_local(air_date, format=time_format)

            if tools.getSetting('general.datedelay') == 'true':
                air_date = tools.datetime_workaround(air_date, time_format, False)
                air_date += datetime.timedelta(days=1)
            else:
                air_date = tools.datetime_workaround(air_date, time_format, False)

            if air_date > datetime.datetime.now():
                return False
            else:
                return True
        except:
            import traceback
            traceback.print_exc()
            # Assume an item is not aired if we do not have any information on it or fail to identify
            return False


import webapp2
import jinja2
from google.appengine.api import users
from google.appengine.ext import ndb
import os
from snippets import Words
# from Tkinter import *

JINJA_ENVIRONMENT=jinja2.Environment(loader=jinja2.FileSystemLoader(os.path.dirname(__file__)),
extensions=['jinja2.ext.autoescape'],
autoescape=True
)

class Add(webapp2.RequestHandler):
    def split(self, word):
        return [char for char in word]
    def merge(self, word):
        new=''
        for x in word:
            new += x
        return new
    def sort(self, word):
        split_data =self.split(word)
        # print(split_data)
        sorted_alphabets = sorted(split_data)
        # print(sorted_alphabets)
        merged_word = self.merge(sorted_alphabets)
        return merged_word
    # def incrementCounter(self, user)
    def get(self):
        self.response.headers['Content-Type'] = 'text/html'
        display_message='Add words in the Anagram Engine'
        Template_values ={
        'display_message': display_message
        }
        template = JINJA_ENVIRONMENT.get_template('add.html')
        self.response.write(template.render(Template_values))

    def post(self):
        action = self.request.get('add_button')
        # print (action)
        user = users.get_current_user()

        # print(myuser)
        add_string=self.request.get('word_input')

        sorted_alphabets = sorted(self.split(add_string.lower()))
        keyword=user.user_id() + self.merge(sorted_alphabets)

        if action == 'Add':
            key = ndb.Key('Words', keyword)
            word = key.get()
            myuser_key= ndb.Key('MyUser', user.user_id())
            myuser=myuser_key.get()
            if word == None:
                word = Words(id=keyword)
                word.count_of_words=0
                word.put()
                myuser.uniqueAnagramCounter=myuser.uniqueAnagramCounter+1
                myuser.put()
            string = keyword
            if string == '' or string == None or len(string)<3:
                self.redirect('/add')
                return
            word_doesnt_exists = True
            List = []
            for i in word.wordsList:
                print(i)
                List.append(i)
            print(List.count(add_string))
            if List.count(add_string.lower())>0:
                word_doesnt_exists=False
                print('word exists')
            if(word_doesnt_exists):
                word.wordsList.append(add_string.lower())
                word.count_of_words=word.count_of_words+1
                word.alphabet_no_List.append(len(add_string))
                word.put()

                myuser.wordCounter= myuser.wordCounter+1
                myuser.put()
        # Code to read from text document
        # root = Tk()
        # root.fileName = filedialog.askopenfilename(filetypes =(("All text file", "*.txt")))
        dict = []
        f = open("words.txt", "r")
        for line in f.readlines():
            # sorted_word_from_text = self.merge(sorted(self.split(line.strip())))
            #
            if(line.rstrip()):
            #     print (sorted_word_from_text)
                dict.append(line.rstrip())

        # user = users.get_current_user()
        file_action = self.request.get('add_from_files')

        if file_action=='Add':
            print(len(dict))
            for i in dict:
                if len(i)>0:
                    keyword1=user.user_id() + self.sort(i)

                    key = ndb.Key('Words', keyword1)
                    word = key.get()
                    # print(word)
                    new_word = False
                    if word!=None:
                        if word.wordsList.count(i)==0:
                            word.wordsList.append(i)
                            word.count_of_words=word.count_of_words+1
                            word.alphabet_no_List.append(len(i))
                            word.put()

                        # word.wordsList.append(i)
                        # word.count_of_words=word.count_of_words+1
                        # word.alphabet_no_List.append(len(i))
                        # word.put()
                    else:
                        new_word=True
                    if(new_word):
                        word = Words(id=keyword1)
                        print(i + " word is added")
                        word.wordsList.append(i)
                        word.count_of_words=1
                        word.alphabet_no_List.append(len(i))
                        word.put()
                        print(i)
                    self.redirect('/add')

        self.redirect('/add')
        # raw_word =self.request.get('word_input')
        # sorted_alphabets = sorted(self.split(raw_word.lower()))
        # user = users.get_current_user()
        # # used as a key to display only certain content to certain
        # keyword =user.user_id()+merge(sorted_alphabets)
        # print(keyword)
        # # if action =='Add':
        #
        #
        #     # use user_id()+ keyword as the key.
        # word_key = ndb.Key('Words', keyword)
        # word = word_key.get()
        #
        # if word==None:
        #     # word = Word(id=keyword)
        #     word.word = raw_word
        #     word.count_of_alphabets = len(raw_word)
        #     word.count_of_words = 1
        #     word.put()
        #     word.word.append(raw_word)
        #     word.put()
        #     self.redirect('/add')
        #
        # else:
        #     word.word.append(raw_word)
        #     word.count_of_alphabets = len(raw_word)
        #     countOfWords = word.count_of_words
        #     word.count_of_words=countOfWords+1
        #     word.put()
        #     self.redirect('/add')


import pycosat
from pprint import pprint

# number of cells in sudoku
NUM_CELLS = 81

# this can be used to interate all squares subset in 3x 3
#    V      V      V
#    1 2 3  4 5 6  7 8 9
#> 1 |0 0 0| 0 0 0| 0 0 0|
#  2 |0 0 0| 0 0 0| 0 0 0|
#  3 |0 0 0| 0 0 0| 0 0 0|
#    ---------------------
#> 4 |0 0 0| 0 0 0| 0 0 0|
#  5 |0 0 0| 0 0 0| 0 0 0|
#  6 |0 0 0| 0 0 0| 0 0 0|
#    --------------------
#> 7 |0 0 0| 0 0 0| 0 0 0|
#  8 |0 0 0| 0 0 0| 0 0 0|
#  9 |0 0 0| 0 0 0| 0 0 0|
SUBSQUARE_BOUDARIES = [1, 4, 7]

# most digit in sudoku
NUM_DIGITS = 9

def get_cell_value(row, column, digit):
   """
     make unique id to a cell

     first number is a row
     second is a column
     third is a digit value

     Ex: the cell (1,3) and (3,1) with digit 7 there are diferentes id
         cell (1,3) = 137
         cell (3,1) = 317
   """
   return row*100+column*10+ digit

def get_base_clauses():
    """
    make all NUM_VARIABLES used in sudoku board, named by id

    Ex:
        for cell (1,1) and digit 7 can be named: 117
        for cell (2,3) and digit 5 can be named: 235
    """
    base_clauses = []

    for row in range(1, NUM_DIGITS+1):
       for column in range(1, NUM_DIGITS+1):
          clauses = []
          for digit in range(1, NUM_DIGITS+1):
             clauses.append(get_cell_value(row,column,digit))
          base_clauses.append(clauses)

    return base_clauses

def get_unique_cells_clauses():
    """
    make clauses guarantee that a cell can just appear once for sudoku board.

    to make this each cell will have the next cell with the clause:
      ~current_digit or ~next_digit

    Example:
        to cell 111 there the clauses:
        (-111,-112),(-111,-113),(-111,-114),...,(-111,-999)
    """
    unique_digits_clauses = []

    for row in range(1, NUM_DIGITS+1):
        for column in range(1,NUM_DIGITS+1):
            for digit in range(1, NUM_DIGITS+1):
                for next_digit in range(digit+1, NUM_DIGITS+1):
                    cell_id = -get_cell_value(row,column,digit)
                    next_cell_id = -get_cell_value(row,column,next_digit)
                    unique_digits_clauses.append([cell_id,next_cell_id])

    return unique_digits_clauses

def get_unique_subset_clauses(board_subset):
    """
    this guarantee that a cell appear only once in the board subset
    """
    subset_clauses = []
    subset = enumerate(board_subset)

    for index, first_tuple in enumerate(board_subset):
        for n_index, n_tuple in enumerate(board_subset):
            if index < n_index:
                for digit in range(1, NUM_DIGITS + 1):
                    clause = [-get_cell_value(
                        first_tuple[0], first_tuple[1], digit),
                        -get_cell_value(
                        n_tuple[0], n_tuple[1], digit)]
                    subset_clauses.append(clause)
    return subset_clauses

def get_row_unique_clauses():
    """
    this guarantee that a cell in row appear only once in the row
    """
    unique_clauses = []

    for row in range(1,NUM_DIGITS +1):
        subset = []
        for column in range(1, NUM_DIGITS+1):
            subset.append((row,column))
        unique_clauses.extend(get_unique_subset_clauses(subset))
    return unique_clauses

def get_columns_unique_clauses():
    """
    this guarantee that a cell in column appear only once in the columns
    """
    unique_clauses = []
    for row in range(1,NUM_DIGITS +1):
        subset = []
        for column in range(1, NUM_DIGITS+1):
            subset.append((column,row))
        unique_clauses.extend(get_unique_subset_clauses(subset))
    return unique_clauses

def get_square_unique_clauses():
    """
    this guarantee that a cell in square appear only once in the squares
    """
    subset_clauses = []

    for row in SUBSQUARE_BOUDARIES:
        for column in SUBSQUARE_BOUDARIES:
            subset = [] # make subset 3x3
            for k in range(9):
                subset.append((row+k%3,column+k//3))
            subset_clauses.extend(get_unique_subset_clauses(subset))

    return subset_clauses

def get_sudoku_clauses():
    """
    mix all defined clauses to guarantee a valid sudoku
    """
    sudoku_clauses = []
    sudoku_clauses.extend(get_base_clauses())
    sudoku_clauses.extend(get_unique_cells_clauses())
    sudoku_clauses.extend(get_row_unique_clauses())
    sudoku_clauses.extend(get_columns_unique_clauses())
    sudoku_clauses.extend(get_square_unique_clauses())

    return sudoku_clauses

def get_single_clauses(sudoku_board):
    """
    This method make a clauses that can be answer true
    to solve the sudoku board
    """
    single_clauses = []
    for row in range(1, NUM_DIGITS+1):
        for column in range(1,NUM_DIGITS+1):
            cell_value = sudoku_board[row-1][column-1]
            if cell_value:
                single_clauses.append(
                    [get_cell_value(row,column,cell_value)])

    return single_clauses

def get_cell_solution(sudoku_solution, row, column):
    """
    verify a cell id in a sudoku solution list
    """
    for digit in range(1, NUM_DIGITS+1):
        if get_cell_value(row,column,digit) in sudoku_solution:
            return digit
    return -1

def solve_sudoku(sudoku_board):
    """
    Generate a sudoku clauses, apply in a pycosat and get sudoku solution
    """
    sudoku_clauses = get_sudoku_clauses()
    single_clauses = get_single_clauses(sudoku_board)
    sudoku_clauses.extend(single_clauses)

    sudoku_solution = set(pycosat.solve(sudoku_clauses))

    for row in range(1, NUM_DIGITS+1):
        for column in range(1, NUM_DIGITS+1):
            sudoku_board[row-1][column-1] = get_cell_solution(
                sudoku_solution, row, column)
    return sudoku_board

def main():
    print ("Sudoku problem:")
    sudoku_problem = [[0, 2, 0, 0, 0, 0, 0, 0, 0],
                      [0, 0, 0, 6, 0, 0, 0, 0, 3],
                      [0, 7, 4, 0, 8, 0, 0, 0, 0],
                      [0, 0, 0, 0, 0, 3, 0, 0, 2],
                      [0, 8, 0, 0, 4, 0, 0, 1, 0],
                      [6, 0, 0, 5, 0, 0, 0, 0, 0],
                      [0, 0, 0, 0, 1, 0, 7, 8, 0],
                      [5, 0, 0, 0, 0, 9, 0, 0, 0],
                      [0, 0, 0, 0, 0, 0, 0, 4, 0]]
    pprint(sudoku_problem)

    print('\nGenerating solution:')
    sudoku_solution = solve_sudoku(sudoku_problem)
    pprint(sudoku_solution)

if __name__ == '__main__':
    main()


from typing import Any, Dict, Iterable

import pandas as pd

from fugue import DataFrame, FugueWorkflow, PandasDataFrame, out_transform, transform
from fugue.constants import FUGUE_CONF_WORKFLOW_CHECKPOINT_PATH

def test_transform():
    pdf = pd.DataFrame([[1, 10], [0, 0], [1, 1], [0, 20]], columns=["a", "b"])

    def f1(df: pd.DataFrame) -> pd.DataFrame:
        return df.sort_values("b").head(1)

    result = transform(pdf, f1, schema="*")
    assert isinstance(result, pd.DataFrame)
    assert result.values.tolist() == [[0, 0]]

    # schema: *
    def f2(df: pd.DataFrame) -> pd.DataFrame:
        return df.sort_values("b").head(1)

    result = transform(pdf, f2)
    assert isinstance(result, pd.DataFrame)
    assert result.values.tolist() == [[0, 0]]

    result = transform(pdf, f2, partition=dict(by=["a"]))
    assert isinstance(result, pd.DataFrame)
    assert sorted(result.values.tolist(), key=lambda x: x[0]) == [[0, 0], [1, 1]]
    result = transform(
        pdf, f2, partition=dict(by=["a"]), force_output_fugue_dataframe=True
    )
    assert isinstance(result, DataFrame)

    ppdf = PandasDataFrame(pdf)
    assert isinstance(transform(ppdf, f2), DataFrame)

    # schema: *
    def f3(df: pd.DataFrame, called: callable) -> pd.DataFrame:
        called()
        return df

    cb = Callback()
    result = transform(pdf, f3, callback=cb.called)
    assert 1 == cb.ct

def test_transform_from_yield(tmpdir):
    # schema: *,x:int
    def f(df: pd.DataFrame) -> pd.DataFrame:
        return df.assign(x=1)

    dag = FugueWorkflow()
    dag.df([[0]], "a:int").yield_dataframe_as("x1")
    dag.df([[1]], "b:int").yield_dataframe_as("x2")
    dag.run("", {FUGUE_CONF_WORKFLOW_CHECKPOINT_PATH: str(tmpdir)})

    result = transform(dag.yields["x1"], f)
    assert isinstance(result, DataFrame)
    assert result.as_array(type_safe=True) == [[0, 1]]

    result = transform(
        dag.yields["x2"],
        f,
        engine_conf={FUGUE_CONF_WORKFLOW_CHECKPOINT_PATH: str(tmpdir)},
    )
    assert isinstance(result, DataFrame)
    assert result.as_array(type_safe=True) == [[1, 1]]

def test_out_transform(tmpdir):
    pdf = pd.DataFrame([[1, 10], [0, 0], [1, 1], [0, 20]], columns=["a", "b"])

    class T:
        def __init__(self):
            self.n = 0

        def f(self, df: Iterable[Dict[str, Any]]) -> None:
            self.n += 1

    t = T()
    out_transform(pdf, t.f)
    assert 1 == t.n

    t = T()
    out_transform(pdf, t.f, partition=dict(by=["a"]))
    assert 2 == t.n

    dag = FugueWorkflow()
    dag.df(pdf).yield_dataframe_as("x1")
    dag.df(pdf).yield_dataframe_as("x2")
    dag.run("", {FUGUE_CONF_WORKFLOW_CHECKPOINT_PATH: str(tmpdir)})

    t = T()
    out_transform(dag.yields["x1"], t.f)
    assert 1 == t.n

    t = T()
    out_transform(
        dag.yields["x2"],
        t.f,
        partition=dict(by=["a"]),
        engine_conf={FUGUE_CONF_WORKFLOW_CHECKPOINT_PATH: str(tmpdir)},
    )
    assert 2 == t.n

    # schema: *
    def f3(df: pd.DataFrame, called: callable) -> pd.DataFrame:
        called()
        return df

    cb = Callback()
    result = out_transform(pdf, f3, callback=cb.called)
    assert 1 == cb.ct

class Callback:
    def __init__(self):
        self.ct = 0

    def called(self) -> None:
        self.ct += 1


from django.test import TestCase
from django.contrib.auth import get_user_model
from django.urls import reverse

from rest_framework import status
from rest_framework.test import APIClient

from core.models import *
from recipe.serializers import RecipeSerializer, RecipeDetailSerializer

RECIPE_URL = reverse('recipe:recipe-list')

# /api/recipe/recipes
# /api/recipe/recipes/1/

def detail_url(recipe_id):
    """return recipe detail url"""
    return reverse('recipe:recipe-detail', args=[recipe_id])

def sample_tag(user, name='Main Course'):
    """Create and return sample tag"""
    return Tag.objects.create(user=user, name=name)

def sample_ingredient(user, name='Cinnamon'):
    """Create and return sample ingredient"""
    return Ingredient.objects.create(user=user, name=name)

def sample_recipe(user,**params):
    """Create and return sample recipe"""

    test_recipe = {
        'title': 'Mushroom Chicken',
        'time_minutes': 10,
        'price': 5.00
    }
    # update will create/update keys in dictionary
    test_recipe.update(params)

    return Recipe.objects.create(user=user, **test_recipe)

class PublicRecipeTests(TestCase):
    """Test publicly avaialble tags API"""
    def setup(self):
        self.client = APIClient()

    def test_auth_required(self):
        res = self.client.get(RECIPE_URL)
        self.assertEqual(res.status_code, status.HTTP_401_UNAUTHORIZED)

class PrivateTagsTests(TestCase):
    """Test the unauthenticated recipe API"""
    def setUp(self):
        self.user = get_user_model().objects.create_user(
            'test12@test.com',
            'testpass',
        )
        self.client = APIClient()
        self.client.force_authenticate(self.user)

    def test_retrieve_recipe(self):
        """Test retrieving a list of recipes"""
        sample_recipe(user=self.user)
        sample_recipe(user=self.user)

        res = self.client.get(RECIPE_URL)

        recipes = Recipe.objects.all().order_by('-id')
        serializer = RecipeSerializer(recipes, many=True)

        self.assertEqual(res.status_code, status.HTTP_200_OK)
        self.assertEqual(res.data, serializer.data)

    def test_recipes_limited(self):
        """Test that recipes are retrieved for user"""

        user2 = get_user_model().objects.create_user(
            'new@test.com',
            'testpass'
        )
        sample_recipe(user=user2)
        sample_recipe(user=self.user)

        res = self.client.get(RECIPE_URL)

        recipes = Recipe.objects.filter(user=self.user)
        serializer = RecipeSerializer(recipes, many=True)

        self.assertEqual(res.status_code, status.HTTP_200_OK)
        self.assertEqual(len(res.data), 1)
        self.assertEqual(res.data, serializer.data)

    def test_view_recipe_detail(self):
        """TEst viewing recipe detail"""
        recipe=sample_recipe(user=self.user)
        recipe.tags.add(sample_tag(user=self.user))
        recipe.ingredients.add(sample_ingredient(user=self.user))

        url = detail_url(recipe.id)
        res = self.client.get(url)

        serializer = RecipeDetailSerializer(recipe)

        self.assertEqual(res.data, serializer.data)

    def test_create_recipe(self):
        """Test creating recipe"""
        new_recipe = {
            'title': 'Cake',
            'time_minutes': 30,
            'price': 15.00
        }

        res = self.client.post(RECIPE_URL, new_recipe)
        self.assertEqual(res.status_code, status.HTTP_201_CREATED)

        recipe = Recipe.objects.get(id=res.data['id'])

        for key in new_recipe.keys():
            self.assertEqual(new_recipe[key], getattr(recipe, key))

    def test_create_recipe_with_tags(self):
        """Creating recipe with tags"""
        tag1 = sample_tag(user=self.user, name='Vegan')
        tag2 = sample_tag(user=self.user, name='Dessert')

        new_recipe = {
            'title': 'CheeseCake',
            'tags': [tag1.id, tag2.id],
            'time_minutes': 30,
            'price': 15.00
        }
        res = self.client.post(RECIPE_URL, new_recipe)
        self.assertEqual(res.status_code, status.HTTP_201_CREATED)

        recipe = Recipe.objects.get(id=res.data['id'])
        tags = recipe.tags.all()

        self.assertEqual(tags.count(), 2)
        self.assertIn(tag1, tags)
        self.assertIn(tag2, tags)

    def test_create_recipe_with_ingredients(self):
        """Creating recipe with ingredients"""
        ingredient1 = sample_ingredient(user=self.user, name='Shrimp')
        ingredient2 = sample_ingredient(user=self.user, name='Ginger')

        new_recipe = {
            'title': 'Prawn curry',
            'ingredients': [ingredient1.id, ingredient2.id],
            'time_minutes': 25,
            'price': 20.00
        }
        res = self.client.post(RECIPE_URL, new_recipe)

        self.assertEqual(res.status_code, status.HTTP_201_CREATED)
        recipe = Recipe.objects.get(id=res.data['id'])
        ingredients = recipe.ingredients.all()

        self.assertEqual(ingredients.count(), 2)
        self.assertIn(ingredient1, ingredients)
        self.assertIn(ingredient2, ingredients)

    def test_partial_update_recipe(self):
        """test updating recipe with patch"""
        recipe = sample_recipe(user=self.user)
        recipe.tags.add(sample_tag(user=self.user))
        new_tag = sample_tag(user=self.user, name='Curry')

        payload = {'title': 'Chicken tikka', 'tags': [new_tag.id]}
        url = detail_url(recipe.id)
        self.client.patch(url,payload)

        recipe.refresh_from_db()
        self.assertEqual(recipe.title, payload['title'])
        tags = recipe.tags.all()
        self.assertEqual(len(tags), 1)
        self.assertIn(new_tag, tags)

    def test_full_update_reciple(self):
        """Test updating a recipe with put"""
        recipe = sample_recipe(user=self.user)
        recipe.tags.add(sample_tag(user=self.user))
        payload = {
        'title': 'Spaghetti',
        'time_minutes': 15,
        'price': 10.00
        }
        url = detail_url(recipe.id)
        self.client.put(url,payload)

        recipe.refresh_from_db()
        self.assertEqual(recipe.title, payload['title'])
        self.assertEqual(recipe.time_minutes, payload['time_minutes'])
        self.assertEqual(recipe.price, payload['price'])

        tags = recipe.tags.all()
        self.assertEqual(len(tags), 0)


import sys
import fileinput
import csv
import operator
import numpy as np
import scipy.spatial.distance as sd
import pickle

#python compare_google.py [google_match-norm] [avg-$count-norm] states.p statescores$count

google = []
matrix = []

infile = sys.argv[2]

google_reader = csv.reader(open(sys.argv[1], 'rb'), delimiter=',')
matrix_reader = csv.reader(open(infile, 'rb'), delimiter=',')
states = pickle.load(open(sys.argv[3], 'rb'))
results = open(sys.argv[4], 'wb')

weeks = 52

def compute_score(m, g):
    sum = 0
    for i in range(m.shape[1]):
        sum += sd.euclidean(m[:,i],g[:,i])

    score = sum/float(m.shape[1])

    return score

#store google data
for row in google_reader:
    row = [(float(x) if x else 0) for x in row]
    google.append(np.array(row))

google = np.array(google)

for row in matrix_reader:
    row = [(float(x) if x else 0) for x in row]
    matrix.append(np.array(row))

matrix = np.array(matrix)

scores = {}
for i in range(google.shape[1]):
    scores[str(i) + ' ' + states[i]] = sd.euclidean(matrix[:,i],google[:,i])

sorted_scores = sorted(scores.iteritems(), key=operator.itemgetter(1))

for name,score in sorted_scores:
    results.write(name+':'+str(score)+'\n')

results.close()


from fastapi import FastAPI
from pydantic import BaseModel
from fastapi.middleware.cors import CORSMiddleware
from reconocer_form import analyze_general_documents
from base_datos import crear_registro
import asyncio
app = FastAPI()

origins = ["*"]

app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials = True,
    allow_methods = ["*"],
    allow_headers = ["*"]
)

class Registro(BaseModel):
    url: str

@app.get("/")
def verRegistros():
    return "registros"

@app.post("/post")
async def crearRegistro(regitsro: Registro):

    repuesta = await analyze_general_documents(regitsro.url)
    registro_guardado = await crear_registro(regitsro.url,repuesta)
    if registro_guardado:
        return repuesta
    return "no se pudo guardar"

@app.put("/post")
def crearRegistro():
    return "hola"

@app.delete("/delete")
def crearRegistro():
    return "hola"


import numpy as np
import torch

with open("pdtSEIKA.csv", "r") as f:
    f_reader = np.loadtxt(f, delimiter=',', dtype=np.float32)
    predict = f_reader
    f.close()
tensor = torch.from_numpy(np.load(r"C:\Users\cchen\PycharmProjects\LearnPyTorch/K05_excluded_xyz.npy"))  # 101778,
# 15,10,10
with open("ai__K05_SEIKA2.csv", "w") as f:
    for item, result in zip(tensor, predict):
        pos = (item[0, 0, 0], item[5, 0, 0], item[10, 0, 0])
        row = "{},{},{},{},{},{},{},{}\n".format(*pos, *result)
        f.write(row)
    f.close()


import os
import copy
from typing import Dict
import numpy as np
import torch
from collections import defaultdict
from hsp.algorithms.population.policy_pool import PolicyPool
from hsp.runner.shared.base_runner import make_trainer_policy_cls
from hsp.utils.shared_buffer import SharedReplayBuffer
from hsp.algorithms.population.utils import _t2n

class TrainerPool:
    """TrainerPool maintains a pool of trainers, each trainer corresponding to one policy, both have the same name.
    For policies that are not trained, use null trainer.
    By specifying mapping from (env_id, agent_id) to trainer_name, TrainerPool creates buffer for each policy.
    """
    def __init__(self, args, policy_pool: PolicyPool, device = torch.device("cpu")):
        self.all_args = args
        self.device = device
        self.policy_pool = policy_pool
        self.trainer_pool = {}
        self.trainer_total_num_steps = defaultdict(int)
        self.use_policy_in_env = dict(args._get_kwargs()).get('use_policy_in_env', False)
        self.__loaded_population = False
        self.__initialized = False

    def policy_config(self, trainer_name):
        return self.policy_pool.policy_config[trainer_name]

    def policy_type(self, trainer_name):
        if trainer_name.startswith("ppo") and trainer_name[-1] in "123":
            return eval(trainer_name[-1])
        elif trainer_name.startswith("policy"):
            # preference policy
            return 4
        else:
            raise RuntimeError(f"Cannot recognize policy type for {trainer_name}.")

    def policy_id(self, trainer_name):
        return int(self.policy_pool.policy_info[trainer_name][1]["id"] * self.policy_pool.num_policies - 1)

    def init_population(self):
        self.on_training = []
        self.best_r = defaultdict(float)
        for policy_name, policy, policy_config, policy_train in self.policy_pool.all_policies():
            # use the same name for trainer and policy
            trainer_name = policy_name
            trainer_cls, _ = make_trainer_policy_cls(policy_config[0].algorithm_name, use_single_network=policy_config[0].use_single_network)
            trainer = trainer_cls(policy_config[0], policy, device = self.device)
            self.trainer_pool[trainer_name] = trainer
            self.best_r[trainer_name] = -1e9
            if policy_train:
                self.on_training.append(trainer_name)

        # trans policies in policy pool to EvalPolicy
        self.policy_pool.trans_to_eval()

        # train info would update when a trainer performs training
        self.train_infos = {}
        self.train_infos.update({f"{trainer_name}-total_num_steps":0 for trainer_name in self.trainer_pool.keys()})

        self.__loaded_population = True

    def reset(self, map_ea2t, n_rollout_threads, num_agents, load_unused_to_cpu=False, **kwargs):
        assert self.__loaded_population
        self.map_ea2t = map_ea2t
        self.n_rollout_threads = n_rollout_threads
        self.num_agents = num_agents

        self.control_agent_count = defaultdict(int)
        self.control_agents = defaultdict(list)
        for (e, a), trainer_name in self.map_ea2t.items():
            self.control_agent_count[trainer_name] += 1
            self.control_agents[trainer_name].append((e, a))

        self.active_trainers = []
        self.buffer_pool: Dict[str, SharedReplayBuffer] = {}
        for trainer_name in self.trainer_pool.keys():
            # set n_rollout_threads as control_agent_count[trainer_name] and num_agents as 1
            if self.control_agent_count[trainer_name] > 0:
                policy_args, obs_space, share_obs_space, act_space = self.policy_config(trainer_name)
                self.buffer_pool[trainer_name] = SharedReplayBuffer(
                    policy_args, 1, obs_space, share_obs_space, act_space,
                    n_rollout_threads=self.control_agent_count[trainer_name])
                self.trainer_pool[trainer_name].to(self.device)
                self.active_trainers.append(trainer_name)
            else:
                if load_unused_to_cpu:
                    self.trainer_pool[trainer_name].to(torch.device("cpu"))
                else:
                    self.trainer_pool[trainer_name].to(self.device)
                self.buffer_pool[trainer_name] = None
        #print("active trainers:", self.active_trainers)
        self.__initialized = True

    def extract_elements(self, trainer_name, x):
        return np.stack([x[e][a] for e, a in self.control_agents[trainer_name]])

    def skip(self, trainer_name):
        # produce actions in parallel envs, skip this trainer
        return (self.use_policy_in_env and trainer_name not in self.on_training) or (trainer_name.startswith("script:"))

    def init_first_step(self, share_obs:np.ndarray, obs:np.ndarray):
        assert self.__initialized
        for trainer_name in self.active_trainers:
            # extract corresponding (e, a) and add num_agent=1 dimension
            obs_lst = np.expand_dims(self.extract_elements(trainer_name, obs), axis=1)
            share_obs_lst = np.expand_dims(self.extract_elements(trainer_name, share_obs), axis=1)
            self.buffer_pool[trainer_name].share_obs[0] = share_obs_lst.copy()
            self.buffer_pool[trainer_name].obs[0] = obs_lst.copy()
        self._step = 0

    def reward_shaping_steps(self):
        """This should differ among algorithms and should be overrided by subclasses.
        """
        reward_shaping_steps = []
        for e in range(self.n_rollout_threads):
            train_tot_num_steps = [self.trainer_total_num_steps[self.map_ea2t[(e, a)]] * int(self.map_ea2t[(e, a)] in self.on_training) for a in range(self.num_agents)]
            reward_shaping_steps.append(max(train_tot_num_steps))
        return reward_shaping_steps

    @torch.no_grad()
    def step(self, step):
        assert self.__initialized
        actions = np.full((self.n_rollout_threads, self.num_agents), fill_value=None).tolist()
        self.step_data = dict()
        for trainer_name in self.active_trainers:
            self.trainer_total_num_steps[trainer_name] += self.control_agent_count[trainer_name]
            self.train_infos[f"{trainer_name}-total_num_steps"] = self.trainer_total_num_steps[trainer_name]

            if self.skip(trainer_name):
                continue

            trainer = self.trainer_pool[trainer_name]
            buffer = self.buffer_pool[trainer_name]

            trainer.prep_rollout()

            value, action, action_log_prob, rnn_states, rnn_states_critic \
                = trainer.policy.get_actions(np.concatenate(buffer.share_obs[step]),
                                np.concatenate(buffer.obs[step]),
                                np.concatenate(buffer.rnn_states[step]),
                                np.concatenate(buffer.rnn_states_critic[step]),
                                np.concatenate(buffer.masks[step]))

            value = np.expand_dims(np.array(_t2n(value)), axis=1)
            action = np.expand_dims(np.array(_t2n(action)), axis=1)
            action_log_prob = np.expand_dims(np.array(_t2n(action_log_prob)), axis=1)
            rnn_states = np.expand_dims(np.array(_t2n(rnn_states)), axis=1)
            rnn_states_critic = np.expand_dims(np.array(_t2n(rnn_states_critic)), axis=1)

            self.step_data[trainer_name] = value, action, action_log_prob, rnn_states, rnn_states_critic

            for i, (e, a) in enumerate(self.control_agents[trainer_name]):
                actions[e][a] = action[i][0]
        return actions

    def insert_data(self, share_obs, obs, rewards, dones, active_masks=None, bad_masks=None, infos=None):
        """
            ndarrays of shape (n_rollout_threads, num_agents, *)
        """
        assert self.__initialized
        self._step += 1
        for trainer_name in self.active_trainers:
            if self.skip(trainer_name):
                continue

            trainer = self.trainer_pool[trainer_name]
            buffer = self.buffer_pool[trainer_name]

            value, action, action_log_prob, rnn_states, rnn_states_critic = self.step_data[trainer_name]

            # (control_agent_count[trainer_name], 1, *)
            obs_lst = np.expand_dims(self.extract_elements(trainer_name, obs), axis=1)
            share_obs_lst = np.expand_dims(self.extract_elements(trainer_name, share_obs), axis=1)
            rewards_lst = np.expand_dims(self.extract_elements(trainer_name, rewards), axis=1)
            dones_lst = np.expand_dims(self.extract_elements(trainer_name, dones), axis=1)

            rnn_states[dones_lst == True] = np.zeros(((dones_lst == True).sum(), buffer.recurrent_N, buffer.hidden_size), dtype=np.float32)
            rnn_states_critic[dones_lst == True] = np.zeros(((dones_lst == True).sum(), *buffer.rnn_states_critic.shape[3:]), dtype=np.float32)
            masks = np.ones((self.control_agent_count[trainer_name], 1, 1), dtype=np.float32)
            masks[dones_lst == True] = np.zeros(((dones_lst == True).sum(), 1), dtype=np.float32)

            bad_masks_lst = active_masks_lst = None
            if bad_masks is not None:
                bad_masks_lst = np.expand_dims(self.extract_elements(trainer_name, bad_masks), axis=1)
            if active_masks is not None:
                active_masks_lst = np.expand_dims(self.extract_elements(trainer_name, active_masks), axis=1)

            if self.all_args.use_task_v_out:
                value = value[:, :, self.policy_id(trainer_name)][:, :, np.newaxis]

            buffer.insert(share_obs_lst, obs_lst, rnn_states, rnn_states_critic, action, action_log_prob, value, rewards_lst, masks, active_masks=active_masks_lst, bad_masks=bad_masks_lst)

            if infos is not None:
                if self.all_args.env_name == "Overcooked" and self.all_args.predict_other_shaped_info:
                    if not hasattr(buffer, "other_shaped_info"):
                        buffer.other_shaped_info = np.zeros((buffer.episode_length + 1, buffer.n_rollout_threads, 1, 12), dtype=np.int32)
                    for i, (e, a) in enumerate(self.control_agents[trainer_name]):
                        buffer.other_shaped_info[self._step, i, 0] = infos[e]["vec_shaped_info_by_agent"][1-a] # insert other agent's shaped info

            # partner policy info
            if self.all_args.env_name == "Overcooked":
                if self.all_args.policy_group_normalization and not hasattr(buffer, "other_policy_type"):
                    buffer.other_policy_type = np.zeros((buffer.episode_length + 1, buffer.n_rollout_threads, 1, 1), dtype=np.int32)
                    for i, (e, a) in enumerate(self.control_agents[trainer_name]):
                        buffer.other_policy_type[:, i, :, :] = self.policy_type(self.map_ea2t[(e, 1-a)])

                if not hasattr(buffer, "other_policy_id"):
                    buffer.other_policy_id = np.zeros((buffer.episode_length + 1, buffer.n_rollout_threads, 1, 1), dtype=np.int32)
                    for i, (e, a) in enumerate(self.control_agents[trainer_name]):
                        buffer.other_policy_id[:, i, :, :] = self.policy_id(self.map_ea2t[(e, 1-a)])

        self.step_data = None

    def compute_advantages(self):
        all_adv = defaultdict(list)
        for trainer_name in self.active_trainers:
            trainer = self.trainer_pool[trainer_name]
            buffer = self.buffer_pool[trainer_name]

            if trainer_name in self.on_training:
                advantages = trainer.compute_advantages(buffer)
                for i, (e, a) in enumerate(self.control_agents[trainer_name]):
                    all_adv[(self.map_ea2t[(e, 0)], self.map_ea2t[(e, 1)], a)].append(advantages[:, i].mean())
        return all_adv

    def train(self, **kwargs):
        assert self.__initialized
        for trainer_name in self.active_trainers:
            trainer = self.trainer_pool[trainer_name]
            buffer = self.buffer_pool[trainer_name]

            if trainer_name in self.on_training:
                trainer.prep_rollout()

                # compute returns
                next_values = trainer.policy.get_values(np.concatenate(buffer.share_obs[-1]),
                                                        np.concatenate(buffer.rnn_states_critic[-1]),
                                                        np.concatenate(buffer.masks[-1]))
                next_values = np.expand_dims(np.array(_t2n(next_values)), axis=1)
                if self.all_args.use_task_v_out:
                    next_values = next_values[:, :, self.policy_id(trainer_name)][:, :, np.newaxis]
                buffer.compute_returns(next_values, trainer.value_normalizer)

                # train
                trainer.prep_training()
                train_info = trainer.train(buffer, turn_on=(self.trainer_total_num_steps[trainer_name] >= self.all_args.critic_warmup_horizon))
                self.train_infos.update({f"{trainer_name}-{k}": v for k, v in train_info.items()})
                self.train_infos.update({f"{trainer_name}-average_episode_rewards": np.mean(buffer.rewards) * buffer.episode_length})

            # place first step observation of next episode
            buffer.after_update()

        return copy.deepcopy(self.train_infos)

    def lr_decay(self, episode, episodes):
        for trainer_name in self.on_training:
            self.trainer_pool[trainer_name].policy.lr_decay(episode, episodes)

    def update_best_r(self, d, save_dir=None):
        for trainer_name, r in d.items():
            trainer = self.trainer_pool[trainer_name]
            if r > self.best_r[trainer_name]:
                self.best_r[trainer_name] = r
                if trainer_name in self.on_training and save_dir is not None:
                    if not os.path.exists(str(save_dir) + "/{}".format(trainer_name)):
                        os.makedirs(str(save_dir) + "/{}".format(trainer_name))
                    #print("save", str(save_dir) + "/{}".format(trainer_name), f"best_r")
                    if self.policy_config(trainer_name)[0].use_single_network:
                        policy_model = trainer.policy.model
                        torch.save(policy_model.state_dict(), str(save_dir) + "/{}/model_best_r.pt".format(trainer_name))
                    else:
                        policy_actor = trainer.policy.actor
                        torch.save(policy_actor.state_dict(), str(save_dir) + "/{}/actor_best_r.pt".format(trainer_name))
                        policy_critic = trainer.policy.critic
                        torch.save(policy_critic.state_dict(), str(save_dir) + "/{}/critic_best_r.pt".format(trainer_name))

    def save(self, step, save_dir):
        for trainer_name in self.on_training:
            trainer = self.trainer_pool[trainer_name]
            if not os.path.exists(str(save_dir) + "/{}".format(trainer_name)):
                os.makedirs(str(save_dir) + "/{}".format(trainer_name))
            trainer_step = self.trainer_total_num_steps[trainer_name]
            #print("save", str(save_dir) + "/{}".format(trainer_name), f"periodic_{trainer_step}")
            if self.policy_config(trainer_name)[0].use_single_network:
                policy_model = trainer.policy.model
                torch.save(policy_model.state_dict(), str(save_dir) + "/{}/model_periodic_{}.pt".format(trainer_name, trainer_step))
            else:
                policy_actor = trainer.policy.actor
                torch.save(policy_actor.state_dict(), str(save_dir) + "/{}/actor_periodic_{}.pt".format(trainer_name, trainer_step))
                policy_critic = trainer.policy.critic
                torch.save(policy_critic.state_dict(), str(save_dir) + "/{}/critic_periodic_{}.pt".format(trainer_name, trainer_step))


import collections
import torchvision.transforms as transforms
import os
import json

try:
    from IPython import embed
except:
    pass

_DATASETS = {}

Dataset = collections.namedtuple(
    'Dataset', ['trainset', 'testset'])

def _add_dataset(dataset_fn):
    _DATASETS[dataset_fn.__name__] = dataset_fn
    return dataset_fn

def _get_transforms(augment=True, normalize=None):
    if normalize is None:
        normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                         std=[0.229, 0.224, 0.225])

    basic_transform = [transforms.ToTensor(), normalize]

    transform_train = []
    if augment:
        transform_train += [
            transforms.RandomResizedCrop(224),
            transforms.RandomHorizontalFlip(),
        ]
    else:
        transform_train += [
            transforms.Resize(256),
            transforms.CenterCrop(224),
        ]
    transform_train += basic_transform
    transform_train = transforms.Compose(transform_train)

    transform_test = [
        transforms.Resize(256),
        transforms.CenterCrop(224),
    ]
    transform_test += basic_transform
    transform_test = transforms.Compose(transform_test)

    return transform_train, transform_test

def _get_mnist_transforms(augment=True, invert=False, transpose=False):
    transform = [
        transforms.ToTensor(),
    ]
    if invert:
        transform += [transforms.Lambda(lambda x: 1. - x)]
    if transpose:
        transform += [transforms.Lambda(lambda x: x.transpose(2, 1))]
    transform += [
        transforms.Normalize((.5,), (.5,)),
        transforms.Lambda(lambda x: x.expand(3, 32, 32))
    ]

    transform_train = []
    transform_train += [transforms.Pad(padding=2)]
    if augment:
        transform_train += [transforms.RandomCrop(32, padding=4)]
    transform_train += transform
    transform_train = transforms.Compose(transform_train)

    transform_test = []
    transform_test += [transforms.Pad(padding=2)]
    transform_test += transform
    transform_test = transforms.Compose(transform_test)

    return transform_train, transform_test

def _get_cifar_transforms(augment=True):
    transform = [
        transforms.ToTensor(),
        transforms.Normalize((0.5071, 0.4867, 0.4408), (0.2675, 0.2565, 0.2761)),
    ]
    transform_train = []
    if augment:
        transform_train += [
            transforms.Pad(padding=4, fill=(125, 123, 113)),
            transforms.RandomCrop(32, padding=0),
            transforms.RandomHorizontalFlip()]
    transform_train += transform
    transform_train = transforms.Compose(transform_train)
    transform_test = []
    transform_test += transform
    transform_test = transforms.Compose(transform_test)
    return transform_train, transform_test

def set_metadata(trainset, testset, config, dataset_name):
    trainset.metadata = {
        'dataset': dataset_name,
        'task_id': config.task_id,
        'task_name': trainset.task_name,
    }
    testset.metadata = {
        'dataset': dataset_name,
        'task_id': config.task_id,
        'task_name': testset.task_name,
    }
    return trainset, testset

@_add_dataset
def inat2018(root, config):
    from dataset.inat import iNat2018Dataset
    transform_train, transform_test = _get_transforms()
    trainset = iNat2018Dataset(root, split='train', transform=transform_train, task_id=config.task_id)
    testset = iNat2018Dataset(root, split='val', transform=transform_test, task_id=config.task_id)
    trainset, testset = set_metadata(trainset, testset, config, 'inat2018')
    return trainset, testset

def load_tasks_map(tasks_map_file):
    assert os.path.exists(tasks_map_file), tasks_map_file
    with open(tasks_map_file, 'r') as f:
        tasks_map = json.load(f)
        tasks_map = {int(k): int(v) for k, v in tasks_map.items()}
    return tasks_map

@_add_dataset
def cub_inat2018(root, config):
    """This meta-task is the concatenation of CUB-200 (first 25 tasks) and iNat (last 207 tasks).

    - The first 10 tasks are classification of the animal species inside one of 10 orders of birds in CUB-200
        (considering all orders except passeriformes).
    - The next 15 tasks are classification of species inside the 15 families of the order of passerifomes
    - The remaining 207 tasks are classification of the species inside each of 207 families in iNat

    As noted above, for CUB-200 10 taks are classification of species inside an order, rather than inside of a family
    as done in the iNat (recall order > family > species). This is done because CUB-200 has very few images
    in each family of bird (expect for the families of passeriformes). Hence, we go up step in the taxonomy and
    consider classification inside a orders and not families.
    """
    NUM_CUB = 25
    NUM_CUB_ORDERS = 10
    NUM_INAT = 207
    assert 0 <= config.task_id < NUM_CUB + NUM_INAT
    transform_train, transform_test = _get_transforms()
    if 0 <= config.task_id < NUM_CUB:
        # CUB
        from dataset.cub import CUBTasks, CUBDataset
        tasks_map_file = os.path.join(root, 'cub/CUB_200_2011', 'final_tasks_map.json')
        tasks_map = load_tasks_map(tasks_map_file)
        task_id = tasks_map[config.task_id]

        if config.task_id < NUM_CUB_ORDERS:
            # CUB orders
            train_tasks = CUBTasks(CUBDataset(root, split='train'))
            trainset = train_tasks.generate(task_id=task_id,
                                            use_species_names=True,
                                            transform=transform_train)
            test_tasks = CUBTasks(CUBDataset(root, split='test'))
            testset = test_tasks.generate(task_id=task_id,
                                          use_species_names=True,
                                          transform=transform_test)
        else:
            # CUB passeriformes families
            train_tasks = CUBTasks(CUBDataset(root, split='train'))
            trainset = train_tasks.generate(task_id=task_id,
                                            task='family',
                                            taxonomy_file='passeriformes.txt',
                                            use_species_names=True,
                                            transform=transform_train)
            test_tasks = CUBTasks(CUBDataset(root, split='test'))
            testset = test_tasks.generate(task_id=task_id,
                                          task='family',
                                          taxonomy_file='passeriformes.txt',
                                          use_species_names=True,
                                          transform=transform_test)
    else:
        # iNat2018
        from dataset.inat import iNat2018Dataset
        tasks_map_file = os.path.join(root, 'inat2018', 'final_tasks_map.json')
        tasks_map = load_tasks_map(tasks_map_file)
        task_id = tasks_map[config.task_id - NUM_CUB]

        trainset = iNat2018Dataset(root, split='train', transform=transform_train, task_id=task_id)
        testset = iNat2018Dataset(root, split='val', transform=transform_test, task_id=task_id)
    trainset, testset = set_metadata(trainset, testset, config, 'cub_inat2018')
    return trainset, testset

@_add_dataset
def imat2018fashion(root, config):
    NUM_IMAT = 228
    assert 0 <= config.task_id < NUM_IMAT
    from dataset.imat import iMat2018FashionDataset, iMat2018FashionTasks
    transform_train, transform_test = _get_transforms()
    train_tasks = iMat2018FashionTasks(iMat2018FashionDataset(root, split='train'))
    trainset = train_tasks.generate(task_id=config.task_id,
                                    transform=transform_train)
    test_tasks = iMat2018FashionTasks(iMat2018FashionDataset(root, split='validation'))
    testset = test_tasks.generate(task_id=config.task_id,
                                  transform=transform_test)
    trainset, testset = set_metadata(trainset, testset, config, 'imat2018fashion')
    return trainset, testset

@_add_dataset
def split_mnist(root, config):
    assert isinstance(config.task_id, tuple)
    from dataset.mnist import MNISTDataset, SplitMNISTTask
    transform_train, transform_test = _get_mnist_transforms()
    train_tasks = SplitMNISTTask(MNISTDataset(root, train=True))
    trainset = train_tasks.generate(classes=config.task_id, transform=transform_train)
    test_tasks = SplitMNISTTask(MNISTDataset(root, train=False))
    testset = test_tasks.generate(classes=config.task_id, transform=transform_test)
    trainset, testset = set_metadata(trainset, testset, config, 'split_mnist')
    return trainset, testset

@_add_dataset
def split_cifar(root, config):
    assert 0 <= config.task_id < 11
    from dataset.cifar import CIFAR10Dataset, CIFAR100Dataset, SplitCIFARTask
    transform_train, transform_test = _get_cifar_transforms()
    train_tasks = SplitCIFARTask(CIFAR10Dataset(root, train=True), CIFAR100Dataset(root, train=True))
    trainset = train_tasks.generate(task_id=config.task_id, transform=transform_train)
    test_tasks = SplitCIFARTask(CIFAR10Dataset(root, train=False), CIFAR100Dataset(root, train=False))
    testset = test_tasks.generate(task_id=config.task_id, transform=transform_test)
    trainset, testset = set_metadata(trainset, testset, config, 'split_cifar')
    return trainset, testset

@_add_dataset
def cifar10_mnist(root, config):
    from dataset.cifar import CIFAR10Dataset
    from dataset.mnist import MNISTDataset
    from dataset.expansion import UnionClassificationTaskExpander
    transform_train, transform_test = _get_cifar_transforms()
    trainset = UnionClassificationTaskExpander(merge_duplicate_images=False)(
        [CIFAR10Dataset(root, train=True), MNISTDataset(root, train=True, expand=True)], transform=transform_train)
    testset = UnionClassificationTaskExpander(merge_duplicate_images=False)(
        [CIFAR10Dataset(root, train=False), MNISTDataset(root, train=False, expand=True)], transform=transform_test)
    return trainset, testset

@_add_dataset
def cifar10(root):
    from torchvision.datasets import CIFAR10
    transform = transforms.Compose([
        transforms.Resize(224),
        transforms.ToTensor(),
        transforms.Normalize((0.5071, 0.4867, 0.4408), (0.2675, 0.2565, 0.2761)),
    ])
    trainset = CIFAR10(root, train=True, transform=transform, download=True)
    testset = CIFAR10(root, train=False, transform=transform)
    return trainset, testset

@_add_dataset
def cifar100(root):
    from torchvision.datasets import CIFAR100
    transform = transforms.Compose([
        transforms.Resize(224),
        transforms.ToTensor(),
        transforms.Normalize((0.5071, 0.4867, 0.4408), (0.2675, 0.2565, 0.2761)),
    ])
    trainset = CIFAR100(root, train=True, transform=transform, download=True)
    testset = CIFAR100(root, train=False, transform=transform)
    return trainset, testset

@_add_dataset
def mnist(root):
    from torchvision.datasets import MNIST
    transform = transforms.Compose([
        lambda x: x.convert("RGB"),
        transforms.Resize(224),
        transforms.ToTensor(),
        # transforms.Normalize((0.5, 0.5, 0.5), (1., 1., 1.)),
    ])
    trainset = MNIST(root, train=True, transform=transform, download=True)
    testset = MNIST(root, train=False, transform=transform)
    return trainset, testset

@_add_dataset
def letters(root):
    from torchvision.datasets import EMNIST
    transform = transforms.Compose([
        lambda x: x.convert("RGB"),
        transforms.Resize(224),
        transforms.ToTensor(),
        # transforms.Normalize((0.5, 0.5, 0.5), (1., 1., 1.)),
    ])
    trainset = EMNIST(root, train=True, split='letters', transform=transform, download=True)
    testset = EMNIST(root, train=False, split='letters', transform=transform)
    return trainset, testset

@_add_dataset
def kmnist(root):
    from torchvision.datasets import KMNIST
    transform = transforms.Compose([
        lambda x: x.convert("RGB"),
        transforms.Resize(224),
        transforms.ToTensor(),
    ])
    trainset = KMNIST(root, train=True, transform=transform, download=True)
    testset = KMNIST(root, train=False, transform=transform)
    return trainset, testset

@_add_dataset
def stl10(root):
    from torchvision.datasets import STL10
    transform = transforms.Compose([
        transforms.Resize(224),
        transforms.ToTensor(),
        transforms.Normalize((0.5071, 0.4867, 0.4408), (0.2675, 0.2565, 0.2761)),
    ])
    trainset = STL10(root, split='train', transform=transform, download=True)
    testset = STL10(root, split='test', transform=transform)
    trainset.targets = trainset.labels
    testset.targets = testset.labels
    return trainset, testset

def get_dataset(root, config=None):
    return _DATASETS[config.name](os.path.expanduser(root), config)


# vim set fileencoding=utf-8
from setuptools import setup

with open('README.rst') as f:
    long_description = f.read()
setup(
    name = 'AnthraxEplasty',
    version = '0.0.3',
    author = 'Szymon Pyżalski',
    author_email = 'zefciu <szymon@pythonista.net>',
    description = 'Anthrax - generating forms from Elephantoplasty objects',
    url = 'http://github.com/zefciu/Anthrax',
    keywords = 'form web orm database',
    long_description = long_description,

    install_requires = ['anthrax', 'Elephantoplasty'],
    tests_require = ['nose>=1.0', 'nose-cov>=1.0'],
    test_suite = 'nose.collector',
    package_dir = {'': 'src'},
    namespace_packages = ['anthrax'],
    packages = [
        'anthrax', 'anthrax.eplasty'
    ],
    classifiers = [
        'Development Status :: 1 - Planning',
        'Intended Audience :: Developers',
        'License :: OSI Approved :: BSD License',
        'Programming Language :: Python :: 3',
        'Programming Language :: Python :: 3.2',
        'Topic :: Internet :: WWW/HTTP :: Dynamic Content',
    ],
    entry_points = """[anthrax.reflector]
eplasty = anthrax.eplasty.reflector:EplastyReflector
[anthrax.field_mixins]
eplasty_unique = anthrax.eplasty.field:UniqueMixin
""",
)


from typing import List

from project.appliances.appliance import Appliance
from project.people.child import Child

class Room:

    def __init__(self, name: str, budget: float, members_count: int):
        self.family_name = name
        self.budget = budget
        self.members_count = members_count
        self.children: List[Child] = []
        self.expenses = 0

    @property
    def expenses(self):
        return self.__expenses

    @expenses.setter
    def expenses(self, value):
        if value < 0:
            raise ValueError("Expenses cannot be negative")

        self.__expenses = value

    def calculate_expenses(self, *args):
        total_expenses = 0

        for list_obj in args:
            for obj in list_obj:
                if isinstance(obj, Appliance):
                    total_expenses += obj.get_monthly_expense()
                else:
                    total_expenses += obj.cost * 30

        self.expenses = total_expenses

    def room_info(self):
        result_str = [
            f"{self.family_name} with {self.members_count} members."
            f" Budget: {self.budget:.2f}$, Expenses: {self.expenses:.2f}$"]

        for idx, child in enumerate(self.children):
            result_str.append(f"--- Child {idx + 1} monthly cost: {(child.cost * 30):.2f}$")

        if hasattr(self, 'appliances'):
            appliances_monthly_cost = sum([a.get_monthly_expense() for a in self.appliances])
            result_str.append(f"--- Appliances monthly cost: {appliances_monthly_cost:.2f}$")

        return '\n'.join(result_str)


"""Create category

Revision ID: bc8fb2b5aaaa
Revises: cf3388347129
Create Date: 2023-05-06 09:44:36.431462

"""
from alembic import op
import sqlalchemy as sa

# revision identifiers, used by Alembic.
revision = 'bc8fb2b5aaaa'
down_revision = 'cf3388347129'
branch_labels = None
depends_on = None

def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table('category',
    sa.Column('id', sa.Integer(), nullable=False),
    sa.Column('name', sa.String(length=50), nullable=True),
    sa.Column('icon', sa.String(length=50), nullable=True),
    sa.Column('description', sa.Text(), nullable=True),
    sa.Column('create_at', sa.DateTime(timezone=True), server_default=sa.text('now()'), nullable=True),
    sa.Column('update_at', sa.DateTime(timezone=True), server_default=sa.text('now()'), nullable=True),
    sa.PrimaryKeyConstraint('id'),
    sa.UniqueConstraint('name')
    )
    op.create_index(op.f('ix_category_id'), 'category', ['id'], unique=False)
    # ### end Alembic commands ###

def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_index(op.f('ix_category_id'), table_name='category')
    op.drop_table('category')
    # ### end Alembic commands ###


from flask_wtf import FlaskForm
from wtforms import StringField, IntegerField, SelectField
from wtforms.validators import InputRequired, Optional

sizes = ['XS','S', 'M', 'L', 'XL']
ratings = [1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0, 4.5, 5.0, 5.5,
           6.0, 6.5, 7.0, 7.5, 8.0, 8.5, 9.0, 9.5, 10.0]

class AddCupcakeForm(FlaskForm):
    """Form for adding cupcakes"""

    flavor = StringField("Flavor", validators=[InputRequired(message="Flavor cannot be blank")])
    size = SelectField("Size", choices=[(size, size) for size in sizes], validators=[InputRequired(message="Size cannot be blank")])
    rating = SelectField("Rating", coerce=float, choices=[(rating, rating) for rating in ratings], validators=[InputRequired(message="Rating cannot be blank")])
    image = StringField("Image Link", validators=[Optional()])


# Import the utils
import math
import sys
import os
sys.path.append(os.path.abspath('../QUtils'))
from qutils import pprint, graph

# Import the QISKit SDK
from qiskit import QuantumCircuit, ClassicalRegister, QuantumRegister
from qiskit import execute

# Create a Quantum Register with 2 qubits.
q = QuantumRegister(3)
# Create a Classical Register with 2 bits.
c = ClassicalRegister(3)
# Create a Quantum Circuit
qc = QuantumCircuit(q, c)

# move the qubits into a superpostion such that when they have an H gate and a Measure
# applied they are equally likely to collapse to 0 or 1
qc.h(q)
qc.s(q)

pi = math.pi
# add the gates to get the normal distribution
"""
# makes 111 less likely
qc.crz(-0.3*pi, q[2], q[1])
qc.crz(-0.3*pi, q[1], q[0])

# some how encourages 100
qc.x(q[2])
qc.crz(-0.4*pi, q[0], q[1])
qc.crz(-0.4*pi, q[1], q[2])
qc.crz(-0.4*pi, q[0], q[2])
qc.x(q[2])

"""

# THE 2 QUBIT CIRCUIT
# I don't get why but this bumps up the ends and drops the middle
# aka the exact opposite of what we want
#qc.crz(-0.5*pi, q[1], q[0])
#qc.crz(-0.5*pi, q[0], q[1])

# this then flips it the right way up
#qc.rz(pi, q[0])

# ATTEMPT AT 3 QUBIT CIRCUIT - kinda works?
# ups 011 and 100, but also 000 and 111
qc.crz(-0.5*pi, q[1], q[0])
qc.crz(-0.5*pi, q[0], q[1])
qc.crz(-0.5*pi, q[2], q[1])
qc.rz(pi, q[2])

qc.h(q)
# Add a Measure gate to see the state.
qc.measure(q, c)
# Compile and run the Quantum circuit on a simulator backend
job_sim = execute(qc, "local_qasm_simulator", shots=7000)
sim_result = job_sim.result()

# Show the results
print("simulation: ", sim_result)
# Returns a dict
pprint(sim_result.get_counts(qc))
graph(sim_result.get_counts(qc))