"""
Assignment 5: K-Means. See the instructions to complete the methods below.
"""

from collections import Counter
from collections import defaultdict
import gzip
import math
import numpy as np

class KMeans(object):

    def __init__(self, k=2):
        """ Initialize a k-means clusterer. Should not have to change this."""
        self.k = k

    def cluster(self, documents, iters=10):
        """
        Cluster a list of unlabeled documents, using iters iterations of k-means.
        Initialize the k mean vectors to be the first k documents provided.
        After each iteration, print:
        - the number of documents in each cluster
        - the error rate (the total Euclidean distance between each document and its assigned mean vector), rounded to 2 decimal places.
        See Log.txt for expected output.
        The order of operations is:
        1) initialize means
        2) Loop
          2a) compute_clusters
          2b) compute_means
          2c) print sizes and error
        """
        self.doc_norm = defaultdict(lambda: 0.0)
        self.documents = documents

        self.mean_norms=[]
        self.mean_vectors = []
        for doc_id in range(len(documents)):
            self.doc_norm[doc_id] = self.sqnorm(documents[doc_id])
            if doc_id < self.k:
                self.mean_vectors.append(documents[doc_id])
                self.mean_norms.append(self.doc_norm[doc_id])

        for j in range(iters):

            self.compute_clusters(documents)
            self.compute_means()

            num_of_docs = []
            for i in self.k_cluster:
                num_of_docs.append(len(self.k_cluster[i]))

            print(num_of_docs)
            print(self.error(documents))

    def compute_means(self):
        """ Compute the mean vectors for each cluster (results stored in an
        instance variable of your choosing)."""
        self.mean_vectors = []
        for i in range(self.k):
            term_freq = Counter()
            for doc_id in self.k_cluster[i]:
                term_freq.update(self.documents[doc_id])
            if(len(self.k_cluster[i]) > 0):
                for term in term_freq:
                    term_freq[term] = 1.0 * term_freq[term] / len(self.k_cluster[i])
                self.mean_vectors.append(term_freq)
            self.mean_norms = []
            for t in self.mean_vectors:
                self.mean_norms.append(self.sqnorm(t))

    def compute_clusters(self, documents):
        """ Assign each document to a cluster. (Results stored in an instance
        variable of your choosing). """
        self.k_cluster = defaultdict(lambda: [])
        for doc_id in range(len(documents)):
            assign_cluster = -1
            min_distance = -1
            for cluster in range(self.k):
                distance = self.distance(documents[doc_id],self.mean_vectors[cluster],self.mean_norms[cluster]+self.doc_norm[doc_id])
                if ( distance < min_distance or assign_cluster == -1):
                    assign_cluster = cluster
                    min_distance = distance
            self.k_cluster[assign_cluster].append(doc_id)

    def sqnorm(self, d):
        """ Return the vector length of a dictionary d, defined as the sum of
        the squared values in this dict. """
        sqsum = 0.0
        for key in d.keys():
            sqsum += (d[key]**2)

        return sqsum

    def distance(self, doc, mean, mean_norm):
        """ Return the Euclidean distance between a document and a mean vector.
        See here for a more efficient way to compute:
        http://en.wikipedia.org/wiki/Cosine_similarity#Properties"""
        distance = mean_norm
        for term in doc:
            distance += - ( 2.0 * doc[term] * mean[term] )

        return float(math.sqrt(distance))

    def error(self, documents):
        """ Return the error of the current clustering, defined as the total
        Euclidean distance between each document and its assigned mean vector."""

        error = 0.0
        self.k_cluster_dist = defaultdict(lambda: [])
        for cluster in self.k_cluster.keys():
            for doc_id in self.k_cluster[cluster]:
                distance = self.distance(documents[doc_id],self.mean_vectors[cluster],self.mean_norms[cluster]+self.doc_norm[doc_id])
                error += distance
                self.k_cluster_dist[cluster].append((documents[doc_id],distance))

        return error

    def print_top_docs(self, n=10):
        """ Print the top n documents from each cluster. These are the
        documents that are the closest to the mean vector of each cluster.
        Since we store each document as a Counter object, just print the keys
        for each Counter (sorted alphabetically).
        Note: To make the output more interesting, only print documents with more than 3 distinct terms.
        See Log.txt for an example."""
        for cluster in self.k_cluster.keys():
            print("CLUSTER ",cluster)
            topdocs = sorted(self.k_cluster_dist[cluster],key=lambda x:x[1])
            count = 0
            for doc_id in range(len(topdocs)):
                if(len(topdocs[doc_id][0]) > 3):
                    str = ' '.join(sorted(topdocs[doc_id][0].keys())).encode('utf-8')
                    print(str.decode('utf-8'))
                    count += 1
                if count == n:
                    break

def prune_terms(docs, min_df=3):
    """ Remove terms that don't occur in at least min_df different
    documents. Return a list of Counters. Omit documents that are empty after
    pruning words.
    >>> prune_terms([{'a': 1, 'b': 10}, {'a': 1}, {'c': 1}], min_df=2)
    [Counter({'a': 1}), Counter({'a': 1})]
    """
    term_doc_freq = defaultdict(lambda: 0)

    for doc in docs:
        for term in doc.keys():
            term_doc_freq[term] += 1

    result = []
    for doc in docs:
        freq = Counter()
        for term in doc.keys():
            if term_doc_freq[term] >= min_df:
                freq[term] += doc[term]
        if freq:
            result.append(freq)

    return result

def read_profiles(filename):
    """ Read profiles into a list of Counter objects.
    DO NOT MODIFY"""
    profiles = []
    with gzip.open(filename, mode='rt', encoding='utf8') as infile:
        for line in infile:
            profiles.append(Counter(line.split()))
    return profiles

def main():
    profiles = read_profiles('profiles.txt.gz')
    print('read', len(profiles), 'profiles.')
    profiles = prune_terms(profiles, min_df=2)
    km = KMeans(k=10)
    km.cluster(profiles, iters=20)
    km.print_top_docs()

if __name__ == '__main__':
    main()


from fastapi import FastAPI
import uvicorn
import json
import pandas as pd

app =FastAPI()

# load the files optimized for the functions requests
function1 = pd.read_parquet('function1.parquet')
function2 = pd.read_parquet('function2.parquet')
function3 = pd.read_parquet('function3.parquet')
function4 = pd.read_parquet('function4.parquet')
function5 = pd.read_parquet('function5.parquet')

# healthcheck
@app.get("/")
def read_root():
    """
    Healthcheck endpoint that returns a JSON response with the status "ok".
    """
    return {"status": "ok"}

# Function 1
@app.get("/PlaytimeByGenre")
def function_1(genre):
    """
    Endpoint to calculate the year with the most playtime for a given genre.

    Args:
        genre (str): The genre for which to find the year with the most playtime.

    Returns:
        dict: A dictionary containing the result with the year of most playtime for the genre.
    """

    if genre in function1['genres'].values:
        max_playtime = function1.loc[function1['genres'] == genre, 'playtime_forever'].max()
        year_with_max_playtime = function1.loc[function1['playtime_forever'] == max_playtime, 'year'].values[0]
        response = {f"Year with the most playtime for Genre {genre}": str(year_with_max_playtime)}
    else:
        response = "Enter a valid genre"
    return response

# Function 2
@app.get("/TotalPlaytimeByGenre")
def function_2(genre):
    """
    Endpoint to calculate the user with the most playtime and their playtime breakdown by year for a given genre.

    Args:
        genre (str): The genre for which to find the user with the most playtime.

    Returns:
        dict: A dictionary containing the user with the most playtime and their playtime breakdown by year.
    """
    if genre in function2['genres'].values:
        genre_data = function2[function2['genres'] == genre]
        total_playtime_by_user = genre_data.groupby('user_id')['playtime_forever'].sum().reset_index()
        user_with_max_playtime = total_playtime_by_user.loc[total_playtime_by_user['playtime_forever'].idxmax(), 'user_id']
        max_playtime = total_playtime_by_user['playtime_forever'].max()
        years_user_with_max_playtime = genre_data[genre_data['user_id'] == user_with_max_playtime]
        playtime_by_year = years_user_with_max_playtime.groupby('year')['playtime_forever'].sum().reset_index()
        playtime_by_year = playtime_by_year.rename(columns={'year': 'Year', 'playtime_forever': 'Playtime (hours)'})
        playtime_data = playtime_by_year.to_dict(orient='records')
        response = {f"User with the most playtime for Genre {genre}": str(user_with_max_playtime),
                    "Max Playtime": max_playtime,
                    "Playtime by Year": playtime_data}
    else:
        response = "Enter a valid genre"
    return response

# Function 3
@app.get("/Top3ItemsBySentiment")
def function_3(year):
    """
    Endpoint to find the top 3 items with the highest sentiment scores for a given year.

    Args:
        year (int): The year for which to find the top items.

    Returns:
        list: A list of dictionaries containing the top 3 items for the given year.
    """
    year = int(year)
    if year in function3['year'].values:
        year_data = function3[function3['year'] == year]
        top3_items = year_data.sort_values(by='sentiment', ascending=False).head(3)
        top3_items = top3_items.reset_index(drop=True)
        top3_items_data = top3_items['item_name'].to_dict()
        top3_items_list = [{"Rank " + str(index + 1): value} for index, value in enumerate(top3_items_data.values())]
        response = top3_items_list
    else:
        response = "Enter a valid year"
    return response

# Function 4
@app.get("/Bottom3ItemsBySentiment")
def function_4(year):
    """
    Endpoint to find the bottom 3 items with the lowest sentiment scores for a given year.

    Args:
        year (int): The year for which to find the bottom items.

    Returns:
        list: A list of dictionaries containing the bottom 3 items for the given year.
    """
    year = int(year)
    if year in function4['year'].values:
        year_data = function4[function4['year'] == year]
        bottom3_items = year_data.sort_values(by='sentiment', ascending=True).head(3)
        bottom3_items = bottom3_items.reset_index(drop=True)
        bottom3_items_data = bottom3_items['item_name'].to_dict()
        bottom3_items_list = [{"Rank " + str(index + 1): value} for index, value in enumerate(bottom3_items_data.values())]
        response = bottom3_items_list
    else:
        response = "Enter a valid year"
    return response

# function 5
@app.get("/SentimentAnalysisByYear")
def function_5(year):
    """
    Endpoint to retrieve sentiment analysis data (negative, neutral, positive) for a given year.

    Args:
        year (int): The year for which to retrieve sentiment analysis data.

    Returns:
        dict: A dictionary containing sentiment analysis data for the given year.
    """
    year = int(year)
    if year in function5['year'].values:
        year_data = function5[function5['year'] == year]
        sentiment_data = {
            "Negative": year_data['Negative'].values[0],
            "Neutral": year_data['Neutral'].values[0],
            "Positive": year_data['Positive'].values[0]
        }
        response = sentiment_data
    else:
        response = "Enter a valid year"
    return response


#   File: RandomWalker.py
#   Author: Nawaf Abdullah
#   Creation Date: 19/Mar/2018
#   Description: Modeling of the random walker problem in 1D, 2D and 3D
#   Reference: Computational Physics by Nichholas J. Giordano & Hisao Nakanishi (2nd Ed)
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import random

class Walker:
    n_points = 0
    x = [0]
    y = [0]
    z = [0]
    z_switch = False

    def __init__(self, N):
        self.n_points = N
        for i in range(N-1):
            self.x.append(0)
            self.y.append(0)
            self.z.append(0)

    def init_npoints(self, N):
        del self.x[:]
        del self.y[:]
        del self.z[:]
        self.__init__(N)

    def one_dimension(self):
        for i in range(self.n_points-1):
            r = random.randint(0, 100)
            if r > 50:
                self.x[i] = self.x[i] + 1
            else:
                self.x[i] = self.x[i] - 1
            self.y[i+1] = self.y[i] + 1
        self.z_switch = False
        return self.x, self.y

    def two_dimension(self):
        for i in range(self.n_points-1):
            if i == 0:
                continue
            r = random.randint(0, 100)
            if r <= 25:
                self.x[i+1] = self.x[i] + 1
                self.y[i+1] = self.y[i]
            elif 25 < r <= 50:
                self.x[i+1] = self.x[i] - 1
                self.y[i+1] = self.y[i]
            elif 50 < r <= 75:
                self.y[i+1] = self.y[i] + 1
                self.x[i+1] = self.x[i]
            else:
                self.y[i+1] = self.y[i] - 1
                self.x[i+1] = self.x[i]
        self.z_switch = False
        return self.x, self.y

    def three_dimension(self):
        for i in range(self.n_points-1):
            if i == 0:
                continue
            r = random.randint(0, 100)
            if r <= 16.66:
                self.x[i+1] = self.x[i] + 1
                self.y[i+1] = self.y[i]
                self.z[i+1] = self.z[i]
            elif 16.66 < r <= 33.32:
                self.x[i+1] = self.x[i] - 1
                self.y[i+1] = self.y[i]
                self.z[i+1] = self.z[i]
            elif 33.32 < r <= 49.98:
                self.y[i+1] = self.y[i] + 1
                self.x[i+1] = self.x[i]
                self.z[i+1] = self.z[i]
            elif 49.98 < r <= 66.64:
                self.y[i+1] = self.y[i] - 1
                self.x[i+1] = self.x[i]
                self.z[i+1] = self.z[i]
            elif 66.64 < r < 83.33:
                self.z[i+1] = self.z[i] + 1
                self.x[i+1] = self.x[i]
                self.y[i+1] = self.y[i]
            else:
                self.z[i+1] = self.z[i] - 1
                self.x[i+1] = self.x[i]
                self.y[i+1] = self.y[i]
        self.z_switch = True
        return self.x, self.y, self.z

    def plot_walker(self):
        if self.z_switch is True:
            ax = plt.axes(projection='3d')
            ax.plot3D(self.x, self.y, self.z, 'gray')
            plt.show()
        else:
            fig1 = plt.figure(1)
            fig1.suptitle("Walker Path")
            plt.xlabel("x-position")
            plt.ylabel("y-position")
            plt.plot(self.x, self.y)
            plt.show()


"""
Setup.py for eScrapper, modules to do scrapping on version control portals,
providing some kind of "API"
"""
#!/usr/bin/env python
# -*- coding: utf-8 -*-
from __future__ import print_function
from setuptools import setup
from setuptools.command.test import test as TestCommand
import io
import os
import sys
import re

HERE = os.path.abspath(os.path.dirname(__file__))

def read(*filenames, **kwargs):
    """ Read a python's file """
    encoding = kwargs.get('encoding', 'utf-8')
    sep = kwargs.get('sep', '\n')
    buf = []
    for filename in filenames:
        with io.open(filename, encoding=encoding) as current_file:
            buf.append(current_file.read())
    return sep.join(buf)

LONG_DESCRIPTION = read('README.rst')

SOURCE = read('escrapper/__init__.py')
PATTERN = re.compile(r'''__version__ = ['"](?P<version>[\d.]+)['"]''')
VERSION = PATTERN.search(SOURCE).group('version')

class PyTest(TestCommand):
    """ Mixin to integrate pytest with setup.py """
    def finalize_options(self):
        TestCommand.finalize_options(self)
        self.test_args = []
        self.test_suite = True

    def run_tests(self):
        import pytest
        errcode = pytest.main(self.test_args)
        sys.exit(errcode)

setup(
    name='escrapper',
    version=VERSION,
    url='http://github.com/esparta/escrapper',
    license='Apache Software License',
    author="Espartaco Palma",
    tests_require=['tox'],
    install_requires=['beautifulsoup4>=4.3.1',
                      'requests>=1.2.3'
                      ],
    cmdclass={'test': PyTest},
    author_email='esparta@gmail.com',
    description='Scrapping tool, can process WebSVN portal ',
    long_description=LONG_DESCRIPTION,
    packages=['escrapper'],
    include_package_data=True,
    platforms='any',
    test_suite='escrapper.test_app',
    classifiers=[
        'Programming Language :: Python',
        'Development Status :: 1 - Beta',
        'Natural Language :: English',
        'Environment :: Web Environment',
        'Intended Audience :: Developers',
        'License :: OSI Approved :: Apache Software License',
        'Operating System :: OS Independent',
        'Topic :: Software Development :: Libraries :: Python Modules',
        'Topic :: Internet :: WWW/HTTP :: Dynamic Content',
        ],
    extras_require={'testing': ['pytest'], }
)


#!/usr/bin/env python3
import os
import sys
import logging
import time
import argparse
import configparser
from itertools import chain
from distutils.util import strtobool
import numpy as np
import scipy.interpolate
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import pandas as pd
import ftplib
import io
import datetime
from collections import OrderedDict

plt.style.use('seaborn-paper')

DEFAULT_LOGGING_LEVEL = logging.INFO
MAX_LOGGING_LEVEL = logging.CRITICAL

ISD_LITE_FTP_BASE_PATH = "/pub/data/noaa/isd-lite"
ISD_LITE_SPEC = [
    [( 0, 4), 'year', int],
    [( 5, 7), 'month', int],
    [( 8,10), 'day', int],
    [(11,13), 'hour', int],
    [(13,19), 'air temp', lambda x: np.nan if x.strip()=='-9999' else int(x)/10],
    [(19,25), 'dewpoint', lambda x: np.nan if x.strip()=='-9999' else int(x)/10],
    [(25,31), 'sea level pressure', lambda x: np.nan if x=='-9999' else int(x)/10],
    [(31,37), 'wind direction', lambda x: np.nan if x=='-9999' else int(x)],
    [(37,43), 'wind speed', lambda x: np.nan if x=='-9999' else int(x)/10],
    [(43,49), 'sky condition total coverage code', lambda x: np.nan if x=='-9999' else int(x)],
    [(49,55), 'liquid percip 1hr', lambda x: np.nan if x=='-9999' else int(x)/10],
    [(55,61), 'liquid percip 6hr', lambda x: np.nan if x=='-9999' else int(x)/10],
    ]

def Main(argv):
    tic_total = time.time()

    # parse cfg_file argument
    conf_parser = argparse.ArgumentParser(description=__doc__,
                                          add_help=False)  # turn off help so later parse (with all opts) handles it
    conf_parser.add_argument('-c', '--cfg-file', type=argparse.FileType('r'), default='main.cfg',
                             help="Config file specifiying options/parameters.\nAny long option can be set by removing the leading '--' and replacing '-' with '_'")
    args, remaining_argv = conf_parser.parse_known_args(argv)
    # build the config (read config files)
    cfg_filename = None
    if args.cfg_file:
        cfg_filename = args.cfg_file.name
        cfg = configparser.ConfigParser(inline_comment_prefixes=('#',';'))
        cfg.optionxform = str # make configparser case-sensitive
        cfg.read_file(chain(("[DEFAULTS]",), args.cfg_file))
        defaults = dict(cfg.items("DEFAULTS"))
        # special handling of paratmeters that need it like lists
        # defaults['make_temperature_plots'] = strtobool(defaults['make_temperature_plots'])
        #        if( 'bam_files' in defaults ): # bam_files needs to be a list
        #            defaults['bam_files'] = [ x for x in defaults['bam_files'].split('\n') if x and x.strip() and not x.strip()[0] in ['#',';'] ]
    else:
        defaults = {}

    # Parse rest of arguments with a new ArgumentParser
    aparser = argparse.ArgumentParser(description=__doc__, parents=[conf_parser])
    # parse arguments
    aparser.add_argument('-f', '--force', action='store_true', default=False,
                         help="Download/fetch temperature data instead of just updating")
    aparser.add_argument('-P', '--make-temperature-plots', type=strtobool, default=False,
                         help="Generate plots when fetching and cleaning temperatures")
    aparser.add_argument('-s', "--station-callsign", help="Callsign of the weather station to use")
    #aparser.add_argument("--start-year", type=int, help="First year to use historic temperature data for")
    aparser.add_argument("--potentially-problematic-gap-size", type=pd.Timedelta,
                         help="Time gaps in temperature data longer than this are noted and filled via day-over-day interpolation")
    aparser.add_argument("--outlier_rolling_sigma_window", type=int, help="0 to just use median instead of median/sigma")
    aparser.add_argument("--outlier_rolling_median_window", type=int, help="Size of window for diff from median (or median/sigma) outlier detection")
    aparser.add_argument("--outlier_thresh", type=float, help="Threshold for outlier detection")
    aparser.add_argument("--outlier_multipass", type=strtobool, help="Do multiple outlier removal passes; use True")

    aparser.add_argument('-v', '--verbose', action='count', default=0,
                        help="increase logging verbosity")
    aparser.add_argument('-q', '--quiet', action='count', default=0,
                        help="decrease logging verbosity")
    aparser.set_defaults(**defaults) # add the defaults read from the config file
    args = aparser.parse_args(remaining_argv)

    # check required arguments @TCC TODO more to fill in here
    if not 'basedir' in args or not args.basedir:
        args.basedir = os.getcwd()
    if not 'station_callsign' in args or not args.station_callsign:
        logging.error("station_callsign parameter is required: 4 letter callsign for the weather station to use")
        sys.exit(2)
    #if not 'start_year' in args or not args.start_year:
    #    args.start_year = None

    # setup logger
    setup_logger(verbose_level=args.verbose-args.quiet)
    logging.info("Using config file: `{}`".format(cfg_filename))
    logging.info('Using passed arguments: '+str(argv))
    logging.info('args: '+str(args))

    # # @TCC TMP -- print out all the options/arguments
    # for k,v in vars(args).items():
    #     print(k,":",v, file=sys.stderr)

    datadir = os.path.join(args.basedir, args.temperature_data_dir)

    ####
    tempdf = fetch_and_process_temperatures(
                                    station_callsign=args.station_callsign,
                                    #start_year=args.start_year,
                                    basedir=args.basedir,
                                    datadir=os.path.join(args.basedir, args.temperature_data_dir),
                                    make_temperature_plots=args.make_temperature_plots,
                                    NOAA_ISD_ftp_host=args.NOAA_ISD_ftp_host,
                                    potentially_problematic_gap_size=args.potentially_problematic_gap_size,
                                    force_download_temps=args.force,
                                    outlier_rolling_sigma_window=args.outlier_rolling_sigma_window,
                                    outlier_rolling_median_window=args.outlier_rolling_median_window,
                                    outlier_thresh=args.outlier_thresh,
                                    outlier_multipass=args.outlier_multipass,
                                )

    ## end
    logging.info("Done: {:.2f} sec elapsed".format(time.time()-tic_total))
    return 0

def fetch_and_process_temperatures(station_callsign,
                                   #start_year,
                                   basedir,
                                   datadir,
                                   make_temperature_plots,
                                   NOAA_ISD_ftp_host,
                                   potentially_problematic_gap_size,
                                   force_download_temps,
                                   outlier_rolling_sigma_window=24*5,  # None or 0 to just use median instead of median/sigma
                                   outlier_rolling_median_window=5,
                                   outlier_thresh=1.5,  # deviation from media/sigma to trigger removal
                                   outlier_multipass=True,  # cycle until no points removed, or False for not
                                   ):
    os.makedirs(datadir, exist_ok=True)
    now_year = datetime.datetime.now().year

    # get the isd_history (aka station_info)
    station_info = get_ISD_history_file(datadir, NOAA_ISD_ftp_host, save_file=True,
                                     force_download=force_download_temps)
    # pick just the info associated with the station we want
    station_info = station_info[station_info['CALL'] == station_callsign]
    logging.info("Station info for "+station_callsign+"\n"+str(station_info))

    # load the last fetch timestamp
    last_fetch_time = None
    last_fetch_fn = os.path.join(datadir, "last_fetch_time_{}.txt".format(station_callsign))
    if not force_download_temps:
        try:
            with open(last_fetch_fn, 'r') as fh:
                last_fetch_time = fh.readline().strip()
                last_fetch_time = time.strptime(last_fetch_time, "%Y-%m-%dT%H:%M:%S%z")
        except FileNotFoundError:
            logging.info("No '{}' file found".format(last_fetch_fn))
            last_fetch_time = None
            pass
        except ValueError as err:
            logging.error("Did not understand '{}': {}".format(last_fetch_fn, str(err)))
            last_fetch_time = None

    isd_files = []
    to_get = []
    for k, row in station_info.iterrows():
        for year in range(int(row['BEGIN'][0:4]), int(row['END'][0:4])+2):  # +2 because isd-history may be out of date
            if year <= now_year:
                isd_files.append([year, "{}-{}-{:04d}.gz".format(row['USAF'], row['WBAN'], year)])
                if last_fetch_time is None or year >= last_fetch_time.tm_year:
                    to_get.append(isd_files[-1])

    logging.info("Number of possible ISD files = {:d}".format(len(isd_files)))
    logging.info("Number of files try and get/update = {:d}".format(len(to_get)))

    ## Download isd-lite files
    # print(*to_get, sep="\n")
    if len(to_get) > 0:
        with ftplib.FTP(host=NOAA_ISD_ftp_host) as ftpconn:
            ftpconn.login()
            for year, fn in to_get:
                url = "{}/{:04d}/{}".format(ISD_LITE_FTP_BASE_PATH,year,fn)
                logging.info("Fetching: '{}'".format(url))
                has_err = False
                with open(os.path.join(datadir, fn), 'wb') as fh:
                    try:
                        ftpconn.retrbinary('RETR ' + url, fh.write)
                    except ftplib.error_perm as err:
                        has_err = True
                        if str(err).startswith('550 '):
                            logging.warning(err)
                        else:
                            raise
                if has_err:
                    os.remove(os.path.join(datadir, fn))

    # save the timestamp to a file
    with open(last_fetch_fn, 'w') as fh:
        print(time.strftime("%Y-%m-%dT%H:%M:%S%z"), file=fh)

    ## Load the isd-lite data from files
    tempdf = None
    for year, fn in isd_files:
        print(fn)
        try:
            d = pd.read_fwf(os.path.join(datadir, fn),
                            colspecs=[x[0] for x in ISD_LITE_SPEC],
                            names=[x[1] for x in ISD_LITE_SPEC],
                            converters=OrderedDict([[x[1], x[2]] for x in ISD_LITE_SPEC]),
                            compression='infer')
            if tempdf is None:
                tempdf = d
            else:
                tempdf = pd.concat([tempdf, d])
        except FileNotFoundError:
            # logging.warning("File not found: '{}'".format(fn))
            pass
    # make a datetime index
    tempdf.index = pd.to_datetime(tempdf[['year', 'month', 'day', 'hour']])
    tempdf.index = tempdf.index.tz_localize('UTC')
    # sort
    tempdf.sort_index(inplace=True)

    # extract just air temp and rename columns
    t = pd.DataFrame(tempdf['air temp'])
    t.columns = ['AT']
    t.index.rename('datetime', inplace=True)

    if make_temperature_plots:
        fig = plt.figure()
        ax = fig.add_subplot(1,1,1)
        ax.plot(t.index, t['AT'], '.', label='raw temperatures')
        ax.set_title(station_callsign)
        ax.set_xlabel("datetime")
        ax.set_ylabel("air temp ['C]")
        plt.savefig(os.path.join(basedir, '{}_AT_orig.png'.format(station_callsign)))

#    # @TCC TMP -- Save before cleaning
#    t.to_hdf(os.path.join(basedir, '{}_AT_orig.h5'.format(station_callsign)),
#              'table', mode='w',
#              data_colums=True, complevel=5, complib='bzip2',
#              dropna=False)

    ### Cleaning temperatures ###

    ## deduplication (should not be needed with ISD-lite data)
    tmp = t[t.index.duplicated(keep=False)].sort_index()
    if len(tmp) > 0:
        logging.warning("{} duplicate temperature entries found".format(len(tmp)))
        t = t[~t.index.duplicated(keep='first')].sort_index()

    ## drop nan values
    tmp = t.shape[0]
    t.dropna(inplace=True)
    logging.info("Dropping {} NaN values".format(tmp-t.shape[0]))
    # keep a copy of the original temperatures for final plot if needed
    if make_temperature_plots:
        ot = t.copy(deep=True)

    ## Outlier removal
    cum_num = 0
    while True:
        if outlier_rolling_sigma_window > 0:
            sigma = t['AT'].rolling(window=outlier_rolling_sigma_window, center=True).std()
        else:
            sigma = 1
        diff = (t['AT']-t['AT'].rolling(window=outlier_rolling_median_window, center=True).median())/sigma
        outlier_mask = diff.abs() > outlier_thresh
        num = np.count_nonzero(outlier_mask)
        cum_num += num
        if num == 0:
            break
        else:
            logging.info("Outlier removal of {} points".format(num))
        t = t[~outlier_mask]
        if not outlier_multipass:
            break

    # plot showing what is being removed
    if make_temperature_plots:
        if cum_num > 0:
            fig = plt.figure()
            ax = fig.add_subplot(1,1,1)
            ax = ot[~ot.index.isin(t.index)].plot(ax=ax, linestyle='none', marker='o', color='r', zorder=8)
            # ax = ot.plot(ax=ax, linestyle='-', linewidth=1, marker=None, color='red')
            ax = t.plot(ax=ax, linestyle='none', marker='.', color='blue')
            ax.legend(['outlier', 'original', 'cleaned'])
            ax.set_title(station_callsign)
            ax.set_ylabel("air temp ['C]")
            fn = '{}_outlier.png'.format(station_callsign)
            fig.savefig(os.path.join(basedir, fn))

    ## "by-hand" fixes for individual datasets
    # @TCC TODO MAKE THIS PART OF CONFIG
    if station_callsign == 'KSNA': # KSNA (Orange County)
        # 2016-08-14 to 2016-08-15 overnight has some >0 values when they should be more like 19-20
        remove_spurious_temps(t, '< 0', '2016-08-14', '2016-08-15', inplace=True)
    if station_callsign == 'KSFO':
        remove_spurious_temps(t, '< 0', '1976-07-16', '1976-07-17', inplace=True)
    if station_callsign == 'KRIV':
        remove_spurious_temps(t, '< 0', '1995-11-15', '1995-11-15', inplace=True)

    ## Larger gap identification
    gaps_filename = os.path.join(basedir, "{}_AT_gaps.tsv".format(station_callsign))
    gaps = t.index.to_series().diff()[1:]  # @TCC had ot instead of t for some reason
    idx = np.flatnonzero(gaps > potentially_problematic_gap_size)
    prob_gaps = gaps[idx]
    # save to file for future reference
    with open(gaps_filename,'w') as fh:
        # output the gaps, biggest to smallest, to review
        print('#', station_callsign, t.index[0].isoformat(), t.index[-1].isoformat(), sep='\t', file=fh)
        print('# Potentially problematic gaps:', len(prob_gaps), file=fh)
        logging.info('Potentially problematic gaps: {}'.format(len(prob_gaps)))
        tmp = prob_gaps.sort_values(ascending=False)
        for i in range(len(tmp)):
            rng = [tmp.index[i]-tmp.iloc[i], tmp.index[i]]
            print(rng[0], rng[1], rng[1]-rng[0], sep='\t', file=fh)
    newidx = pd.date_range(start=t.index[0].round('d')+pd.Timedelta('0h'),
                           end=t.index[-1].round('d')-pd.Timedelta('1s'),
                           freq='1h', tz='UTC')
    # Simple linear interpolation
    at_interp_func = scipy.interpolate.interp1d(t.index.astype('int64').values,
                                           t['AT'].values,
                                           kind='linear',
                                           fill_value=np.nan, #(0,1)
                                           bounds_error=False)
    nt = pd.DataFrame({'AT':at_interp_func(newidx.astype('int64').values)},
                       index=newidx)
    # Fill those gaps using day-to-day (at same hour) interpolation
    gap_pad = pd.Timedelta('-10m') # contract the gaps a bit so we don't remove good/decent edge values
    t = nt.copy(deep=True) # operate on a copy so we can compare with nt
    # fill the gap ranges with nan (replacing the default interpolation)
    for i in range(len(prob_gaps)):
        rng = [prob_gaps.index[i]-prob_gaps.iloc[i], prob_gaps.index[i]]
        t[rng[0]-gap_pad:rng[1]+gap_pad] = np.nan
    # reshape so each row is a whole day's (24) data points
    rows = int(t.shape[0]/24)
    foo = pd.DataFrame(t.iloc[:rows*24].values.reshape((rows,24)))
    # simple linear interpolation
    foo.interpolate(metnod='time', limit=24*60, limit_direction='both', inplace=True)
    # reshape back
    t = pd.DataFrame({'AT':foo.stack(dropna=False).values}, index=t.index[:rows*24])

    # save final cleaned temperatures
    outfn = "{}_AT_cleaned".format(station_callsign)
    print("Saving cleaned temp data to:",outfn)
    t.to_hdf(os.path.join(basedir, outfn+'.h5'), 'table', mode='w',
              data_columns=True, complevel=5, complib='bzip2',
              dropna=False)

    # plot temperatures
    if make_temperature_plots:
        r1 = t.index[0]
        r2 = t.index[-1]
        fig = plt.figure()
        ax = fig.add_subplot(1,1,1)
        ax.plot(ot.loc[r1:r2].index, ot.loc[r1:r2]['AT'], linestyle='none', marker='.', label='raw')
        ax.plot(nt.loc[r1:r2].index, nt.loc[r1:r2]['AT'], linestyle='-', marker=None, lw=1, label='interpolated')
        for i in range(len(prob_gaps)):
            if i == 0: # only label first segment
                label = 'filled'
            else:
                label = ''
            rng = [tmp.index[i]-tmp.iloc[i], tmp.index[i]]
            ax.plot(t.loc[rng[0]:rng[1]].index, t.loc[rng[0]:rng[1]]['AT'], '.-', lw=1, color='r', label=label)
        ax.set_xlim((r1,r2))
        ax.set_xlabel('DateTime')
        ax.set_ylabel('Temperature [$\degree$C]')
        ax.set_title(station_callsign)
        ax.legend()
        fig.savefig(os.path.join(basedir, '{}_cleaning.png'.format(station_callsign)))

    return t

############################################################3

def setup_logger(verbose_level):
    fmt=('%(levelname)s %(asctime)s [%(module)s:%(lineno)s %(funcName)s] :: '
            '%(message)s')
    logging.basicConfig(format=fmt, level=max((0, min((MAX_LOGGING_LEVEL,
                        DEFAULT_LOGGING_LEVEL-(verbose_level*10))))))

def get_ISD_history_file(datadir, NOAA_ISD_ftp_host, save_file=True, force_download=False):
    ## get (fetch or load cached) isd-history.txt file
    tmp = None
    if not force_download: # try to load cached isd-history file
        try:
            tmp = read_isd_history_stations_list(os.path.join(datadir, 'isd-history.txt'))
        except FileNotFoundError:
            pass
    if tmp is None: # fetch
        logging.info("Fetching isd-history.txt")
        with ftplib.FTP(host=NOAA_ISD_ftp_host) as ftpconn:
            ftpconn.login()
            ftp_file = "/pub/data/noaa/isd-history.txt"
            # read the whole file and save it to a BytesIO (stream)
            response = io.BytesIO()
            try:
                ftpconn.retrbinary('RETR ' + ftp_file, response.write)
            except ftplib.error_perm as err:
                if str(err).startswith('550 '):
                    logging.error(err)
                else:
                    raise
        if save_file:
            response.seek(0)  # jump back to the beginning of the stream
            with open(os.path.join(datadir, 'isd-history.txt'),'wb') as fh:
                fh.write(response.getvalue())
        response.seek(0)  # jump back to the beginning of the stream
        tmp = read_isd_history_stations_list(response)
    return tmp

def read_isd_history_stations_list(filename, skiprows=22):
    """Read and parse stations information from isd_history.txt file"""
    fwfdef = (( ('USAF', (6, str)),
                ('WBAN', (5, str)),
                ('STATION NAME', (28, str)),
                ('CTRY', (4, str)),
                ('ST', (2, str)),
                ('CALL', (5, str)),
                ('LAT', (7, str)),
                ('LON', (8, str)),
                ('EVEV', (7, str)),
                ('BEGIN', (8, str)),
                ('END', (8, str)),
                ))
    names = []
    colspecs = []
    converters = {}
    i = 0
    for k,v in fwfdef:
        names.append(k)
        colspecs.append((i, i+v[0]+1))
        i += v[0]+1
        converters[k] = v[1]
    stdf = pd.read_fwf(filename, skiprows=skiprows,
                       names=names,
                       colspecs=colspecs,
                       converters=converters)
    return stdf

def remove_spurious_temps(ot, query_op, date1, date2=None, plot=True, inplace=False):
    if date2 is None:
        date2 = date1
    #ax = ot.loc[date1:date2].plot(ax=None, linestyle='-', marker='o') # plot
    out_t = ot.drop(ot.loc[date1:date2].query('AT {}'.format(query_op)).index, inplace=inplace)
    if inplace:
        out_t = ot
    #out_t.loc[date1:date2].plot(ax=ax, linestyle='-', marker='*') # plot'
    #ax.set_title("Remove AT {}, range=[{}:{}]".format(query_op, date1, date2))
    return out_t

#########################################################################
# Main loop hook... if run as script run main, else this is just a module
if __name__ == "__main__":
    sys.exit(Main(argv=None))


#! /usr/bin/env python3
import os
import requests

os.chdir('/data')

dictionary = {"title": "", "name": "", "date": "", "feedback": ""}

for files in os.listdir():
    with open(files,'r') as file:
        data = file.read().split('\n')
        dictionary["title"] = data[0];
        dictionary["name"] = data[1];
        dictionary["date"] = data[2];
        dictionary["feedback"] = data[3];
        response = requests.post("http://<IP of server>/feedback/",json=dictionary)
        print(response.status_code)
        if (response.status_code == 201):
          print("Feedback Added")
        else :
          response.raise_for_status()


import torch
import torch.nn as nn
from torch.autograd import Variable as V
import torch.nn.functional as F
import cv2
import numpy as np
from metrics import Evaluator
from loss import dice_bce_loss,edge_loss,DualTaskLoss
class MyFrame():
    def __init__(self, net,lr=2e-4, evalmode = False,batchsize=1):
        self.net = net().cuda()
        self.net = torch.nn.DataParallel(self.net, device_ids=range(torch.cuda.device_count()))
        self.optimizer = torch.optim.Adam(params=self.net.parameters(), lr=lr)
        #self.optimizer = torch.optim.RMSprop(params=self.net.parameters(), lr=lr)
        self.log_vars = nn.Parameter(torch.zeros((2))).cuda()
        self.segloss = dice_bce_loss()
        self.edgeloss=edge_loss()
        # self.Dualloss=DualTaskLoss()
        # self.edgeattention=ImageBasedCrossEntropyLoss2d(1)
        self.old_lr = lr
        # self.confusion_matrix = np.zeros((2,) * 2)
        self.batchsize=batchsize
        # self.evaluator = Evaluator(2)
        if evalmode:
            for i in self.net.modules():
                if isinstance(i, nn.BatchNorm2d):
                    i.eval()

    def set_input(self, img_batch, mask_batch=None,edge_batch=None,img_id=None):
        self.img = img_batch
        self.mask = mask_batch
        self.edge = edge_batch
        self.img_id = img_id

    def test_one_img(self, img):
        pred = self.net.forward(img)

        pred[pred>0.5] = 1
        pred[pred<=0.5] = 0

        mask = pred.squeeze().cpu().data.numpy()
        return mask

    def test_batch(self):
        self.forward(volatile=True)
        mask =  self.net.forward(self.img).cpu().data.numpy().squeeze(1)
        mask[mask>0.5] = 1
        mask[mask<=0.5] = 0

        return mask, self.img_id

    def test_one_img_from_path(self, path):
        img = cv2.imread(path)
        img = np.array(img, np.float32)/255.0 * 3.2 - 1.6
        img = torch.Tensor(img).cuda()

        mask = self.net.forward(img).squeeze().cpu().data.numpy()#.squeeze(1)
        mask[mask>0.5] = 1
        mask[mask<=0.5] = 0

        return mask

    def forward(self, volatile=False):
        self.img = self.img.cuda()
        if self.mask is not None:
            self.mask = self.mask.cuda()
            self.edge=self.edge.cuda()

    def optimize(self):
        self.forward()
        self.optimizer.zero_grad()
        pred,prededge = self.net.forward(self.img)
        # pred=F.log_softmax(pred, dim=1)
        # loss1 = self.segloss(self.mask, pred)
        precision1 = torch.exp(-self.log_vars[0]).cuda()
        loss1 = torch.sum(precision1 * (self.segloss(self.mask, pred)) ** 2. + self.log_vars[0], -1)
        # loss2 = self.edgeloss(prededge,self.edge)
        precision2 = torch.exp(-self.log_vars[1]).cuda()
        loss2= torch.sum(precision2 * (self.edgeloss(prededge,self.edge)) ** 2. + self.log_vars[1], -1)
        # loss3 = self.Dualloss(pred,self.mask.long(),128)
        # loss4 = self.edge_attention(pred,self.mask,self.labeledge)
        loss = loss1 + loss2
        # loss=loss1+loss2*20
        loss.backward()
        self.optimizer.step()
        return loss.data[0],loss1,loss2

    def save(self, path):
        torch.save(self.net.state_dict(), path)

    def load(self, path):
        self.net.load_state_dict(torch.load(path))

    def update_lr(self, new_lr, mylog, factor=False):
        if factor:
            new_lr = self.old_lr / new_lr
        for param_group in self.optimizer.param_groups:
            param_group['lr'] = new_lr

        #print >> mylog, 'update learning rate: %f -> %f' % (self.old_lr, new_lr)
        print ('update learning rate: %f -> %f' % (self.old_lr, new_lr))
        self.old_lr = new_lr

    # def getedge(self):
    #     mask=self.mask.cpu().numpy().transpose((0,2,3,1)).astype(np.uint8)
    #     canny = np.zeros((self.batchsize, 1, 1024, 1024))
    #     for i in range(self.batchsize):
    #         canny[i] = cv2.Canny(mask[i][0], 10, 100)
    #     canny = torch.from_numpy(canny).cuda().float()
    #     return canny

    # def edge_attention(self,input, target, edge):
    #     n, c, h, w = input.size()
    #     filler = torch.ones_like(target) * 255
    #     return self.edgeattention(input,torch.where(edge.max(1)[0] > 0.8, target, filler).long())

# class ImageBasedCrossEntropyLoss2d(nn.Module):
#
#     def __init__(self, classes, weight=None, size_average=True, ignore_index=255,
#                  norm=False, upper_bound=1.0):
#         super(ImageBasedCrossEntropyLoss2d, self).__init__()
#         self.num_classes = classes
#         self.nll_loss = nn.NLLLoss2d(weight, size_average, ignore_index)
#         self.norm = norm
#         self.upper_bound = upper_bound
#         self.batch_weights = False
#
#     def calculateWeights(self, target):
#         hist = np.histogram(target.flatten(), range(
#             self.num_classes + 1), normed=True)[0]
#         if self.norm:
#             hist = ((hist != 0) * self.upper_bound * (1 / hist)) + 1
#         else:
#             hist = ((hist != 0) * self.upper_bound * (1 - hist)) + 1
#         return hist
#
#     def forward(self, inputs, targets):
#         target_cpu = targets.data.cpu().numpy()
#         if self.batch_weights:
#             weights = self.calculateWeights(target_cpu)
#             self.nll_loss.weight = torch.Tensor(weights).cuda()
#
#         loss = 0.0
#         for i in range(0, inputs.shape[0]):
#             if not self.batch_weights:
#                 weights = self.calculateWeights(target_cpu[i])
#                 self.nll_loss.weight = torch.Tensor(weights).cuda()
#
#             loss += self.nll_loss(F.log_softmax(inputs[i].unsqueeze(0)),
#                                   targets[i].unsqueeze(0))
#         return loss


# coding: utf-8

import logging
import pickle

from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.decomposition import TruncatedSVD
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import Normalizer
from sklearn import metrics

from sklearn.cluster import KMeans

def doc_iter(filepath):
    with open(filepath, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            yield line

def lsa_kmeans(raw_doc, n_components, n_clusters):
    logger = logging.getLogger("lsa_kmeans")
    logger.info("Loading stop words...")
    stop_words = [word.strip() for word in open("stop_words/stop_words.txt", "r", encoding="utf-8")]

    logger.info("Compute tf-idf")
    vectorizer = TfidfVectorizer(stop_words=stop_words, ngram_range=(1, 1),
                                 max_df=1.0, min_df=1, max_features=None,
                                 vocabulary=None, norm="l2", use_idf=True,
                                 smooth_idf=True, sublinear_tf=True)
    tfidf = vectorizer.fit_transform(raw_doc)

    logger.info("SVD for lsa")
    svd = TruncatedSVD(n_components=n_components, algorithm="randomized",
                       n_iter=5, random_state=None, tol=0.0)
    normalizer = Normalizer(copy=False)
    lsa = make_pipeline(svd, normalizer)
    X = lsa.fit_transform(tfidf)

    logger.info("clustering documents")
    clf = KMeans(algorithm="auto", copy_x=True, init="k-means++", max_iter=300,
                 n_clusters=n_clusters, n_init=10, n_jobs=-2, precompute_distances="auto",
                 random_state=None, tol=0.0001, verbose=0)
    clf.fit(X)
    return clf

def save_model(clf, path):
    print("Save KMeans model to %s" % path)
    with open(path, "wb") as f:
        pickle.dump(clf, f)
    f.close()

def load_model(path):
    print("Load KMeans model to %s" % path)
    with open(path, "rb") as f:
        clf = pickle.load(f)
    f.close()
    return clf

def categories():
    return ["life_art_culture",                 # 生活、艺术与文化
            "chinese_culture",                  # 中华文化
            "social",                           # 社会
            "religion_belief",                  # 宗教及信仰
            "world",                            # 世界各地
            "social_sciences",                  # 人文与社会科学
            "natural_science",                  # 自然与自然科学
            "engineering_technology",           # 工程、技术与应用科学
            "general_works",                    # 常用列表
            "topic"                             # 主题首页
           ]


"""
1. Создает 50 zip-архивов, в каждом 100 xml файлов со случайными данными следующей структуры:

<root>
<var name=’id’ value=’<случайное уникальное строковое значение>’/>
<var name=’level’ value=’<случайное число от 1 до 100>’/>
<objects><object name=’<случайное строковое значение>’/>
<object name=’<случайное строковое значение>’/>
…
</objects>
</root>

В тэге objects случайное число (от 1 до 10) вложенных тэгов object.
"""
import os
import random
import string
import zipfile

import xml.etree.ElementTree as ET

def create_xml_file():
    root = ET.Element('root')

    id = ET.Element('var', name="id")
    id.text = ''.join(random.choices(string.ascii_lowercase, k=5))
    root.append(id)

    level = ET.Element('var', name="level")
    level.text = str(random.randint(1, 100))
    root.append(level)

    objects = ET.Element('objects')

    for i in range(0, random.randint(1, 10)):
        name = ''.join(random.choices(string.ascii_lowercase, k=10))
        objects.append(ET.Element('object', name=name))

    root.append(objects)
    indent(root)

    return root

def indent(elem, level=0):
    i = "\n" + level*"  "
    if len(elem):
        if not elem.text or not elem.text.strip():
            elem.text = i + "  "
        if not elem.tail or not elem.tail.strip():
            elem.tail = i
        for elem in elem:
            indent(elem, level+1)
        if not elem.tail or not elem.tail.strip():
            elem.tail = i
    else:
        if level and (not elem.tail or not elem.tail.strip()):
            elem.tail = i

def create_xml_files():
    for n in range(1, 101):
        ET.ElementTree(create_xml_file()).write(f"test_{n}.xml")

def create_folder(name: str):
    if not os.path.isdir(name):
        os.mkdir(name)
    os.chdir(name)

def create_archive(n: int):
    with zipfile.ZipFile(f"archive_{n}.zip", 'w', zipfile.ZIP_DEFLATED, True) as myzip:
        for root, dirs, files in os.walk(f"folder_{n}"):
            for file in files:
                myzip.write(os.path.join(root, file), arcname=os.path.join('', file))

def create_archives():
    create_folder("archives")

    for i in range(1, 51):
        create_folder(f"folder_{i}")
        create_xml_files()
        os.chdir("..")
        create_archive(i)
        os.remove(f"folder_{i}")

if __name__ == '__main__':
    create_archives()


# coding: utf-8

import time

from datetime import datetime
from functools import partial
from math import log10

class Ranking(object):
    def rank(self, iterable):
        return sorted(iterable, key=lambda obj: self.calculate_score(obj), reverse=True)

    def calculate_score(self, obj):
        raise NotImplemented

    @staticmethod
    def get_ranking_attr(obj, attr):
        if isinstance(attr, str):
            return getattr(obj, attr)

        if not callable(attr):
            return attr

        return attr(obj)

class RedditRanking(Ranking):
    def __init__(self, score_attr, created_at_attr):
        self.score_attr = score_attr
        self.created_at_attr = created_at_attr

    def calculate_score(self, obj):
        score = self.get_ranking_attr(obj, self.score_attr)
        order = log10(max(abs(score), 1))

        if score > 0:
            sign = 1
        elif score < 0:
            sign = -1
        else:
            sign = 0

        created_at = self.get_ranking_attr(obj, self.created_at_attr)
        seconds = time.mktime(created_at.timetuple()) - 1134028003

        return round(sign * order + seconds / 45000, 7)

class HackerNewsBaseRanking(Ranking):
    DEFAULT_GRAVITY = 1.8

    def __init__(self, votes_attr, created_at_attr, gravity=DEFAULT_GRAVITY):
        self.votes_attr = votes_attr
        self.created_at_attr = created_at_attr
        self.gravity = gravity

    def _get_object_info(self, obj):
        votes = self.get_ranking_attr(obj, self.votes_attr)

        created_at = self.get_ranking_attr(obj, self.created_at_attr)
        delta = datetime.now() - created_at
        hour_age = delta.days * 24 + delta.seconds / 3600.0

        return votes, hour_age

class HackerNewsRanking(HackerNewsBaseRanking):
    def calculate_score(self, obj):
        votes, hour_age = self._get_object_info(obj)
        return (votes - 1) / pow((hour_age + 2), self.gravity)

class HackerNewsExtendedRanking(HackerNewsBaseRanking):
    """
    gravity_conf:

    {'<feature>': [weight, score_limit]}

    """

    def __init__(
            self,
            votes_attr, created_at_attr,
            gravity_conf=None, min_gravity=1.5, max_gravity=2,
            annotate_items=False
    ):
        super(HackerNewsExtendedRanking, self).__init__(votes_attr, created_at_attr)

        self.MIN_GRAVITY = min_gravity
        self.MAX_GRAVITY = max_gravity
        self.GRAVITY_RANGE = self.MAX_GRAVITY - self.MIN_GRAVITY

        if gravity_conf is None:
            self.gravity = lambda obj: self.DEFAULT_GRAVITY

        else:
            self.gravity = partial(self._calculate_gravity, gravity_conf=gravity_conf)
            self.feature_max_impacts = self._get_feature_max_impacts(gravity_conf)

        self.annotate_items = annotate_items

    def _get_feature_max_impacts(self, gravity_conf):
        feature_max_impact_map = dict((feature, 0) for feature in gravity_conf)

        sum_weight = sum([weight for weight, _ in gravity_conf.values()])
        for feature, (weight, _) in gravity_conf.items():
            feature_max_impact_map[feature] = float(weight) / sum_weight * self.GRAVITY_RANGE

        return feature_max_impact_map

    def _calculate_gravity(self, obj, gravity_conf=None):
        gravity_impact = 0

        for feature, (weight, score_limit) in gravity_conf.items():
            feature_score = self.get_ranking_attr(obj, feature)

            if feature_score >= score_limit:
                gravity_impact += self.feature_max_impacts[feature]
                continue

            gravity_impact += float(feature_score) / score_limit * self.feature_max_impacts[feature]

        obj.gravity = self.MAX_GRAVITY - gravity_impact
        return self.MAX_GRAVITY - gravity_impact

    def calculate_score(self, obj):
        votes, hour_age = self._get_object_info(obj)
        gravity = self.gravity(obj)
        gravity_coeff = obj.gravity_coeff if hasattr(obj, 'gravity_coeff') else 1

        rank = (votes - 1) / pow((hour_age + 2), gravity * gravity_coeff)

        if self.annotate_items:
            obj.trending_rank = rank
            obj.trending_gravity = gravity

        return rank


from shapely.geometry import Point, box

class Environment(object):

    def __init__(self, width, height):
        self.robots = []
        self.obstacles = []
        self.width = width
        self.height = height

        self.boundary = box(0, 0, width, height)

    def addRobot(self, x, y):
        poly = Point(x, y).buffer(30.0)

        self.robots.append(poly)

    def addObstable(self, x, y, width, height):
        self.obstacles.append(box(x, y, x + width, y + height))

    def colides(self, ref):

        point = Point(ref)

        if self.boundary.intersects(point):
            for robot in self.robots:
                if robot.intersects(point):
                    # print("robot colides")
                    return True

            for obstacle in self.obstacles:
                if obstacle.intersects(point):
                    # print("obstacle colides")
                    return True

        else:
            raise Exception('Invalid point')

        return False


import os
import pickle

import torch
import torch.nn as nn
import torch.nn.functional as F
import learn2learn as l2l
import numpy as np
from tqdm import tqdm

from argparser import parse_args
from activations import Sparsemax
from utils import fix_seed, get_cv_score, check_both_none_or_not_none
from metrics import cluster_acc, cluster_ari

def run(args=None):
    args = parse_args(args)
    device = torch.device(args.device)
    fix_seed(args.seed)
    if not os.path.exists(args.exp_path):
        os.makedirs(args.exp_path)

    phi1 = np.load(args.phi1_path).astype(np.float32)
    phi2 = np.load(args.phi2_path).astype(np.float32)
    assert check_both_none_or_not_none(args.phi1_path_val, args.phi2_path_val)
    if args.phi1_path_val is not None:
        phi1_val = np.load(args.phi1_path_val).astype(np.float32)
        phi2_val = np.load(args.phi2_path_val).astype(np.float32)
    else:
        phi1_val = np.copy(phi1)
        phi2_val = np.copy(phi2)
    y_true_val = np.load(args.gt_labels_path)
    assert phi1.shape[0] == phi2.shape[0]
    assert phi1_val.shape[0] == phi2_val.shape[0]
    assert phi1_val.shape[0] == y_true_val.shape[0]
    n_train = phi1.shape[0]
    d1, d2 = phi2.shape[1], phi1.shape[1]
    subset_size = min(n_train, args.subset_size)

    # Instantiate linear layer for the inner optimization (Equation 5)
    inner_linear = nn.Linear(d1, args.k, bias=True).to(device)
    inner_linear = l2l.algorithms.MAML(inner_linear, lr=args.inner_lr)

    # Instantiate task encoder with orthogonal weights parametrization (Equation 3)
    task_encoder = nn.Linear(d2, args.k, bias=False).to(device)
    task_encoder = nn.utils.parametrizations.orthogonal(task_encoder)

    all_parameters = list(task_encoder.parameters())
    optimizer = torch.optim.Adam(all_parameters, lr=args.outer_lr)
    scheduler = torch.optim.lr_scheduler.MultiStepLR(
        optimizer, milestones=[100, 200],
        gamma=0.1 if args.anneal else 1.0
    )
    old_lr = args.outer_lr
    tau = args.tau
    sparsemax_act = Sparsemax(dim=1)
    for i in tqdm(range(args.num_iters)):
        optimizer.zero_grad()
        mean_train_error = 0.0
        mean_valid_error = 0.0
        mean_valid_acc = 0.0
        mean_train_acc = 0.0
        mean_label_dist = 0.0
        mean_sparsity = 0.0

        for j in range(args.num_subsets):
            # Sample X_tr and X_te
            subset = np.random.choice(n_train, size=subset_size, replace=False)
            subset_tr = subset[:int(subset_size * args.train_fraction)]
            subset_te = subset[int(subset_size * args.train_fraction):]

            phi1_tr = torch.from_numpy(phi1[subset_tr]).to(device)
            phi1_te = torch.from_numpy(phi1[subset_te]).to(device)
            phi2_tr = torch.from_numpy(phi2[subset_tr]).to(device)
            phi2_te = torch.from_numpy(phi2[subset_te]).to(device)

            # Get labels using current task encoder
            task_labels_tr = sparsemax_act(task_encoder(phi1_tr) / tau)
            task_labels_te = sparsemax_act(task_encoder(phi1_te) / tau)
            task_labels_all = torch.cat((task_labels_tr, task_labels_te))

            """
            Perform inner optimization from the random initialization or
            from fixed w0 (corresponds to Cold Start BLO for Equation 5)
            """

            if args.rand_init:
                inner_linear.reset_parameters()
            learner = inner_linear.clone()

            for step in range(args.adaptation_steps):
                train_error = F.cross_entropy(learner(phi2_tr), task_labels_tr)
                learner.adapt(train_error)

            # Compute HUME's objective (Equation 7)
            label_dist = task_labels_all.mean(0)
            entr = torch.special.entr(label_dist)
            valid_error = F.cross_entropy(learner(phi2_te), task_labels_te)

            # Accumulate gradients across args.num_subsets
            (valid_error - args.H_reg * entr.sum()).backward()

            # Compute training stats
            mean_train_error += train_error.item()
            mean_train_acc += torch.eq(
                learner(phi2_tr).argmax(1),
                task_labels_tr.argmax(1)
            ).float().mean().item()
            mean_valid_error += valid_error.item()
            mean_valid_acc += torch.eq(
                learner(phi2_te).argmax(1),
                task_labels_te.argmax(1)
            ).float().mean().item()
            mean_label_dist += label_dist.detach().cpu().numpy()
            mean_sparsity += task_labels_all[torch.arange(task_labels_all.shape[0]),
                                             task_labels_all.argmax(1)].mean().item()

        # Average gradients over args.num_subsets and update the task encoder parameters
        for p in all_parameters:
            p.grad.data.mul_(1.0 / args.num_subsets)
            print(f"Grad norm: {torch.norm(p.grad.data).item()}")
        nn.utils.clip_grad_norm_(task_encoder.parameters(), 1.0)
        optimizer.step()
        scheduler.step()

        # Anneal step size and temperature
        if scheduler.get_last_lr()[0] != old_lr:
            print("Annealed Learning rate")
            old_lr = scheduler.get_last_lr()[0]
            print("Annealed Temperature")
            tau = tau / 10
            print()

        # Print train stats
        print("Train stats:")
        print(f"Mean TrainError {mean_train_error / args.num_subsets}")
        print(f"Mean ValidError {mean_valid_error / args.num_subsets}")
        print(f"Mean TrainAcc {mean_train_acc / args.num_subsets}")
        print(f"Mean ValidAcc {mean_valid_acc / args.num_subsets}")
        print(f"Mean Sparsity {mean_sparsity / args.num_subsets}")
        print("Mean Label Dist:", mean_label_dist / args.num_subsets)
        print()

        # Print val stats
        out_all_val = task_encoder(torch.from_numpy(phi1_val).to(device))
        preds_all_val = torch.argmax(out_all_val, dim=1).detach().cpu().numpy()
        print("Val metrics:")
        print("Num found clusters:", len(np.unique(preds_all_val)))
        print(f"Cluster ACC epoch {i}:", cluster_acc(preds_all_val, y_true_val))
        print(f"Cluster ARI epoch {i}:", cluster_ari(preds_all_val, y_true_val))
        print()

        if args.save_all:
            torch.save(task_encoder.state_dict(), args.exp_path + f"linear_task_{i}_{args.seed}.pt")

    # Compute cross-validation accuracy w.r.t. found task and save the results
    out_all_val = task_encoder(torch.from_numpy(phi1_val).to(device))
    task_val = torch.argmax(out_all_val, dim=1).detach().cpu().numpy()
    final_cv_score = get_cv_score(phi2_val, task_val)
    with open(args.exp_path + f"results_{args.seed}.pickle", "wb") as handle:
        pickle.dump({"CV_Score": final_cv_score}, handle, protocol=pickle.HIGHEST_PROTOCOL)
    torch.save(task_encoder.state_dict(), args.exp_path + f"linear_task_{args.seed}.pt")

if __name__ == '__main__':
    run()


import math

import torch
from torch.nn import ReplicationPad3d
import torchvision

from common.i3d_utils import inflator as inflate

bb_to_tv_function = {"res18": torchvision.models.resnet18,
                    "res34": torchvision.models.resnet34,
                    "res50": torchvision.models.resnet50,
                    "res101": torchvision.models.resnet101
                }
bb_out_dims = {"res18": 512, # Backbone_out_dims
                "res34": 512,
                "res50": 2048,
                "res101": 2048
            }

class I3ResNet(torch.nn.Module):
    def __init__(self, config):
        """
        Args:
            conv_class: Whether to use convolutional layer as classifier to
                adapt to various number of frames
        """
        super(I3ResNet, self).__init__()
        frame_nb=16
        conv_class = config.NETWORK.CONV_CLASS

        resnet2d = bb_to_tv_function[config.NETWORK.BACKBONE](pretrained=config.NETWORK.BACKBONE_LOAD_PRETRAINED)
        self.conv_class = conv_class

        self.conv1 = inflate.inflate_conv(
            resnet2d.conv1, time_dim=3, time_padding=1, center=True)
        self.bn1 = inflate.inflate_batch_norm(resnet2d.bn1)
        self.relu = torch.nn.ReLU(inplace=False)
        self.maxpool = inflate.inflate_pool(
            resnet2d.maxpool, time_dim=3, time_padding=1, time_stride=2)

        self.layer1 = inflate_reslayer(resnet2d.layer1, config.NETWORK.BACKBONE)
        self.layer2 = inflate_reslayer(resnet2d.layer2, config.NETWORK.BACKBONE)
        self.layer3 = inflate_reslayer(resnet2d.layer3, config.NETWORK.BACKBONE)
        self.layer4 = inflate_reslayer(resnet2d.layer4, config.NETWORK.BACKBONE)

        if conv_class:
            self.avgpool = inflate.inflate_pool(resnet2d.avgpool, time_dim=1)
            self.classifier = torch.nn.Conv3d(
                in_channels=bb_out_dims[config.NETWORK.BACKBONE],
                out_channels=config.NETWORK.NUM_CLASSES,
                kernel_size=(1, 1, 1),
                bias=True)
        else:
            raise NotImplementedError
            final_time_dim = int(math.ceil(frame_nb / 16))
            self.avgpool = inflate.inflate_pool(
                resnet2d.avgpool, time_dim=final_time_dim)
            self.fc = inflate.inflate_linear(resnet2d.fc, 1)

    def forward(self, x, other1, other2):
        x = x.permute(0,2,1,3,4)
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)

        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)

        if self.conv_class:
            x = self.avgpool(x)
            x = self.classifier(x)
            x = x.squeeze(3)
            x = x.squeeze(3)
            x = x.mean(2)
        else:
            x = self.avgpool(x)
            x_reshape = x.view(x.size(0), -1)
            x = self.fc(x_reshape)
        return x

def inflate_reslayer(reslayer2d, backbone):
    reslayers3d = []
    for layer2d in reslayer2d:
        layer3d = Bottleneck3d(layer2d, backbone)
        reslayers3d.append(layer3d)
    return torch.nn.Sequential(*reslayers3d)

class Bottleneck3d(torch.nn.Module):
    def __init__(self, bottleneck2d, backbone):
        super(Bottleneck3d, self).__init__()

        spatial_stride = bottleneck2d.conv2.stride[0]

        self.conv1 = inflate.inflate_conv(
            bottleneck2d.conv1, time_dim=1, center=True)
        self.bn1 = inflate.inflate_batch_norm(bottleneck2d.bn1)

        self.conv2 = inflate.inflate_conv(
            bottleneck2d.conv2,
            time_dim=3,
            time_padding=1,
            time_stride=spatial_stride,
            center=True)
        self.bn2 = inflate.inflate_batch_norm(bottleneck2d.bn2)

        if backbone not in ["res18", "res34"]:
            self.conv3 = inflate.inflate_conv(
                bottleneck2d.conv3, time_dim=1, center=True)
        else:
            self.conv3 = None

        if backbone not in ["res18", "res34"]:
            self.bn3 = inflate.inflate_batch_norm(bottleneck2d.bn3)
        else:
            self.bn3 = None

        self.relu = torch.nn.ReLU(inplace=False)

        if bottleneck2d.downsample is not None:
            self.downsample = inflate_downsample(
                bottleneck2d.downsample, time_stride=spatial_stride)
        else:
            self.downsample = None

        assert (self.conv3 != None and self.bn3 != None) or (self.conv3 == None and self.bn3 == None)
        self.stride = bottleneck2d.stride

    def forward(self, x):
        residual = x
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)

        if self.conv3 != None:
            out = self.conv3(out)
        if self.bn3 != None:
            out = self.bn3(out)

        if self.downsample is not None:
            residual = self.downsample(x)

        out += residual
        out = self.relu(out)
        return out

def inflate_downsample(downsample2d, time_stride=1):
    downsample3d = torch.nn.Sequential(
        inflate.inflate_conv(
            downsample2d[0], time_dim=1, time_stride=time_stride, center=True),
        inflate.inflate_batch_norm(downsample2d[1]))
    return downsample3d


import csv
import math
import json
import numpy as np
import random
import os
import argparse
from tqdm import tqdm
import torch
import torch.nn.functional as F
from transformers import AutoTokenizer, AutoModel
from torch.utils.data import Dataset, DataLoader, Sampler, default_collate
from torch.optim import AdamW, Adam, SGD
from torch.optim.lr_scheduler import LambdaLR
from typing import List
from copy import deepcopy
from contextlib import nullcontext
from collections import defaultdict

from guidelines.guidelines import Guidelines
from models.guided_bert_clues import GuidedBertClues
from experiments.utils import RunningMean, Mean

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model_name_or_path", type=str,
                        default="bert-base-uncased")
    parser.add_argument("--output_dir", type=str, default=None)
    parser.add_argument("--data_dir", type=str, default=None)
    parser.add_argument("--num_train_epochs", type=int, default=20)
    parser.add_argument("--learning_rate", type=float, default=3e-5)
    parser.add_argument("--weight_decay", type=float, default=1e-2)
    parser.add_argument("--seed", type=int, default=1)
    parser.add_argument("--fp16", action="store_true")
    parser.add_argument("--per_device_train_batch_size", type=int, default=16)
    parser.add_argument("--per_device_eval_batch_size", type=int, default=16)
    parser.add_argument("--gradient_accumulation_steps", type=int, default=1)
    parser.add_argument("--eval_accumulation_steps", type=int,
                        default=1)

    parser.add_argument("--guided_model", type=str, default="main",
                        choices=["main"])
    parser.add_argument("--phase", type=str, default="program",
                        choices=["program", "program-only", "similarity",
                                 "attributes-plain"])
    parser.add_argument("--subtask", type=str, default="clues_real")
    # parser.add_argument("--program_length", type=int, default=10)
    parser.add_argument("--single_task", type=str, default=None)
    parser.add_argument("--zero_shot", action="store_true")
    parser.add_argument("--individual_scores", action="store_true")
    parser.add_argument("--normalize_datasets", action="store_true")
    parser.add_argument("--optimizer", type=str, default="AdamW")
    parser.add_argument("--scheduler", type=str, default="")
    parser.add_argument("--encoder_feature_layer", type=str,
                        default="transformer-encoder")
    parser.add_argument("--guideline_feature_layer", type=str,
                        default="transformer-encoder")
    parser.add_argument("--encoder_feature_layer_num_head",
                        type=int, default=8)
    parser.add_argument("--guideline_feature_layer_num_head",
                        type=int, default=8)
    parser.add_argument("--model_version", type=str, default="2")
    parser.add_argument("--batch_per_epoch", type=int, default=2500)
    parser.add_argument("--remove_label_text", action="store_true")
    parser.add_argument("--num_feature_layer", type=int, default=1)
    parser.add_argument("--data_variance", type=str, default=None)
    parser.add_argument("--mask_out_templates", type=int,
                        nargs="*", default=[])

    args = parser.parse_args()
    seed = args.seed
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)
    random.seed(seed)
    os.makedirs(args.output_dir, exist_ok=True)
    print(args)
    return args

def get_data(data_dir, subtask, single_task, zero_shot, remove_label_text):
    def load_jsonl(filename):
        with open(filename, "r") as f:
            return [json.loads(line) for line in f.readlines()]

    def load_json(filename):
        with open(filename, "r") as f:
            return json.load(f)

    def load_csv(filename):
        with open(filename, "r") as f:
            reader = csv.reader(f)
            return list(reader)

    def revert_explanations(explanations, entails, target_key):
        return {
            _key: [
                (_expl, entails is None or entails[_expl] == "entail")
                for _expl, _key2 in explanations.items()
                if _key2 == _key
            ]
            for _key in target_key
        }

    def remove_label_text_fn(explanations):
        new_explanations = {}
        for _key, sentences in explanations.items():
            new_explanations[_key] = []
            for _s, _l in sentences:
                _s = _s.replace(_key, "this")
                _s = _s.replace(_key[0].lower()+_key[1:], "this")
                _s = _s.replace(_key[0].upper()+_key[1:], "this")
                _s = _s.replace(_key.lower(), "this")
                _s = _s.replace(_key.upper(), "this")
                new_explanations[_key].append((_s, _l))
        return new_explanations

    def get_data_clues_real():
        this_data = {
            sub2task: {
                split: load_jsonl(
                    f"{data_dir}/clues_real/{sub1task}/"
                    f"{sub2task}/{split}.jsonl")
                for split in ("train", "dev", "test")
            }
            for sub1task in os.listdir(f"{data_dir}/clues_real")
            for sub2task in os.listdir(f"{data_dir}/clues_real/{sub1task}")
        }
        zero_shot_tasks = (
            "banknote-authentication", "tic-tac-toe-endgame", "car-evaluation",
            "contraceptive-method-choice", "indian-liver-patient",
            "travel-insurance", "entrepreneur-competency",
            "award-nomination-result", "coin-face-value", "coin-metal",
            "driving-championship-points", "election-outcome", "hotel-rating",
            "manifold-orientability", "soccer-club-region",
            "soccer-league-type",
        )
        wikipedia_tasks = (
            "award-nomination-result", "coin-metal", "soccer-league-type",
            "driving-championship-points", "hotel-rating", "coin-face-value",
            "election-outcome", "manifold-orientability", "soccer-club-region",
        )
        if zero_shot:
            for _task, _task_data in this_data.items():
                if _task in zero_shot_tasks:
                    if _task in wikipedia_tasks:
                        _task_data["test"] = (
                            _task_data["train"] + _task_data["dev"] +
                            _task_data["test"])
                    _task_data["train"] = []
                    _task_data["dev"] = []
                else:
                    _task_data["train"] = \
                        _task_data["train"] + _task_data["test"]
                    _task_data["test"] = []
        explanations = load_json(
            f"{data_dir}/clues_real_utils/real_lbl_indicator_map.json")
        entail = load_json(
            f"{data_dir}/clues_real_utils/real_exp_entail_indicator.json")
        task_details = load_json(
            f"{data_dir}/clues_real_utils/real_task_details.json")
        assert set(explanations.keys()) == set(this_data.keys())
        for sub2task, _expl in explanations.items():
            _entail = entail[sub2task]
            target_key = task_details[sub2task]["task_lbls"]
            this_data[sub2task]["target_key"] = target_key
            _expl = revert_explanations(
                _expl, _entail, target_key)
            if remove_label_text:
                _expl = remove_label_text_fn(_expl)
            this_data[sub2task]["explanations"] = _expl
        if single_task is not None:
            this_data = {single_task: this_data[single_task]}
        return this_data

    def get_data_clues_syn():
        this_data = {
            sub1task: {
                split: load_jsonl(
                    f"{data_dir}/clues_syn/{sub1task}/{split}.jsonl")
                for split in ("train", "valid", "test")
            }
            for sub1task in os.listdir(f"{data_dir}/clues_syn")
        }
        if zero_shot:
            this_data = {
                sub1task: {
                    "train": this_task_data["train"] + this_task_data["valid"]
                    + this_task_data["test"],
                    "valid": [],
                    "test": []
                } for sub1task, this_task_data in this_data.items()
            }
        targets = {
            "bird species": ["wug", "blicket", "dax", "toma", "pimwit", "zav",
                             "speff", "tulver", "gazzer", "fem", "fendle",
                             "tupa",
                             "not blicket", "not dax", "not fem", "not fendle",
                             "not pimwit", "not speff", "not toma", "not tupa",
                             "not wug", "not zav",
                             ],
            "animal species": ["wug", "blicket", "dax", "toma", "pimwit",
                               "zav", "speff", "tulver", "gazzer", "fem",
                               "fendle", "tupa",
                               "not blicket", "not dax", "not fendle",
                               "not speff", "not toma", "not tulver",
                               "not tupa", "not wug", "not zav"],
            "rain tomorrow": ["yes", "no", "not yes", "not no"],
            "final position": ["1", "2", "3", "4", "Not Qualified",
                               "not 1", "not 2", "not 3", "not 4",
                               "not Not Qualified"],
            "relevance score": ["1", "2", "3", "4", "5",
                                "not 1", "not 2", "not 4"]
        }
        for sub1task in os.listdir(f"{data_dir}/clues_syn"):
            explanations = [
                _expl[0] for _expl in load_csv(
                    f"{data_dir}/clues_syn/{sub1task}/explanations.csv"
                )
            ]
            target_type = list(set(targets.keys()).intersection(
                set(this_data[sub1task]["train"][0].keys())))
            assert len(target_type) == 1
            target_type = target_type[0]
            target_key = targets[target_type]
            this_data[sub1task]["target_key"] = target_key
            for split in ("train", "valid", "test"):
                for item in this_data[sub1task][split]:
                    item["lbl"] = item.pop(target_type)
                    item["target_key"] = target_key
            explanation_keys = []
            for _expl in explanations:
                this_target_key = [_key for _key in target_key if
                                   _expl.endswith(_key) or
                                   _expl.endswith(_key+".")]
                assert len(this_target_key) in [1, 2], f"{_expl}, {target_key}"
                if len(this_target_key) == 2:
                    assert this_target_key[1] == "not " + this_target_key[0]
                    this_target_key.pop(0)
                explanation_keys.append(this_target_key[0])
            explanations = {_expl: _key for _expl, _key in
                            zip(explanations, explanation_keys)}
            this_data[sub1task]["explanations"] = revert_explanations(
                explanations, None, target_key)
        return this_data

    def mix_data(list_of_data, valid_to_dev=True):
        final_data = {}
        id_iter = iter(range(1000000))
        for each_data in list_of_data:
            for sub1task, this_task_data in each_data.items():
                if valid_to_dev and "valid" in this_task_data:
                    this_task_data["dev"] = this_task_data.pop("valid")
                final_data[str(next(id_iter))+"_"+sub1task] = this_task_data
        return final_data

    if subtask == "clues_real":
        return get_data_clues_real()
    elif subtask == "clues_syn":
        return get_data_clues_syn()
    else:
        _clues_real = get_data_clues_real()
        _clues_syn = get_data_clues_syn()
        if subtask == "real:syn=1:1":
            return mix_data([_clues_real, _clues_syn])
        elif subtask == "real:syn=1:3":
            return mix_data([_clues_real, _clues_syn, _clues_syn, _clues_syn])
        elif subtask == "real:syn=1:5":
            return mix_data([_clues_real, _clues_syn, _clues_syn, _clues_syn,
                             _clues_syn, _clues_syn])
    raise NotImplementedError()

def datum_to_sentence(datum, data_variance):
    def piece_fn(_key, _value):
        if data_variance is None:
            return f"{_key} is {_value}"
        elif data_variance.startswith("replace:"):
            return f"{_key} {data_variance[8:]} {_value}"
        elif data_variance.startswith("append:") or \
                data_variance.startswith("prepend"):
            return f"{_key} is {_value}"

    fat = (
        "For this item, " +
        ", ".join([
            piece_fn(_key, _value)
            for _key, _value in datum.items()
            if _key not in ['lbl', 'target_key']
        ])
    )
    if data_variance is None or data_variance.startswith("replace"):
        fat = fat + "."
    elif data_variance.startswith("append"):
        fat = fat + data_variance.split(":")[-1]
    elif data_variance.startswith("prepend"):
        fat = data_variance.split(":")[-1] + fat
    return fat

class CLUES_Dataset(Dataset):
    def __init__(self, data, tokenizer, split, data_variance):
        self.data = data
        self.tokenizer = tokenizer
        self.split = split
        self.task_names = list(data.keys())
        self.task_names.sort()
        self.task_data_num = {
            _task: len(data[_task][split])
            for _task in self.task_names
        }
        self.datum_indices = [
            _task
            for _task in self.task_names
            for _j in data[_task][split]
        ]
        self.task_start = {_task: self.datum_indices.index(_task)
                           for _task in self.task_names if _task in
                           self.datum_indices}
        self.num_data = len(self.datum_indices)
        self.data_variance = data_variance

    def __len__(self):
        return self.num_data

    def __getitem__(self, index):
        task_name = self.datum_indices[index]
        datum = self.data[task_name][self.split][
            index-self.task_start[task_name]]
        datum_sentence = datum_to_sentence(datum, self.data_variance)
        label = datum["lbl"]
        target_key = datum["target_key"]
        guidelines = self.data[task_name]["guidelines"]
        return datum_sentence, target_key.index(str(label)), task_name, \
            guidelines.inputs, guidelines.target_text_matrix, \
            guidelines.entail_indicators

class CLUES_BatchSampler(Sampler[List[int]]):
    def __init__(self, dataset, batch_size, shuffle, normalize_datasets,
                 batch_per_epoch):
        self.dataset = dataset
        self.batch_size = batch_size
        self.shuffle = shuffle
        self.normalize_datasets = normalize_datasets
        assert shuffle or not normalize_datasets
        self.all_batches = []
        task_num_batches = defaultdict(lambda: 0)
        task_names = []
        batch = []
        prev_task_name = None
        for idx, (_, _, task_name, _, _, _) in enumerate(dataset):
            if (prev_task_name is None or task_name == prev_task_name) \
                    and len(batch) < self.batch_size:
                batch.append(idx)
            else:
                self.all_batches.append(batch)
                task_num_batches[prev_task_name] += 1
                task_names.append(prev_task_name)
                batch = [idx]
            prev_task_name = task_name
        self.all_batches.append(batch)
        task_num_batches[task_name] += 1
        task_names.append(task_name)
        self.sampling_weights = [
            1/len(task_num_batches)/task_num_batches[_task_name]
            for _task_name in task_names
        ]
        self.batch_per_epoch = batch_per_epoch

    def __iter__(self):
        all_batches = deepcopy(self.all_batches)
        if self.shuffle:
            if not self.normalize_datasets:
                random.shuffle(all_batches)
            else:
                all_batches = np.random.choice(
                    np.array(all_batches, dtype=object),
                    p=self.sampling_weights,
                    size=self.batch_per_epoch
                )
        for batch in all_batches:
            yield batch

    def __len__(self):
        if self.shuffle and self.normalize_datasets:
            return self.batch_per_epoch
        else:
            return len(self.all_batches)

def CLUES_collate(batch):
    collated = default_collate(batch)
    collated = [
        collated[0],
        collated[1],
        collated[2][0],
        {_key: _value[0] for _key, _value in collated[3].items()},
        collated[4][0],
        torch.stack([x[0] for x in collated[5]]),
    ]
    # inputs, labels, task_name, guideline_inputs,
    # guideline_matrix, entail_indicators
    return collated

def postprocess_batch(batch, tokenizer, device):
    inputs = tokenizer(list(batch[0]), padding=True, max_length=256,
                       return_tensors="pt")
    batch[0] = {
        _key: _value.to(device)
        for _key, _value in inputs.items()
    }
    if "token_type_ids" not in batch[0]:
        batch[0]["token_type_ids"] = None
    batch[1] = batch[1].to(device)
    batch[3] = {
        "guideline_"+_key: _value.to(device)
        for _key, _value in batch[3].items()
    }
    if "guideline_token_type_ids" not in batch[3]:
        batch[3]["guideline_token_type_ids"] = None
    return batch

def one_epoch(dataloader, model, optimizer, scheduler, tokenizer,
              device, split, i_epoch, individual_scores):
    def mean_of_mean(means):
        values = [_mean.value for _mean in means]
        return sum(values) / len(values)

    print(f"{split} epoch {i_epoch}")
    pbar = tqdm(dataloader, leave=False)
    if split == "train":
        loss_mean = RunningMean()
        acc_mean = RunningMean()
    elif not individual_scores:
        loss_mean = Mean()
        acc_mean = Mean()
    else:
        each_loss_mean = defaultdict(lambda: Mean())
        each_acc_mean = defaultdict(lambda: Mean())

    for batch in pbar:
        optimizer.zero_grad()
        batch = postprocess_batch(batch, tokenizer, device)
        inputs, labels, task_name, guideline_inputs, target_text_matrix, \
            entail_indicators = batch
        batch_size = inputs["input_ids"].shape[0]
        with torch.no_grad() if split != "train" else nullcontext():
            scores = model(**inputs, **guideline_inputs,
                           target_text_matrix=target_text_matrix,
                           entail_indicators=entail_indicators)
        loss = F.cross_entropy(
            scores, labels, reduction="none")
        acc = (scores.argmax(1) == labels).float()
        if split == "train":
            (loss.sum() / batch_size).backward()
            optimizer.step()
            scheduler.step()
        if not individual_scores:
            loss_mean.update(loss.detach().tolist())
            acc_mean.update(acc.detach().tolist())
        else:
            each_loss_mean[task_name].update(loss.detach().tolist())
            each_acc_mean[task_name].update(acc.detach().tolist())
            loss_mean = mean_of_mean(each_loss_mean.values())
            acc_mean = mean_of_mean(each_acc_mean.values())
        pbar.set_description(f"loss: {loss_mean:.4f}, acc: {acc_mean:.4f}")
    print(f"{split} epoch-{i_epoch}:"
          f" average loss: {loss_mean}, average accuracy: {acc_mean}")
    if not individual_scores:
        loss_mean = loss_mean.value
        acc_mean = acc_mean.value
    else:
        print(f"individual loss: {dict(each_loss_mean)}")
        print(f"individual acc: {dict(each_acc_mean)}", flush=True)
    return loss_mean, acc_mean

def save_model(model, optimizer, scheduler, start_epoch, filename):
    print(f"saving model to {filename}")
    torch.save(
        [model.state_dict(), optimizer.state_dict(), scheduler.state_dict(),
         start_epoch],
        filename
    )

def main(args):
    data = get_data(args.data_dir, args.subtask, args.single_task,
                    args.zero_shot, args.remove_label_text)
    device = torch.device("cpu") if torch.cuda.device_count() == 0 else \
        torch.device("cuda:0")
    tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path)
    bert = AutoModel.from_pretrained(args.model_name_or_path)

    for task in data.values():
        task["guidelines"] = Guidelines(
            task["explanations"], tokenizer, task["target_key"], True)

    train_dataset = CLUES_Dataset(data, tokenizer, "train", args.data_variance)
    dev_dataset = CLUES_Dataset(
        data, tokenizer, "valid" if args.subtask == "clues_syn" else "dev",
        args.data_variance)
    test_dataset = CLUES_Dataset(data, tokenizer, "test", args.data_variance)
    train_dataloader = DataLoader(
        train_dataset, collate_fn=CLUES_collate,
        batch_sampler=CLUES_BatchSampler(
            train_dataset, args.per_device_train_batch_size,
            True, args.normalize_datasets, args.batch_per_epoch))
    dev_dataloader = DataLoader(
        dev_dataset, collate_fn=CLUES_collate,
        batch_sampler=CLUES_BatchSampler(
            dev_dataset, args.per_device_eval_batch_size, False, False,
            args.batch_per_epoch
        ))
    test_dataloader = DataLoader(
        test_dataset, collate_fn=CLUES_collate,
        batch_sampler=CLUES_BatchSampler(
            test_dataset, args.per_device_eval_batch_size, False, False,
            args.batch_per_epoch
        ))

    model = GuidedBertClues(
        bert, args.phase,
        args.encoder_feature_layer, args.guideline_feature_layer,
        args.encoder_feature_layer_num_head,
        args.guideline_feature_layer_num_head,
        args.num_feature_layer,
        args.model_version, args.mask_out_templates
    ).to(device)
    optimizer = {
        "AdamW": AdamW,
        "Adam": Adam,
        "SGD": SGD
    }[args.optimizer](model.parameters(), args.learning_rate,
                      weight_decay=args.weight_decay)
    warmup_steps = 0.1 * args.num_train_epochs * len(train_dataloader)
    scheduler = LambdaLR(
        optimizer, lr_lambda={
            "": lambda step: 1,
            "warmup": lambda step: min(
                (step+1) ** (-0.5), (step+1) * warmup_steps ** (-1.5)
            ) * warmup_steps ** 0.5
        }[args.scheduler]
    )
    last_model = args.output_dir+"/last_model.pt"
    start_epoch = 0
    if os.path.exists(last_model):
        print("loading from checkpoint:", last_model)
        ckpt = torch.load(last_model)
        model.load_state_dict(ckpt[0])
        optimizer.load_state_dict(ckpt[1])
        scheduler.load_state_dict(ckpt[2])
        start_epoch = ckpt[3] + 1

    best_dev_acc = -math.inf
    best_epoch = -1
    best_epoch_dev_loss = math.inf
    best_epoch_test_acc = -math.inf
    best_epoch_test_loss = math.inf
    # test_loss, test_acc = one_epoch(
    #     test_dataloader, model, optimizer, scheduler, tokenizer, device,
    #     "test", start_epoch-1, args.individual_scores)
    for i_epoch in tqdm(range(start_epoch, args.num_train_epochs)):
        print("="*64)
        one_epoch(train_dataloader, model, optimizer, scheduler, tokenizer,
                  device, "train", i_epoch, False)
        dev_loss, dev_acc = one_epoch(
            dev_dataloader, model, optimizer, scheduler, tokenizer, device,
            "dev", i_epoch, args.individual_scores)
        test_loss, test_acc = one_epoch(
            test_dataloader, model, optimizer, scheduler, tokenizer, device,
            "test", i_epoch, args.individual_scores)

        ckpt_file = f"{args.output_dir}/epoch{i_epoch}_model.pt"
        save_model(model, optimizer, scheduler, i_epoch, ckpt_file)
        save_model(model, optimizer, scheduler, i_epoch, last_model)
        if dev_acc > best_dev_acc:
            best_model_file = f"{args.output_dir}/best_model.pt"
            print("best dev accuracy")
            save_model(
                model, optimizer, scheduler, i_epoch, best_model_file)
            best_dev_acc = dev_acc
            best_epoch = i_epoch
            best_epoch_dev_loss = dev_loss
            best_epoch_test_acc = test_acc
            best_epoch_test_loss = test_loss

    if start_epoch != args.num_train_epochs:
        print("final results:")
        print(dict(zip(["best_epoch", "best_dev_acc", "best_epoch_dev_loss",
                        "best_epoch_test_acc", "best_epoch_test_loss"],
                       [best_epoch, best_dev_acc, best_epoch_dev_loss,
                        best_epoch_test_acc, best_epoch_test_loss])))
    else:
        test_loss, test_acc = one_epoch(
            test_dataloader, model, optimizer, scheduler, tokenizer, device,
            "test", start_epoch, args.individual_scores)

if __name__ == "__main__":
    args = parse_args()
    main(args)


from datetime import datetime, timedelta

from django.core.management.base import BaseCommand, CommandError
from django.conf import settings
from django.db.models import F

from tasks.models import Task

class Command(BaseCommand):

    CLAIMED_TIME_LIMIT = int(getattr(settings, "CLAIMED_TIME_LIMIT", 72))
    DONE_TIME_LIMIT = int(getattr(settings, "DONE_TIME_LIMIT", 72))

    def mark_claimed(self):

        time_offset = datetime.now() - timedelta(hours=self.CLAIMED_TIME_LIMIT)
        claimed_tasks = Task.objects.filter(state=Task.STATE_CLAIMED)
        claimed_tasks = claimed_tasks.filter(claimedTime__lte=time_offset)

        claimed_tasks.update(claimedTime=None, claimedBy=None, state=Task.STATE_NEW)

    def mark_done(self):
        time_offset = datetime.now() - timedelta(hours=self.DONE_TIME_LIMIT)
        done_tasks = Task.objects.filter(state=Task.STATE_DONE)
        done_tasks = done_tasks.filter(doneTime__lte=time_offset)
        done_tasks.update(verifiedTime=datetime.now(), verifiedBy=F('owner'), state=Task.STATE_VERIFIED)

    def handle(self, **options):
        if self.CLAIMED_TIME_LIMIT >= 0:
            self.mark_claimed()
        if self.DONE_TIME_LIMIT >= 0:
            self.mark_done()