import time
import pickle
import os
import sys as system_module_lol
from eddb import eddb_prime
from eddb.station import station_loader, Station
from eddb.system import system_loader, System
try:
    import matplotlib.pyplot as plt
except ModuleNotFoundError:
    import matplotlib
    matplotlib.use('agg')
    from matplotlib import pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import logging
from eddb.progress_tracker import generate_bar
from progressbar import UnknownLength
from eddb.pricelist import listing_loader, StationPriceList
from eddb.commodity import Commodities
logger = logging.getLogger('RareGraph')
logger.handlers = []
logger.addHandler(logging.StreamHandler(stream=system_module_lol.stdout))
logger.setLevel(logging.INFO)

commodity_controller = Commodities()

class RareNode:
    def __init__(self, system: System, station: Station, prices: StationPriceList, terminus=False, echo=False):
        # graph literally has no reason to be two-way. why? no idea
        self.terminus = terminus
        self.system = system
        self.station = station
        self.prices = prices
        if terminus and not echo:  # new root node automatically creates graph-end terminus as well
            end = RareNode(None, None, None, terminus=True, echo=True)
            self.link_x_forward = end
            self.link_x_backward = None
            self.link_y_forward = end
            self.link_y_backward = None
            self.link_z_forward = end
            self.link_z_backward = None
            end.link_x_backward = self
            end.link_y_backward = self
            end.link_z_backward = self
            end.link_x_forward = None
            end.link_y_forward = None
            end.link_z_forward = None
        if not terminus:
            self.commodities = []

    @property
    def is_terminus(self):
        return self.terminus

    def add_commodity(self, name):
        comm = commodity_controller.get_by_name(name)
        if comm is None:
            comm = {'id': None, 'name': name}
        self.commodities.append(comm)

    def next_x_gt(self, node):
        if self.link_x_forward.terminus:
            return True
        if self.link_x_forward.system.x >= node.system.x:
            return True
        return False

    def next_y_gt(self, node):
        if self.link_y_forward.terminus:
            return True
        if self.link_y_forward.system.y >= node.system.y:
            return True
        return False

    def next_z_gt(self, node):
        if self.link_z_forward.terminus:
            return True
        if self.link_z_forward.system.y >= node.system.y:
            return True
        return False

    def __str__(self):
        return '<terminus>' if not self.system else f'{self.system.name}:{self.station.name}'

    def __eq__(self, other):
        return self.system == other.system and self.station == other.station

    def insert_x(self, node):
        swap = self.link_x_forward
        self.link_x_forward = node
        node.link_x_forward = swap
        swap.link_x_backward = node
        node.link_x_backward = self

    def insert_y(self, node):
        swap = self.link_y_forward
        self.link_y_forward = node
        node.link_y_forward = swap
        swap.link_y_backward = node
        node.link_y_backward = self

    def insert_z(self, node):
        swap = self.link_z_forward
        self.link_z_forward = node
        node.link_z_forward = swap
        swap.link_z_backward = node
        node.link_z_backward = self

# TODO: remove system graph from here and do a separate package
class RareGraph:
    def __init__(self, nodes: list = []):

        logger.info('Graph engine initializing')
        logger.info('Sorting nodes')
        # sorting to speed up insertion
        # graph contains lesser X-Y-Z as top nodes; all nodes are expected to have 3 links (XYZ)
        nodes = sorted(nodes, key=lambda t_node: (t_node.system.x, t_node.system.y, t_node.system.z))
        self.index = dict()
        self.nodes = list()
        logger.info('Creating root node')
        self.root = RareNode(None, None, None, terminus=True)
        logger.info('Iterating through nodes')
        bar = generate_bar(nodes.__len__(), 'Inserting systems into graph')
        bar.start()
        for node in nodes:  # todo: don't recalculate graph on each insert may be?!
            self.insert(node)
            bar.update(bar.value + 1)
        bar.finish()

    def insert(self, node):
        # three passes to link the node
        logger.info(f'Parsing {node.system.name}:{node.station.name}')

        ref_x = self.root
        ref_y = self.root
        ref_z = self.root
        logger.info('Parse_X')
        bar = generate_bar(self.nodes.__len__(), 'Parse_X')
        bar.start()
        while True:
            if ref_x.next_x_gt(node):
                ref_x.insert_x(node)
                break
            else:
                ref_x = ref_x.link_x_forward
            bar.update(bar.value + 1)
        bar.finish()

        bar = generate_bar(self.nodes.__len__(), 'Parse_Y')
        bar.start()
        while True:
            if ref_y.next_y_gt(node):
                ref_y.insert_y(node)
                break
            else:
                ref_y = ref_y.link_y_forward
            bar.update(bar.value + 1)
        bar.finish()

        bar = generate_bar(self.nodes.__len__(), 'Parse_Z')
        bar.start()
        while True:
            if ref_z.next_z_gt(node):
                ref_z.insert_z(node)
                break
            else:
                ref_z = ref_z.link_z_forward
            bar.update(bar.value + 1)
        bar.finish()

        # keep a full list for future ref
        self.nodes.append(node)

        # index this shit up.
        self.index[node.__str__()] = node

    def plot_to(self, x, y, z):
        nodes = []
        ref = self.root
        for i in range(x):
            if ref.link_x_forward.terminus:
                break
            ref = ref.link_x_forward
            nodes.append(ref)
        ref = self.root
        for i in range(y):
            if ref.link_y_forward.terminus:
                break
            ref = ref.link_y_forward
            nodes.append(ref)
        ref = self.root
        for i in range(z):
            if ref.link_z_forward.terminus:
                break
            ref = ref.link_z_forward
            nodes.append(ref)
        self.plot(nodes)

    def plot(self, nodes=None):
        if nodes is None:
            nodes = self.nodes
        ax = plt.gca(projection="3d")
        x = list()
        y = list()
        z = list()
        for node in nodes:
            x.append(node.system.x)
            y.append(node.system.y)
            z.append(node.system.z)

        ax.scatter(x, y, z, c='r', s=100)
        ax.plot(x, y, z, color='r')
        plt.show()

    def query(self, system, station):
        logger.info(f'Querying {system}:{station}')
        for node in self.nodes:
            if node.system.name == system and node.station.name == station:
                return node
        return None

    def closest(self, system):
        distance = system.distance(self.nodes[0].system)
        ref = self.nodes[0]
        for node in self.nodes:
            chk = system.distance(node.system) < distance
            if chk < distance:
                ref = node
                distance = chk
        return ref

    def __generate(self, route, visited=[], sell_distance=150, max_ly=None):
        visited.append(route[-1]["node"])
        res = os.system('cls')
        if res != 0:
            os.system('clear')
        print([r["node"].system.name for r in route])
        # determine list of possible next nodes
        dup_possible_next_jump = [route[-1]["node"].link_x_forward, route[-1]["node"].link_x_backward,
                     route[-1]["node"].link_y_forward, route[-1]["node"].link_y_backward,
                     route[-1]["node"].link_z_forward, route[-1]["node"].link_z_backward]
        # clear terminus
        dup_possible_next_jump = [jump for jump in dup_possible_next_jump if not jump.terminus]
        # removing duplicates
        possible_next_jump = []
        for jump in dup_possible_next_jump:
            if jump not in possible_next_jump:
                possible_next_jump.append(jump)

        # filter visited which should be sold
        filtered_visited = []
        # bar = generate_bar(visited.__len__(), 'Removing sellable visited')
        # bar.start()
        for visit in visited:
            if visit.system.distance(route[-1]["node"].system) < sell_distance:
                filtered_visited.append(visit)
            # bar.update(bar.value + 1)
        # bar.finish()
        # filter by distance
        to_skip = []
        # bar = generate_bar(possible_next_jump.__len__(), 'Filtering visited by closing distance and max jump distance')
        # bar.start()
        for jump in possible_next_jump:
            for visit in visited:
                if visit.system.distance(route[-1]["node"].system) > visit.system.distance(jump.system):
                    to_skip.append(jump)
            if max_ly is not None:
                if route[-1]["node"].system.distance(jump.system) > max_ly:
                    to_skip.append(jump)

            # bar.update(bar.value + 1)
        possible_next_jump = [jump for jump in possible_next_jump if jump not in to_skip]
        # bar.finish()
        # if none are left - return false, we've hit a dead-end

        if possible_next_jump.__len__() == 0:
            print('    ' * route.__len__(), f'Terminated!')
            # return visited nodes
            return False

        # if any have been visited: loop is complete!
        # bar = generate_bar(possible_next_jump.__len__(), 'Checking loop possibility')
        # bar.start()
        for jump in possible_next_jump:
            if jump in [r["node"] for r in route]:
                profits = route[-1]["node"].prices - jump.prices
                route.append({"node": jump, "distance": route[-1]["node"].system.distance(jump.system),
                              "commodity": profits[0], "profit": profits[1], "updated": jump.prices.updated()})
                """({"node": jump["node"], "distance": jump["node"].system.distance(route[-1]["node"].system, ),
                  "commodity": jump["commodity"], "profit": jump["profit"],
                  "updated": jump["node"].prices.updated()})"""
                # todo: fix this profits returning either 0 or None bullshit =\\
                return True
            # bar.update(bar.value + 1)
        # bar.finish()

        # check profits:
        # bar = generate_bar(possible_next_jump.__len__(), 'Populating prices')
        # bar.start()
        jumps_and_profits = []
        for jump in possible_next_jump:
            comm_data = route[-1]["node"].prices - jump.prices
            jumps_and_profits.append({"node": jump, "profit": 0 if comm_data[1] is None else comm_data[1], "commodity": comm_data[0],
                                      "distance": jump.system.distance(route[-1]["node"].system)})
            # bar.update(bar.value + 1)
        jumps_and_profits = sorted(jumps_and_profits, key=lambda x: x["distance"])
        # bar.finish()
        # else take the best profit one and re-run GENERATE (doing this in a loop here
        # print(jumps_and_profits)
        ref_len = route.__len__()
        best_profit = 0
        best_profit_jump = jumps_and_profits[0]
        for jump in jumps_and_profits:
            route.append({"node": jump["node"], "distance": jump["distance"],
                          "commodity": jump["commodity"], "profit": jump["profit"],
                          "updated": jump["node"].prices.updated()})
            loop = self.__generate(route, visited=filtered_visited, sell_distance=sell_distance, max_ly=max_ly)
            # function has altered ROUTE
            if loop:  # if we have looped somewhere down the line, current route is cool
                return loop
            # if we do not loop, correct the route back:
            # know that we can't know how many nodes down the line the latest FALSE for a loop returned
            profit = 0  # check generated route cumulative profit
            for proposed_route_jump in route[ref_len:]:
                profit += proposed_route_jump["profit"]
            if profit > best_profit:
                best_profit_jump = jump
            while route.__len__() > ref_len:
                route.pop()
        # this loop on the first node will guarantee that if there's a loop it will get created.
        # if all failed to loop re-run for the best profit again and return
        # at this point, our route and visited are as we entered the algo

        jump = best_profit_jump
        route.append({"node": jump["node"], "distance": jump["node"].system.distance(route[-1]["node"].system),
                      "commodity": jump["commodity"], "profit": jump["profit"],
                      "updated": jump["node"].prices.updated()})
        # while we use the __generate function, it will not pop our current system or mess with the route.
        return self.__generate(route, filtered_visited, sell_distance=sell_distance, max_ly=max_ly)

    def node_by_commodity(self, cid=None, name=None):
        if not cid and not name:
            return None
        for node in self.nodes:
            if node.terminus:
                continue
            if cid is not None:
                if cid in [q['id'] for q in node.commodities]:
                    return node
            if name is not None:
                if name in [q['name'] for q in node.commodities]:
                    return node

    def generate_from(self, starting_node, current_system=None, first_jump_prices=None, sell_distance=150, max_ly=None):
        """
        format helper:
        [{"node": self.graph.nodes[0], "distance": 50, "commodity": "lulz", "profit": 400, "updated": 18},
         {"node": self.graph.nodes[1], "distance": 56, "commodity": "lulz2", "profit": 1400, "updated": 144}]"""
        route = [{"node": starting_node, "distance": 0, "commodity": None, "profit": None, "updated": starting_node.prices.updated()}]

        looped = self.__generate(route, sell_distance=sell_distance, max_ly=max_ly)  # visited systems are not returned because they shuffle about.

        if first_jump_prices is not None:
            best_deal = first_jump_prices - starting_node.prices
            route.insert(0, {"node": None, "distance": 0, "commodity": best_deal[0], "profit": best_deal[1], "updated": starting_node.prices.updated()})
            route[1]["distance"] = current_system.distance(route[1]["node"].system)

        for jump in route:
            jump['sell'] = []

        for jump in route:  # find where to sell what
            if jump["node"] is None:
                continue
            distance = 0
            ref = None
            for check in route:  #fixme if node does not exist (first-jump-thing)
                chk_dst = jump["node"].system.distance(check["node"].system)
                if chk_dst > distance:
                    distance = chk_dst
                    ref = check
            ref['sell'].append(jump['node'])
        if not looped:  # pretend we do have a loop
            route[0]['distance'] = route[-1]['node'].system.distance(route[0]['node'].system)
        return route

class RareLoader:
    def __init__(self, sell_distance=150):
        self.sell_distance = sell_distance

    def prime(self, ship_size, max_distance_from_star, filter_permit=True):
        logger.info('Opening file for RareLoader')
        f = open(f'{os.path.dirname(os.path.realpath(__file__))}/reference.tsv', 'r', encoding='utf-8')
        header = f.readline().split(';')
        reference = list()
        logger.info('Reading reference')
        for line in f.readlines():
            entry = dict()
            logger.debug(f'Processing {entry}')
            for e, k in zip(line.split(';'), header):
                logger.debug(f'{k} is {e}')
                entry[k.strip()] = e.strip()
            reference.append(entry)

        f.close()

        # load up systems to get id's
        # todo: may be do sort-merge here? although not really necessary on a small list
        logger.info('Loading systems')
        systems = system_loader(names=[q['System'] for q in reference], filter_needs_permit=filter_permit)

        # load up station infos
        station_filter = list()
        # pair up each reference with a system id
        logger.info('Creating station filter')
        for ref in reference:
            for sys in systems:
                if ref['System'] == sys.name:
                    station_filter.append((ref['Port'], sys.id))
                    break
        logger.info('Loading stations')
        stations = station_loader(refs=station_filter, landing_pad=ship_size, distance_to_star=max_distance_from_star)
        prices_list = listing_loader([q.sid for q in stations])
        nodes = []

        logger.info('Creating graph nodes')
        for ref in reference:
            system = None
            for sys in systems:
                if ref['System'] == sys.name:
                    system = sys
                    break
            if system is None:
                logger.warning(f'System {ref["System"]} not found, skipping')
                continue
                # raise NameError(f'System >:{ref["System"]}:< not found')
            station = None
            for st in stations:
                if ref['Port'] == st.name:
                    station = st
                    break
            if station is None:
                logger.warning(f'Station {ref["Port"]} not found')
                # raise NameError(f'Station {ref["Port"]} not found')
                continue
            dup = False
            for node in nodes:
                if node.system.id == system.id and node.station.sid == station.sid:
                    node.add_commodity(ref["Commodity"])
                    dup = True
                    break
            if dup:
                continue
            prices = None
            for listing in prices_list:
                if listing.station == station.sid:
                    prices = listing
            if prices is None:
                logger.warning(f'Skipping {ref["Commodity"]} because {station.name} has no listings')
                continue
            logger.debug(f'Created node for {ref["Commodity"]}')

            nodes.append(RareNode(system, station, prices))
            nodes[-1].add_commodity(ref['Commodity'])
        logger.info('Creating graph')
        self.graph = RareGraph(nodes=nodes)

        # create graph of Systems with stations attached
        # initialize current state if requested
        logger.info('Dumping backup data')
        pickle.dump(self.graph, open(f'{os.path.dirname(os.path.realpath(__file__))}/graph_dump.pcl', 'wb'))

    def init(self, ship_size='None', max_distance_from_star=-1, filter_needs_permit=False):
        self.ship_size = ship_size
        self.max_distance_from_star = max_distance_from_star

        logger.info('Rare loader initializing')
        logger.info('Trying to unpickle populated data')
        results = []
        for api in ['commodities.json', 'stations.jsonl', 'listings.csv', 'systems.csv']:
            results.append(eddb_prime.recache(api))
        if any(results):
            logger.info('Recache needed')
            self.prime(self.ship_size, self.max_distance_from_star, filter_needs_permit)
        else:
            try:
                self.graph = pickle.loads(
                    open(f'{os.path.dirname(os.path.realpath(__file__))}/graph_dump.pcl', 'rb+').read())
            except FileNotFoundError:
                logger.info('Failed to load up backup, reloading')
                self.prime(ship_size, max_distance_from_star, filter_needs_permit)

    def alternative_determine_next(self, cur_system, cur_station, limit=10, visited=[]):   # expecting current
        # cur_system = self.route[-1][0]
        ref_dict = {}
        node_dict = {}
        prices = None
        for node in self.graph.nodes:
            if node.terminus:
                continue
            if node.system == cur_system:
                prices = node.prices
                continue
            ref_dict[node.__str__()] = cur_system.distance(node.system)
            node_dict[node.__str__()] = node
        if prices is None:
            prices = StationPriceList(cur_station.sid)
            prices.populate()
        nodelist = sorted(ref_dict, key=lambda k: ref_dict[k])
        nodelist = [node_dict[q] for q in nodelist]
        to_skip = []
        for node in nodelist:
            for visit in visited:
                if cur_system.distance(visit.system) > node.system.distance(visit.system):  # distance closes!
                    to_skip.append(node)
        logger.debug(f'Filtering systems from {nodelist.__len__()}')
        nodelist = [node for node in nodelist if node not in to_skip]
        logger.debug(f'Resulting list is {nodelist.__len__()}')
        ret_obj = []
        for node in nodelist[:limit]:
            profit = prices - node.prices
            ret_obj.append({'system': node.system.name, 'station': node.station.name,
                            'commodity': [q['name'] for q in node.commodities],
                            'buy': '<none>' if profit[0] is None else profit[0].name,
                            'profit': profit[1], 'distance': ref_dict[node.__str__()],
                            'updated': (time.time() - node.system.updated_at) / 60 / 60})
        return ret_obj

    def query(self, system, station):
        return self.graph.query(system, station)

    def generate(self, system, station, node=None, max_ly=None):  # TODO: add limiters for something. max distance would be ok, but it's always "next over" so dunno. return trips are a failure though.
        if node is None:
            node = self.graph.closest(system)
            prices = StationPriceList(station)
            prices.populate()
        else:
            prices = None
        route = self.graph.generate_from(node, current_system=system, first_jump_prices=prices, sell_distance=self.sell_distance, max_ly=max_ly)
        return route

    def check_sell(self, current_system, visited):
        dst = current_system.distance(visited.system)
        if dst >= self.sell_distance:
            return dst
        else:
            return None

    def query_by_commodity(self, cid=None, name=None):
        return self.graph.node_by_commodity(cid=cid, name=name)

if __name__ == '__main__':
    chona = System(name='Chona')
    chona.populate()
    rl = RareLoader()
    rl.init('L', 50000, True)
    distance = 50000
    ref = None
    for node in rl.graph.nodes:
        if chona.distance(node.system) < distance:
            distance = chona.distance(node.system)
            ref = node.system
    print(ref.name, distance)


#! /usr/bin/python3
import pandas as pd
import matplotlib.pyplot as plot

df = pd.read_csv("avocado.csv")

#print(df.head(3))
#print(df["AveragePrice"].head())
albany_df = df[ df['region'] == 'Albany' ]

#print(albany_df.head())

albany_df.set_index('Date', inplace=True, drop=True)
print(albany_df.head())

plot.imshow(albany_df['AveragePrice'])
plot.show()


import cv2
import numpy as np
from utils import PriorBox, py_cpu_nms, decode, decode_landm

class RetinaDetection:
    def __init__(self, cfg, model_path="./checkpoints/faceDetector_RFB.onnx"):
        self.cfg = cfg
        self.model_path = model_path
        self.load(self.model_path)

    def load(self, model_path):
        self.model = cv2.dnn.readNetFromONNX(model_path)

    def forward(self, image):
        """image: should be BGR image"""
        img, scale, resize, im_height, im_width = self.preprocess(image)
        self.model.setInput(img)
        loc, conf, landms = self.model.forward(["output0", "532","531"])
        #post_process
        priorbox = PriorBox(self.cfg, image_size=(im_height, im_width))
        prior_data = priorbox.forward()
        boxes = decode(np.squeeze(loc,axis=0), prior_data, self.cfg['variance'])
        boxes = boxes * scale / resize
        scores = np.squeeze(conf, axis=0)[:, 1]
        landms = decode_landm(np.squeeze(landms, axis=0), prior_data, self.cfg['variance'])
        scale1 = np.array([img.shape[3], img.shape[2], img.shape[3], img.shape[2],
                                img.shape[3], img.shape[2], img.shape[3], img.shape[2],
                                img.shape[3], img.shape[2]])
        landms = landms * scale1 / resize
        # ignore low scores
        inds = np.where(scores > self.cfg["confidence_threshold"])[0]
        boxes = boxes[inds]
        landms = landms[inds]
        scores = scores[inds]

        # keep top-K before NMS
        order = scores.argsort()[::-1][:self.cfg["top_k"]]
        boxes = boxes[order]
        landms = landms[order]
        scores = scores[order]

        # do NMS
        dets = np.hstack((boxes, scores[:, np.newaxis])).astype(np.float32, copy=False)
        keep = py_cpu_nms(dets, self.cfg["nms_threshold"])
        # keep = nms(dets, args.nms_threshold,force_cpu=args.cpu)
        dets = dets[keep, :]
        landms = landms[keep]

        # keep top-K faster NMS
        dets = dets[:self.cfg["keep_top_k"], :]
        landms = landms[:self.cfg["keep_top_k"], :]

        dets = np.concatenate((dets, landms), axis=1)
        return dets

    def preprocess(self,img):
        target_size = self.cfg["long_side"]
        max_size = self.cfg["long_side"]
        im_shape = img.shape
        im_size_min = np.min(im_shape[0:2])
        im_size_max = np.max(im_shape[0:2])
        resize = float(target_size) / float(im_size_min)
        # prevent bigger axis from being more than max_size:
        if np.round(resize * im_size_max) > max_size:
            resize = float(max_size) / float(im_size_max)
        if self.cfg["origin_size"]:
            resize = 1
        if resize != 1:
            img = cv2.resize(img, None, None, fx=resize, fy=resize, interpolation=cv2.INTER_LINEAR)
        im_height, im_width, _ = img.shape
        scale = np.array([img.shape[1], img.shape[0], img.shape[1], img.shape[0]])
        img -= (104, 117, 123)
        img = img.transpose(2, 0, 1)
        img = np.expand_dims(img, axis=0)
        return img, scale, resize, im_height, im_width


import pandas as pd
import matplotlib.pyplot as plt
import lux
import cudf

def convert_to_list(x):
    """
    "a" --> ["a"]
    ["a","b"] --> ["a","b"]
    """
    if type(x) != list:
        return [x]
    else:
        return x

def pandas_to_lux(df):
    from lux.core.frame import LuxDataFrame

    values = df.values.tolist()
    ldf = LuxDataFrame(values, columns=df.columns)
    return ldf

def get_attrs_specs(intent):
    if intent is None:
        return []
    spec_obj = list(filter(lambda x: x.value == "", intent))
    return spec_obj

def get_filter_specs(intent):
    if intent is None:
        return []

    spec_obj = list(filter(lambda x: x.value != "", intent))
    return spec_obj

def check_import_lux_widget():
    import pkgutil

    if pkgutil.find_loader("luxwidget") is None:
        raise Exception(
            "luxwidget is not installed. Run `pip install luxwidget' to install the Jupyter widget.\nSee more at: https://github.com/lux-org/lux-widget"
        )

def get_agg_title(clause):
    attr = str(clause.attribute)
    if clause.aggregation is None:
        if len(attr) > 25:
            return attr[:15] + "..." + attr[-10:]
        return f"{attr}"
    elif attr == "Record":
        return f"Number of Records"
    else:
        if len(attr) > 15:
            return f"{clause._aggregation_name.capitalize()} of {attr[:15]}..."
        return f"{clause._aggregation_name.capitalize()} of {attr}"

def check_if_id_like(df, attribute):
    import re

    # Strong signals
    # so that aggregated reset_index fields don't get misclassified
    high_cardinality = df.cardinality[attribute] > 500
    attribute_contain_id = re.search(r"id|ID|iD|Id", str(attribute)) is not None
    almost_all_vals_unique = df.cardinality[attribute] >= 0.98 * len(df)
    is_string = isinstance(df[attribute],cudf.core.column.string.StringMethods) #pd.api.types.is_string_dtype(df[attribute])
    if is_string:
        # For string IDs, usually serial numbers or codes with alphanumerics have a consistent length (eg., CG-39405) with little deviation. For a high cardinality string field but not ID field (like Name or Brand), there is less uniformity across the string lengths.
        if len(df) > 50:
            #if lux.config.executor.name == "PandasExecutor":
            sampled = df[attribute].sample(50, random_state=99)
#             else:
#                 from lux.executor.SQLExecutor import SQLExecutor

#                 sampled = SQLExecutor.execute_preview(df, preview_size=50)
        else:
            sampled = df[attribute]
        str_length_uniformity = sampled.apply(lambda x: type(x) == str and len(x)).std() < 3
        return (
            high_cardinality
            and (attribute_contain_id or almost_all_vals_unique)
            and str_length_uniformity
        )
    else:
        if len(df) >= 2:
            series = df[attribute]
            diff = series.diff()
            # if diff.iloc[1]*(diff.shape[0]-1)==diff.iloc[1:].sum():
            #     print("even soacle")
            # diff = cudf.DataFrame(diff.dropna())
            # print("different here", diff)
            # col = diff.columns[0]
            # df_compare = diff.reset_index()[col]
            # new_df = diff.loc[1].repeat(diff.shape[0])
            # new_df= new_df.reset_index()
            #if new_df[1].equals( df_compare):
            new_diff = diff.to_pandas()
            if all(new_diff.iloc[1:] == new_diff.iloc[1]):
                evenly_spaced = True
            else:  evenly_spaced = False

        else:
            evenly_spaced = True
        if attribute_contain_id:
            almost_all_vals_unique = df.cardinality[attribute] >= 0.75 * len(df)
        return high_cardinality and (almost_all_vals_unique or evenly_spaced)

def check_if_id_like_for_sql(df, attribute):
    return df.cardinality[attribute] >= 0.98 * len(df)

def like_nan(val):
    if isinstance(val, str):
        return val.lower() == "nan"
    elif isinstance(val, float) or isinstance(val, int):
        import math

        return math.isnan(val)

def like_geo(val):
    return isinstance(val, str) and val.lower() in {"state", "country"}

def matplotlib_setup(w, h):
    plt.ioff()
    fig, ax = plt.subplots(figsize=(w, h))
    ax.set_axisbelow(True)
    ax.grid(color="#dddddd")
    ax.spines["right"].set_color("#dddddd")
    ax.spines["top"].set_color("#dddddd")
    return fig, ax

def is_numeric_nan_column(series):
    if series.dtype == object:
        if series.hasnans:
            series = series.dropna()
        try:
            return True, series.astype("float")
        except Exception as e:
            return False, series
    else:
        return False, series


"""
    程序功能： 游戏推荐展示
    作者：覃瀚辉
    时间：2021.12.15
"""

import wx
import csv

class MyFrame(wx.Frame):
    def __init__(self):
        super().__init__(parent=None, title='音乐推荐')
        self.data = self.loadSimpDat()
        self.com_game = self.recommend()
        panel = wx.Panel(self)
        panel.SetBackgroundColour([38,50,56])
        my_sizer = wx.BoxSizer(wx.VERTICAL)
        self.text_ctrl = wx.TextCtrl(panel,style=wx.TE_MULTILINE)
        font1 = wx.Font(15, wx.MODERN, wx.LIGHT, wx.NORMAL, False, 'Times New Roman')
        self.text_ctrl.SetFont(font1)
        self.text_ctrl.AppendText('Enter the games you have played')
        self.text_ctrl.SetBackgroundColour([255, 255, 229])
        my_sizer.Add(self.text_ctrl, 5, wx.ALL | wx.EXPAND, 5)
        my_btn = wx.Button(panel, label='推荐')
        my_btn.Bind(wx.EVT_BUTTON, self.on_press)
        my_sizer.Add(my_btn, 0, wx.ALL | wx.CENTER, 5)
        self.text_ctrl2 = wx.TextCtrl(panel,style=wx.TE_MULTILINE)
        self.text_ctrl2.SetFont(font1)
        self.text_ctrl2.AppendText('Recommended Games')
        self.text_ctrl2.SetBackgroundColour([255, 255, 229])
        my_sizer.Add(self.text_ctrl2, 5, wx.ALL | wx.EXPAND, 5)
        panel.SetSizer(my_sizer)
        self.Show()

    def loadSimpDat(self):
        """
        :return:返回原始数据
        """
        dataSet = []
        with open('outputdata/like_song.csv', 'r', encoding='utf-8') as csv_file:
            csv_reader = csv.reader(csv_file)
            for row in csv_reader:
                dataSet.append(row)
        return dataSet

    def recommend(self):
        """
        :return:可能玩的游戏
        """
        dataSet = []
        with open('outputdata/tuijian_song.csv', 'r', encoding='utf-8') as csv_file:
            csv_reader = csv.reader(csv_file)
            for row in csv_reader:
                dataSet.append(row)
        return dataSet

    def RecommendDault(self):
        """
        :return:默认推荐游戏(排名前20)
        """
        dataSet = []
        with open('outputdata/tuijian_default.csv', 'r', encoding='utf-8') as csv_file:
            csv_reader = csv.reader(csv_file)
            for row in csv_reader:
                dataSet.append(row)
        return dataSet

    def on_press(self, event):
        keyword_str = self.text_ctrl.Value
        keyword_str = keyword_str.strip()#除去首尾空格
        keyword_list = keyword_str.split('\n')
        index = 0
        tuijian_str = ''
        tuijian_list = []
        for row in self.data:
            # 如果全部符合 直接加入
            if keyword_list == row:
                for va in self.com_game[index]:
                    va = va.strip()
                    tuijian_list.append(va)
            # 部分符合也推荐
            for keyword in keyword_list:
                if (keyword in row) or (keyword == row):
                    for va in self.com_game[index]:
                        va = va.strip()
                        tuijian_list.append(va)
            index += 1
        tuijian_list = list(set(tuijian_list))
        # 没找到推荐的，就推荐默认的
        if len(tuijian_list) == 0:
            Dault_String = ''
            Dault_list = self.RecommendDault()
            index = 0
            for row in Dault_list:
                Dault_String += str(index + 1) + '. '
                Dault_String += str(row[0])
                Dault_String += '\n'
                index += 1
            self.text_ctrl2.SetValue(Dault_String)
        else:
            #找到推荐了 就推荐有的
            index = 0
            for va in tuijian_list:
                if va != '':
                    tuijian_str += str(index + 1) + '. '
                    tuijian_str += va
                    tuijian_str += '\n'
                    tuijian_str.strip()
                    index += 1
            self.text_ctrl2.SetValue(tuijian_str)
if __name__ == '__main__':
    app = wx.App()
    frame = MyFrame()
    app.MainLoop()


#!/usr/bin/env python3
# --------------------------------------------------
# Fábio Sousa.
# SAVI, December 2022.
# --------------------------------------------------

import matplotlib.pyplot as plt
import torch
from model import Model
from dataset import Dataset

def main():

    # -----------------------------------------------------------------
    # Initialization
    # -----------------------------------------------------------------

    #Create dataset
    dataset = Dataset(3000, -0.3, 14)

    loader = torch.utils.data.DataLoader(dataset=dataset, batch_size=256, shuffle=True)
    #It's good practice to guarantee that the batches have an appropriate size

    # #Draw training data
    # plt.plot(dataset.xs_np, dataset.ys_np_labels,'g.', label = 'labels')
    # plt.legend(loc='best')
    # plt.show()

    #Define hyper parameters
    device = 'cuda:0' if torch.cuda.is_available() else 'cpu' # cuda: 0 index of gpu

    model = Model() #Instantiate model
    model.to(device) #Move the model variable to the gpu if one exists

    learning_rate = 0.01
    maximum_num_epochs = 50
    criterion = torch.nn.MSELoss()
    optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)

    # -----------------------------------------------------------------
    # Training
    # -----------------------------------------------------------------

    #An epoch is an iteration of what comes next (of each prediction)
    idx_epoch = 0

    while True:

        #Train batch by batch
        for batch_idx, (xs_ten, ys_ten_labels) in enumerate(loader):

            #.to(device) keeps all variables in the same GPU or CPU
            xs_ten = xs_ten.to(device)
            ys_ten_labels = ys_ten_labels.to(device)

            #Apply the network to get the predicted ys
            ys_ten_predicted = model.forward(xs_ten)

            #Compute the error based on the predictions
            loss = criterion(ys_ten_predicted, ys_ten_labels)

            #Update the model, i.e. the neural network's weights
            optimizer.zero_grad() #Resets the weights to make sure we are not accumulating
            loss.backward() #Propagates the loss error into each neuron
            optimizer.step() #Update the weights

            #Report
            print('Epoch ' + str(idx_epoch) + ', Loss ' + str(loss.item()))

        idx_epoch += 1 #Go to next epoch

        #Termination criteria
        if idx_epoch > maximum_num_epochs:
            print('Finished training. Reached maximum number of epochs.')
            break

    # -----------------------------------------------------------------
    # Finalization
    # -----------------------------------------------------------------

    #Run the model once to get ys_predicted
    ys_ten_predicted = model.forward(dataset.xs_ten.to(device))
    ys_np_predicted = ys_ten_predicted.cpu().detach().numpy()

    #Draw data in a plot
    plt.plot(dataset.xs_np, dataset.ys_np_labels,'g.', label = 'Original')
    plt.plot(dataset.xs_np, ys_np_predicted,'rx', label = 'Predicted')
    plt.legend(loc='best')
    plt.grid('on')
    plt.title('Deep Learning for line fitting')
    plt.show()

if __name__ == "__main__":
    main()


"""
Fine-tuning the library models for causal language modeling (llama, opt, ...)
on a text file or a dataset without using HuggingFace Trainer.
"""

import logging
import math
import os
import shutil

import tensorflow.compat.v2 as tf
import torch
import transformers
from accelerate import Accelerator
from accelerate.logging import get_logger
from accelerate.utils import recursively_apply, set_seed
from peft import LoraConfig, inject_adapter_in_model, prepare_model_for_kbit_training
from torch.nn import functional as F
from tqdm.auto import tqdm
from transformers import (
    MODEL_MAPPING,
    AutoConfig,
    AutoModelForCausalLM,
    BitsAndBytesConfig,
    get_scheduler,
)
from utils import data_utils as dutils
from utils.args import parse_args

from neurocache import NeurocacheModelForCausalLM, OnDeviceCacheConfig
from neurocache.neurocache import get_attribute

# Make sure that tensorflow is not reserving GPUs.
tf.config.experimental.set_visible_devices([], "GPU")

logger = get_logger(__name__)

MODEL_CONFIG_CLASSES = list(MODEL_MAPPING.keys())

def init_neurocache_weights_(args, neurocache):
    _model = AutoModelForCausalLM.from_pretrained(
        args.model_name_or_path,
        config=AutoConfig.from_pretrained(
            args.model_name_or_path,
        ),
    )

    if hasattr(_model, "base_model"):
        _model = _model.model
    if hasattr(_model, "model"):
        _model = _model.model
    if hasattr(_model, "decoder"):
        _model = _model.decoder
    layers = get_attribute(_model, "layers")

    for i, idx in enumerate(neurocache.config.attention_layers):
        self_attn_module = get_attribute(layers[idx], "self_attn")
        neurocache.cache_attns[i]._init_proj_weights(self_attn_module)

    del _model

def prevent_full_backward_pass(_model):
    """
    Allow gradient checkpointing while preventing full backward pass on the model.
    """

    # Prevent full backward pass on the model. This is a workaround for the issue that
    # the full backward pass on the model will cause the memory and GPU usage to increase.

    # This is required because the default implementation of transformers does checkpointing
    # on all of the layers by default. This cause the backward pass to be done on all of the
    # layers, which is not necessary for our use case.

    # We fix this by enabling gradient checkpointing on the model, and then detaching the
    # output of the bottom layer. This will prevent the backward pass from propagating to
    # the bottom layer, and thus prevent the full backward pass on the model.

    # Note: this does not work by default due to PyTorch checkpoint function default use_reentrant
    # parameter to be True. We modify it to false in the transformers library.
    # modeling_utils.py:
    # After this line:
    # from torch.utils.checkpoint import checkpoint
    # Add this line:
    # checkpoint = functools.partial(checkpoint, use_reentrant=False)

    neurocache_config = _model.base_cache.config
    bottom_layer_idx = min(neurocache_config.attention_layers + neurocache_config.cache_layers)

    if hasattr(_model, "base_model"):
        _model = _model.model
    if hasattr(_model, "model"):
        _model = _model.model
    if hasattr(_model, "decoder"):
        _model = _model.decoder
    layers = get_attribute(_model, "layers")

    def forward_hook(module, args, output):
        return tuple(a.detach() if isinstance(a, torch.Tensor) else a for a in output)

    bottom_layer = layers[bottom_layer_idx - 1]
    bottom_layer.register_forward_hook(forward_hook)

def initialize_model(args, accelerator):
    if args.nf4:
        compute_dtype = getattr(torch, args.cache_dtype)
        quant_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_use_double_quant=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=compute_dtype,
            device_map="auto",
        )
    else:
        quant_config = None

    config = AutoConfig.from_pretrained(
        args.model_name_or_path,
    )
    config.use_cache = False
    config._flash_attn_2_enabled = True

    model = AutoModelForCausalLM.from_pretrained(
        args.model_name_or_path,
        config=config,
        quantization_config=quant_config,
        torch_dtype=getattr(torch, args.cache_dtype) if not args.nf4 else None,
    )

    # Prepare the model for kbit training.
    if args.nf4:
        model = prepare_model_for_kbit_training(model)

    # Set up the neurocache config and wrap the model with the neurocache model.
    if args.cache_layers is not None:
        cache_layers = [int(x) for x in args.cache_layers.split(",")]
    else:
        cache_layers = [model.config.num_hidden_layers * 3 // 4]

    if args.attention_layers is not None:
        attention_layers = [int(x) for x in args.attention_layers.split(",")]
    else:
        attention_layers = list(range(min(cache_layers), model.config.num_hidden_layers))

    if not args.disable_lora:
        # Apply LoRA to the main model to adapt it to using neurocache.
        lora_layers = attention_layers
        lora_config = LoraConfig(
            r=args.lora_r,
            lora_alpha=args.lora_alpha,
            target_modules=args.lora_modules.split(","),
            lora_dropout=args.lora_dropout,
            layers_to_transform=lora_layers,
            bias="none",
            task_type="CAUSAL_LM",
        )
        logger.info(f"LoRA Config: {lora_config}")
        model = inject_adapter_in_model(lora_config, model, "neurocache")
    else:
        # freeze the model
        for param in model.parameters():
            param.requires_grad = False

    if not args.disable_neurocache:
        if not args.pretrained_neurocache:
            neurocache_config = OnDeviceCacheConfig(
                attention_layers=attention_layers,
                cache_layers=cache_layers,
                cache_size=args.cache_size,
                cache_dtype=args.cache_dtype,
                compression_factor=args.compression_factor,
                neighborhood_size=args.neighborhood_size,
                context_size=args.context_size,
                topk=args.topk,
                cache_type=args.cache_type,
            )
            model = NeurocacheModelForCausalLM(model, neurocache_config)
            logger.info(f"Neurocache Config: {neurocache_config}")
            logger.info("Initializing neurocache weights from pretrained model.")
            init_neurocache_weights_(args, model.base_cache)
        else:
            logger.info(f"Loading neurocache from {args.pretrained_neurocache}")
            model = NeurocacheModelForCausalLM.from_pretrained(
                model, args.pretrained_neurocache, cache_type="ONDEVICE", is_training=True
            )
        model = model.to(accelerator.device)

    if args.disable_grad_checkpointing:
        model.base_model.gradient_checkpointing_disable()
    else:
        model.base_model.gradient_checkpointing_enable()
        model.base_model.enable_input_require_grads()
        prevent_full_backward_pass(model)

    print_trainable_parameters(model, accelerator)
    return model

def calculate_metrics(eval_loss):
    try:
        perplexity = math.exp(eval_loss)
    except OverflowError:
        perplexity = float("inf")
    bits_per_token = eval_loss * 1.442695
    return bits_per_token, perplexity

def print_trainable_parameters(model, accelerator):
    """
    Prints the number of trainable parameters in the model.
    """
    trainable_params = 0
    all_param = 0
    for _, param in model.named_parameters():
        all_param += param.numel()
        if param.requires_grad:
            accelerator.print(_, param.shape)
            trainable_params += param.numel()
    accelerator.print(
        f"trainable params: {trainable_params} || all params: {all_param} "
        f"|| trainable%: {100 * trainable_params / all_param}"
    )

def initialize_dataloader(batch_size, task_config, split, accelerator, resume_step=0):
    distributed_batch_size = batch_size * accelerator.num_processes
    per_device_batch_size = batch_size

    if accelerator.is_local_main_process:
        dataset = dutils.LongTextDataset(
            task_config, split, distributed_batch_size, skip=resume_step
        )
    else:
        dataset = dutils.PlaceholderDataset(
            task_config, split, distributed_batch_size, skip=resume_step
        )

    dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, pin_memory=True)

    def slice_fn(data, tensor_slice, process_index=None, num_processes=None):
        """split batch across processes."""
        nonlocal per_device_batch_size

        def _slice_tensor(tensor, tensor_slice):
            # remove the additional batch dim that is added by the dataloader
            return tensor.squeeze(0)[tensor_slice]

        tensor_slice = slice(
            process_index * per_device_batch_size, (process_index + 1) * per_device_batch_size
        )
        return recursively_apply(_slice_tensor, data, tensor_slice)

    # Prepare dataloaders
    dataloader = accelerator.prepare_data_loader(dataloader, slice_fn_for_dispatch=slice_fn)
    return dataloader

def run_evaluation(args, model, dataloader, accelerator, global_step=0, prefix="eval"):
    logger.info("***** Running evaluation *****")
    logger.info(f"    Steps: {global_step}")
    logger.info(f"    Prefix: {prefix}")

    if not args.disable_neurocache:
        if isinstance(model, NeurocacheModelForCausalLM):
            model.base_cache.reinitialize_cache()
        elif isinstance(model, torch.nn.parallel.DistributedDataParallel):
            model.module.base_cache.reinitialize_cache()

    model.eval()
    losses = []
    progress_bar = tqdm(
        range(args.max_eval_steps if args.max_eval_steps > 0 else None),
        disable=not accelerator.is_local_main_process,
    )

    for step, batch in enumerate(dataloader):
        with torch.no_grad():
            target_ids = batch.pop("labels")
            epoch = batch.pop("epochs")[0].item()

            if args.disable_neurocache:
                batch.pop("start_of_sequence")

            outputs = model(**batch)
            loss = F.cross_entropy(
                outputs.logits.view(-1, outputs.logits.size(-1)).to(torch.float32),
                target_ids.view(-1),
                ignore_index=0,  # TODO: get this id from tokenizer
            )

        losses.append(accelerator.gather_for_metrics(loss.repeat(args.per_device_eval_batch_size)))
        if step >= args.max_eval_steps and args.max_eval_steps > 0:
            break

        if epoch > 0:
            logger.warning("Epoch > 0 for evaluation, this is unexpected.")

        progress_bar.update(1)

    eval_loss = torch.cat(losses).mean().item()
    bits_per_token, perplexity = calculate_metrics(eval_loss)

    logger.info(
        f"global_step: {global_step}, perplexity: {perplexity}, loss: {eval_loss}, bits_per_token: {bits_per_token}"
    )

    if args.with_tracking:
        accelerator.log(
            {
                f"{prefix}_perplexity": perplexity,
                f"{prefix}_bits_per_token": bits_per_token,
                f"{prefix}_loss": eval_loss,
            },
            step=global_step,
        )

    model.train()
    if not args.disable_neurocache:
        if isinstance(model, NeurocacheModelForCausalLM):
            model.base_cache.reinitialize_cache()
        elif isinstance(model, torch.nn.parallel.DistributedDataParallel):
            model.module.base_cache.reinitialize_cache()
    return eval_loss, bits_per_token, perplexity

def main():
    args = parse_args()

    print(args)

    # Initialize the accelerator. We will let the accelerator handle device placement for us in this example.
    # If we're using tracking, we also need to initialize it here and it will by default pick up all supported trackers
    # in the environment
    accelerator_log_kwargs = {}

    if args.with_tracking:
        accelerator_log_kwargs["log_with"] = args.report_to
        accelerator_log_kwargs["project_dir"] = args.output_dir

    accelerator = Accelerator(
        gradient_accumulation_steps=args.gradient_accumulation_steps,
        dispatch_batches=True,
        split_batches=True,
        **accelerator_log_kwargs,
    )

    # Make one log on every process with the configuration for debugging.
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        level=logging.INFO,
    )
    logger.info(accelerator.state, main_process_only=False)
    if accelerator.is_local_main_process:
        transformers.utils.logging.set_verbosity_info()
    else:
        transformers.utils.logging.set_verbosity_error()

    # If passed along, set the training seed now.
    if args.seed is not None:
        set_seed(args.seed)

    # Handle the repository creation
    if accelerator.is_main_process:
        if args.output_dir is not None:
            os.makedirs(args.output_dir, exist_ok=True)
    accelerator.wait_for_everyone()

    # Task configuration
    task_config = dutils.TrainingTaskConfig.from_args(args)

    # Set up seqio tasks
    if args.tokenizer_name is not None:
        dutils.setup_tasks(args.tokenizer_name)
    else:
        dutils.setup_tasks(args.model_name_or_path)

    # Initialize the model
    model = initialize_model(args, accelerator)

    # Optimizer
    # Split weights in two groups, one with weight decay and the other not.
    no_decay = ["bias", "norm.weight"]
    optimizer_grouped_parameters = [
        {
            "params": [
                p
                for n, p in model.named_parameters()
                if (not any(nd in n for nd in no_decay) and p.requires_grad)
            ],
            "weight_decay": args.weight_decay,
        },
        {
            "params": [
                p
                for n, p in model.named_parameters()
                if (any(nd in n for nd in no_decay) and p.requires_grad)
            ],
            "weight_decay": 0.0,
        },
    ]
    optimizer = torch.optim.AdamW(optimizer_grouped_parameters, lr=args.learning_rate)

    # Scheduler and math around the number of training steps.
    lr_scheduler = get_scheduler(
        name=args.lr_scheduler_type,
        optimizer=optimizer,
        num_warmup_steps=args.num_warmup_steps,
        num_training_steps=args.max_train_steps,
    )

    # Initialize the datasets
    eval_dataloader = initialize_dataloader(
        args.per_device_eval_batch_size, task_config, task_config.eval_split, accelerator
    )

    # Prepare everything with our `accelerator`.
    model, optimizer, lr_scheduler = accelerator.prepare(model, optimizer, lr_scheduler)

    # We need to initialize the trackers we use, and also store our configuration.
    # The trackers initializes automatically on the main process.
    if args.with_tracking:
        experiment_config = vars(args)
        # TensorBoard cannot log Enums, need the raw value
        experiment_config["lr_scheduler_type"] = experiment_config["lr_scheduler_type"].value
        accelerator.init_trackers("clm_no_trainer", experiment_config)

    total_batch_size = (
        args.per_device_train_batch_size
        * accelerator.num_processes
        * args.gradient_accumulation_steps
    )

    logger.info("***** Running training *****")
    logger.info(f"  Instantaneous batch size per device = {args.per_device_train_batch_size}")
    logger.info(
        f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}"
    )
    logger.info(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
    logger.info(f"  Total optimization steps = {args.max_train_steps}")
    # Only show the progress bar once on each machine.
    progress_bar = tqdm(range(args.max_train_steps), disable=not accelerator.is_local_main_process)

    # Potentially load in the weights and states from a previous save
    completed_steps, resume_step = 0, 0
    if args.resume_from_checkpoint:
        checkpoint_path = args.resume_from_checkpoint.rstrip("/")

        accelerator.load_state(checkpoint_path)
        accelerator.print(f"Resumed from checkpoint: {checkpoint_path}")

        # need to multiply `gradient_accumulation_steps` to reflect real steps
        training_difference = os.path.splitext(os.path.basename(checkpoint_path))[0]
        resume_step = (
            int(training_difference.replace("step_", "")) * args.gradient_accumulation_steps
        )
        completed_steps = resume_step // args.gradient_accumulation_steps

    if args.save_and_quit:
        if accelerator.is_main_process:
            unwrapped_model = accelerator.unwrap_model(model)
            unwrapped_model.save_pretrained(args.output_dir)
        return

    # update the progress_bar if load from checkpoint
    if not args.only_evaluate:
        train_dataloader = initialize_dataloader(
            args.per_device_train_batch_size,
            task_config,
            task_config.train_split,
            accelerator,
            resume_step=resume_step,
        )

        progress_bar.update(completed_steps)

        model.train()
        logging_loss, logging_steps = 0, 0
        for step, batch in enumerate(train_dataloader):
            with accelerator.accumulate(model):
                target_ids = batch.pop("labels")
                epoch = batch.pop("epochs")[0].item()

                if args.disable_neurocache:
                    batch.pop("start_of_sequence")

                # Forward pass
                outputs = model(**batch)
                loss = F.cross_entropy(
                    outputs.logits.view(-1, outputs.logits.size(-1)).to(torch.float32),
                    target_ids.view(-1),
                    ignore_index=0,  # TODO: get this id from tokenizer
                )

                # Backward pass
                accelerator.backward(loss)

                # Clip gradients
                if accelerator.sync_gradients:
                    accelerator.clip_grad_norm_(model.parameters(), args.max_grad_norm)

                optimizer.step()
                lr_scheduler.step()
                optimizer.zero_grad()

                # Logging
                logging_loss += loss.detach().float()
                logging_steps += 1

            # Checks if the accelerator has performed
            # an optimization step behind the scenes
            if accelerator.sync_gradients:
                progress_bar.update(1)
                completed_steps += 1

                if completed_steps % args.logging_steps == 0:
                    bits_per_token, perplexity = calculate_metrics(logging_loss / logging_steps)
                    # Log metrics with formatting floats
                    progress_bar.set_description(
                        f"step {completed_steps:5d}, "
                        f"train_loss = {logging_loss / logging_steps:.4f}, "
                        f"train_perplexity = {perplexity:.4f}, "
                        f"learning_rate = {optimizer.param_groups[0]['lr']:.4e}"
                    )

                    if args.with_tracking:
                        accelerator.log(
                            {
                                "train_perplexity": perplexity,
                                "train_bits_per_token": bits_per_token,
                                "train_loss": logging_loss / logging_steps,
                                "train_epoch": epoch,
                                "train_lr": optimizer.param_groups[0]["lr"],
                            },
                            step=completed_steps,
                        )

                    logging_steps, logging_loss = 0, 0

                if completed_steps % args.evaluate_every_steps == 0:
                    run_evaluation(
                        args,
                        model,
                        eval_dataloader,
                        accelerator,
                        global_step=completed_steps,
                        prefix="validation",
                    )

                # Save model checkpoint
                if completed_steps % args.checkpointing_steps == 0:
                    output_dir = f"step_{completed_steps}"
                    output_dir = os.path.join(args.output_dir, output_dir)
                    accelerator.save_state(output_dir)

                    # keep only last n checkpoints
                    if args.keep_n_checkpoints is not None:
                        checkpoint_paths = [
                            os.path.join(args.output_dir, path)
                            for path in os.listdir(args.output_dir)
                            if path.startswith("step_")
                        ]
                        checkpoint_paths = sorted(
                            checkpoint_paths, key=lambda path: int(path.split("_")[-1])
                        )
                        for ckpt in checkpoint_paths[: -args.keep_n_checkpoints]:
                            logger.info(f"Removing old checkpoint: {ckpt}")
                            # remove checkpoint
                            shutil.rmtree(ckpt, ignore_errors=True)

                if completed_steps >= args.max_train_steps:
                    break

        if args.with_tracking:
            accelerator.end_training()

        # Save the model
        if args.output_dir is not None:
            accelerator.wait_for_everyone()
            unwrapped_model = accelerator.unwrap_model(model)
            if accelerator.is_main_process:
                if isinstance(unwrapped_model, NeurocacheModelForCausalLM):
                    unwrapped_model.save_pretrained(args.output_dir)
                elif not args.disable_lora:
                    from peft import get_peft_model_state_dict

                    state_dict = get_peft_model_state_dict(model, adapter_name="neurocache")
                    torch.save(state_dict, os.path.join(args.output_dir, "pytorch_model.bin"))
                else:
                    state_dict = unwrapped_model.state_dict()
                    torch.save(state_dict, os.path.join(args.output_dir, "pytorch_model.bin"))
    else:
        run_evaluation(
            args, model, eval_dataloader, accelerator, completed_steps, prefix="validation"
        )

    test_dataloader = initialize_dataloader(
        args.per_device_eval_batch_size, task_config, task_config.test_split, accelerator
    )
    run_evaluation(args, model, test_dataloader, accelerator, completed_steps, prefix="test")

if __name__ == "__main__":
    main()


import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.linear_model import Ridge
from sklearn.linear_model import RidgeCV
import matplotlib.pyplot as plt

data = pd.read_csv('./data/CCPP.csv')
data_x = data[['AT', 'V', 'AP', 'RH']]
data_y = data[['PE']]

x_train, x_test, y_train, y_test = train_test_split(data_x, data_y, random_state=1)

if __name__ == '__main__':
    # 岭回归
    model = Ridge(alpha=1)
    model.fit(x_train, y_train)
    print('theta0: ', model.intercept_)
    print('coefficients: ', model.coef_)

    # 超参数调优
    print()
    ridgecv = RidgeCV(alphas=[0.01, 0.1, 0.5, 1, 3, 5, 10, 30, 50, 100])
    ridgecv.fit(x_train, y_train)
    print(ridgecv.alpha_)

    # 超参数alpha与模型参数关系
    X = 1. / (np.arange(1, 11) + np.arange(0, 10)[:, np.newaxis])
    y = np.ones(10)
    alphas = np.logspace(-10, -2, 200)
    coefs = []
    for alpha in alphas:
        model = Ridge(alpha=alpha, fit_intercept=False)
        model.fit(X, y)
        coefs.append(model.coef_)
    ax = plt.gca()
    ax.plot(alphas, coefs)
    ax.set_xscale('log')
    ax.set_xlim(ax.get_xlim()[::-1])
    plt.xlabel('alpha')
    plt.ylabel('weight')
    plt.axis('tight')
    plt.show()


import torch
from torch import nn
from sklearn.metrics import mean_squared_error
import os
import sys
import networks
import gc

from vgg19 import vgg, VGG_loss
from networks import EFDM_loss

class InfusionNet(nn.Module):
    def __init__(self, config):
        super(InfusionNet, self).__init__()

        self.config = config
        self.device = torch.device('cuda:{}'.format(config.gpu) if torch.cuda.is_available() else 'cpu')

        self.lr = config.lr
        self.lambda_percept = config.lambda_percept
        self.lambda_const_style = config.lambda_const_style

        self.alpha_in = config.alpha_in
        self.alpha_out = config.alpha_out

        torch.cuda.empty_cache()

        # Encoder & Generator
        self.netE = networks.define_network(net_type='Encoder', alpha_in=self.alpha_in, alpha_out=self.alpha_out,
                                            init_type=self.config.init_type, gpu_ids=self.config.gpu, config = config)
        self.netS = networks.define_network(net_type='StyleEncoder', alpha_in=self.alpha_in, alpha_out=self.alpha_out,
                                            init_type=self.config.init_type, gpu_ids=self.config.gpu, config = config)
        self.netG = networks.define_network(net_type='Generator', alpha_in=self.alpha_in, alpha_out=self.alpha_out,
                                            init_type=self.config.init_type, gpu_ids=self.config.gpu, config = config)

        # Loss
        self.vgg_loss = VGG_loss(config, vgg)
        self.efdm_loss = EFDM_loss()

        # Optimizer
        self.optimizer_E = torch.optim.Adam(self.netE.parameters(), lr=self.config.lr, betas=(self.config.beta1, 0.99))
        self.optimizer_S = torch.optim.Adam(self.netS.parameters(), lr=self.config.lr, betas=(self.config.beta1, 0.99))
        self.optimizer_G = torch.optim.Adam(self.netG.parameters(), lr=self.config.lr, betas=(self.config.beta1, 0.99))

        # Scheduler
        self.E_scheduler = networks.get_scheduler(self.optimizer_E, config)
        self.S_scheduler = networks.get_scheduler(self.optimizer_S, config)
        self.G_scheduler = networks.get_scheduler(self.optimizer_G, config)

    def forward(self, data):
        self.real_A = data['real_A'].to(self.device)
        self.real_B = data['style_img'].to(self.device)

        # 1. Encode content & style information from images
        self.content_A, _ = self.netE(self.real_A)
        self.style_A, _ = self.netS(self.content_A)
        self.content_B, self.cont_feat_B = self.netE(self.real_B)
        self.style_B, self.sty_feat_B = self.netS(self.content_B)

        # 2. Generate infusioned images
        self.trs_AtoB, self.trs_AtoB_high, self.trs_AtoB_low, self.sty_kernel_preds = self.netG(self.content_A, self.style_B)

        # 3. Style description of infusioned images
        self.content_trs_AtoB, self.cont_feat_trs_AtoB = self.netE(self.trs_AtoB)
        self.style_trs_AtoB, self.sty_feat_trs_AtoB = self.netS(self.content_trs_AtoB)

        ## Predict kernel of infusioned image
        _, _, _, self.trs_kernel_preds = self.netG(self.content_A, self.style_trs_AtoB)

    """ Generator loss """
    def calc_G_loss(self):
        ## 1. Perceptual Loss using VGG
        self.G_percept, self.neg_idx = self.vgg_loss.perceptual_loss(self.real_A, self.real_B, self.trs_AtoB)
        self.G_percept *= self.lambda_percept

        ## 2. Style Contrastive Loss
        self.G_const_style = self.efdm_loss(self.cont_feat_B, self.sty_feat_B, self.cont_feat_trs_AtoB, self.sty_feat_trs_AtoB, self.neg_idx) * self.lambda_const_style

        self.G_loss = self.G_percept + self.G_const_style

    def train_step(self, data):
        self.set_requires_grad([self.netE, self.netS, self.netG], True)

        self.forward(data)
        self.calc_G_loss()

        self.optimizer_E.zero_grad()
        self.optimizer_S.zero_grad()
        self.optimizer_G.zero_grad()
        self.G_loss.backward()
        self.optimizer_E.step()
        self.optimizer_S.step()
        self.optimizer_G.step()

        train_dict = {}
        train_dict['G_loss'] = self.G_loss
        train_dict['G_Percept'] = self.G_percept
        train_dict['G_Const_Style'] = self.G_const_style

        train_dict['style_img'] = self.real_B
        train_dict['fake_AtoB'] = self.trs_AtoB
        train_dict['fake_AtoB_high'] = self.trs_AtoB_high
        train_dict['fake_AtoB_low'] = self.trs_AtoB_low

        return train_dict

    def val(self, data):
        with torch.no_grad():
            self.forward(data)

            val_dict = {}
            val_dict['real_A'] = self.real_A
            val_dict['fake_AtoB'] = self.trs_AtoB
            val_dict['real_B'] = self.real_B

        return val_dict

    def test(self, data):
        with torch.no_grad():
            real_A = data['real_A'].to(self.device)
            real_B = data['style_img'].to(self.device)

            # 1. Encode content & style information from images
            content_A, _ = self.netE(real_A)
            content_B, _ = self.netE(real_B)
            style_B, _ = self.netS(content_B)

            # 2. Generate infusioned images
            trs_AtoB, trs_AtoB_high, trs_AtoB_low, _ = self.netG(content_A, style_B)

            test_dict = {}
            test_dict['style_img'] = real_B
            test_dict['fake_AtoB'] = trs_AtoB
            test_dict['fake_AtoB_high'] = trs_AtoB_high
            test_dict['fake_AtoB_low'] = trs_AtoB_low

        return test_dict

    def style_mixing(self, data):
        with torch.no_grad():
            # B1: high style image / B2: low style image
            real_A = data['real_A'].to(self.device)
            real_B1_h = data['style_high'].to(self.device)
            real_B2_l = data['style_low'].to(self.device)

            # 1. Encode content & style information from images
            content_A = self.netE(real_A)
            content_B1_h = self.netE(real_B1_h)
            content_B2_l = self.netE(real_B2_l)
            style_B1_h, _ = self.netS(content_B1_h)[0]
            _, style_B2_l = self.netS(content_B2_l)[0]
            style_B = style_B1_h, style_B2_l

            # 2. Generate infusioned images
            trs_AtoB, trs_AtoB_high, trs_AtoB_low, _ = self.netG(content_A, style_B)

            mix_dict = {}
            mix_dict['style_high_img'] = real_B1_h
            mix_dict['style_low_img'] = real_B2_l
            mix_dict['fake_AtoB'] = trs_AtoB
            mix_dict['fake_AtoB_high'] = trs_AtoB_high
            mix_dict['fake_AtoB_low'] = trs_AtoB_low

        return mix_dict

    def set_requires_grad(self, nets, requires_grad=False):
        if not isinstance(nets, list):
            nets = [nets]
        for net in nets:
            if net is not None:
                for param in net.parameters():
                    param.requires_grad = requires_grad


import torch
import torch.nn as nn
# from utils import load_state_dict_from_url
from meta_module import *

DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

model_urls = {
    'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
    'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
    'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
    'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
    'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
    'resnext50_32x4d': 'https://download.pytorch.org/models/resnext50_32x4d-7cdf4587.pth',
    'resnext101_32x8d': 'https://download.pytorch.org/models/resnext101_32x8d-8ba56ff5.pth',
    'wide_resnet50_2': 'https://download.pytorch.org/models/wide_resnet50_2-95faca4d.pth',
    'wide_resnet101_2': 'https://download.pytorch.org/models/wide_resnet101_2-32ee1156.pth',
}

class BasicBlock(MetaModule):
    expansion = 1

    def __init__(self, in_planes, planes, stride=1):
        super(BasicBlock, self).__init__()
        self.conv1 = MetaConv2d(
            in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
        self.bn1 = MetaBatchNorm2d(planes)
        self.conv2 = MetaConv2d(planes, planes, kernel_size=3,
                               stride=1, padding=1, bias=False)
        self.bn2 = MetaBatchNorm2d(planes)

        self.shortcut = MetaSequential()
        if stride != 1 or in_planes != self.expansion*planes:
            self.shortcut = MetaSequential(
                MetaConv2d(in_planes, self.expansion*planes,
                          kernel_size=1, stride=stride, bias=False),
                MetaBatchNorm2d(self.expansion*planes)
            )

    def forward(self, x):
        out = F.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        out += self.shortcut(x)
        out = F.relu(out)
        return out

class Bottleneck(MetaModule):
    expansion = 4

    def __init__(self, in_planes, planes, stride=1):
        super(Bottleneck, self).__init__()
        self.conv1 = MetaConv2d(in_planes, planes, kernel_size=1, bias=False)
        self.bn1 = MetaBatchNorm2d(planes)
        self.conv2 = MetaConv2d(planes, planes, kernel_size=3,
                               stride=stride, padding=1, bias=False)
        self.bn2 = MetaBatchNorm2d(planes)
        self.conv3 = MetaConv2d(planes, self.expansion *
                               planes, kernel_size=1, bias=False)
        self.bn3 = MetaBatchNorm2d(self.expansion*planes)

        self.shortcut = MetaSequential()
        if stride != 1 or in_planes != self.expansion*planes:
            self.shortcut = MetaSequential(
                MetaConv2d(in_planes, self.expansion*planes,
                          kernel_size=1, stride=stride, bias=False),
                MetaBatchNorm2d(self.expansion*planes)
            )

    def forward(self, x):
        out = F.relu(self.bn1(self.conv1(x)))
        out = F.relu(self.bn2(self.conv2(out)))
        out = self.bn3(self.conv3(out))
        out += self.shortcut(x)
        out = F.relu(out)
        return out

class ResNet(MetaModule):
    def __init__(self, block, num_blocks, num_classes=10, **kwargs):
        super(ResNet, self).__init__()
        self.in_planes = 64

        # print('happy!!!')
        # raise

        self.conv1 = MetaConv2d(3, 64, kernel_size=3,
                               stride=1, padding=1, bias=False)
        self.bn1 = MetaBatchNorm2d(64)
        self.layer1 = self._make_layer(block, 64, num_blocks[0], stride=1)
        self.layer2 = self._make_layer(block, 128, num_blocks[1], stride=2)
        self.layer3 = self._make_layer(block, 256, num_blocks[2], stride=2)
        self.layer4 = self._make_layer(block, 512, num_blocks[3], stride=2)
        self.linear = MetaLinear(512*block.expansion, num_classes)

    def _make_layer(self, block, planes, num_blocks, stride):
        strides = [stride] + [1]*(num_blocks-1)
        layers = []
        for stride in strides:
            layers.append(block(self.in_planes, planes, stride))
            self.in_planes = planes * block.expansion
        return MetaSequential(*layers)

    def _forward_impl(self, x):
        out = F.relu(self.bn1(self.conv1(x)))
        out = self.layer1(out)
        out = self.layer2(out)
        out = self.layer3(out)
        out = self.layer4(out)
        out = F.avg_pool2d(out, 4) # (input, kernal_size)
        out = out.view(out.size(0), -1)
        out = self.linear(out)

        return out

    def forward(self, x):
        return self._forward_impl(x)

class cResNet(ResNet):
    def __init__(self, arch, block, layers, **kwargs):
        super(cResNet,self).__init__( block, layers, **kwargs)
        self.name = arch
        self.loss = nn.NLLLoss()
        # self.
        # DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    def forward(self,loss):

        inp, out = loss.sample()
        # inp = Variable(inp.view(inp.size()[0], 32*32*3)).to(DEVICE)
        # inp = inp.to(self.DEVICE)
        # out = Variable(out).to(self.DEVICE)

        # cur_layer = 0
        # while f'mat_{cur_layer}' in self.layers:
        #     inp = self.activation(self.layers[f'mat_{cur_layer}'](inp))
        #     cur_layer += 1

        pred = F.log_softmax(self._forward_impl(inp), dim=1)
        l = self.loss(pred, out)
        return l

    def cal_test_acc(self, testloader):

        correct = 0
        total = 0
        with torch.no_grad():
            for i,data in enumerate(testloader):

                images, labels = data
                images, labels = images.to(DEVICE), labels.to(DEVICE)

                outputs = self._forward_impl(images)

                _, predicted = torch.max(outputs.data, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()

        acc = 100 * correct / total

        print(f'\ntest acc is: {acc}\n')
        return acc

def _resnet(arch, block, layers, pretrained, progress, **kwargs):
    # model = ResNet(block, layers, **kwargs)
    model = cResNet(arch, block, layers, **kwargs).to(DEVICE)  #-wz
    if pretrained:
        state_dict = load_state_dict_from_url(model_urls[arch],
                                              progress=progress)
        model.load_state_dict(state_dict)
    return model

def resnet18(pretrained=False, progress=True, **kwargs):
    r"""ResNet-18 model from
    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_

    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
        progress (bool): If True, displays a progress bar of the download to stderr
    """
    return _resnet('resnet18', BasicBlock, [2, 2, 2, 2], pretrained, progress,
                   **kwargs)

def resnet34(pretrained=False, progress=True, **kwargs):
    r"""ResNet-34 model from
    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_

    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
        progress (bool): If True, displays a progress bar of the download to stderr
    """
    return _resnet('resnet34', BasicBlock, [3, 4, 6, 3], pretrained, progress,
                   **kwargs)

def resnet50(pretrained=False, progress=True, **kwargs):
    r"""ResNet-50 model from
    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_

    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
        progress (bool): If True, displays a progress bar of the download to stderr
    """
    return _resnet('resnet50', Bottleneck, [3, 4, 6, 3], pretrained, progress,
                   **kwargs)

def resnet101(pretrained=False, progress=True, **kwargs):
    r"""ResNet-101 model from
    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_

    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
        progress (bool): If True, displays a progress bar of the download to stderr
    """
    return _resnet('resnet101', Bottleneck, [3, 4, 23, 3], pretrained, progress,
                   **kwargs)

def resnet152(pretrained=False, progress=True, **kwargs):
    r"""ResNet-152 model from
    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_

    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
        progress (bool): If True, displays a progress bar of the download to stderr
    """
    return _resnet('resnet152', Bottleneck, [3, 8, 36, 3], pretrained, progress,
                   **kwargs)

def resnext50_32x4d(pretrained=False, progress=True, **kwargs):
    r"""ResNeXt-50 32x4d model from
    `"Aggregated Residual Transformation for Deep Neural Networks" <https://arxiv.org/pdf/1611.05431.pdf>`_

    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
        progress (bool): If True, displays a progress bar of the download to stderr
    """
    kwargs['groups'] = 32
    kwargs['width_per_group'] = 4
    return _resnet('resnext50_32x4d', Bottleneck, [3, 4, 6, 3],
                   pretrained, progress, **kwargs)

def resnext101_32x8d(pretrained=False, progress=True, **kwargs):
    r"""ResNeXt-101 32x8d model from
    `"Aggregated Residual Transformation for Deep Neural Networks" <https://arxiv.org/pdf/1611.05431.pdf>`_

    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
        progress (bool): If True, displays a progress bar of the download to stderr
    """
    kwargs['groups'] = 32
    kwargs['width_per_group'] = 8
    return _resnet('resnext101_32x8d', Bottleneck, [3, 4, 23, 3],
                   pretrained, progress, **kwargs)

def wide_resnet50_2(pretrained=False, progress=True, **kwargs):
    r"""Wide ResNet-50-2 model from
    `"Wide Residual Networks" <https://arxiv.org/pdf/1605.07146.pdf>`_

    The model is the same as ResNet except for the bottleneck number of channels
    which is twice larger in every block. The number of channels in outer 1x1
    convolutions is the same, e.g. last block in ResNet-50 has 2048-512-2048
    channels, and in Wide ResNet-50-2 has 2048-1024-2048.

    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
        progress (bool): If True, displays a progress bar of the download to stderr
    """
    kwargs['width_per_group'] = 64 * 2
    return _resnet('wide_resnet50_2', Bottleneck, [3, 4, 6, 3],
                   pretrained, progress, **kwargs)

def wide_resnet101_2(pretrained=False, progress=True, **kwargs):
    r"""Wide ResNet-101-2 model from
    `"Wide Residual Networks" <https://arxiv.org/pdf/1605.07146.pdf>`_

    The model is the same as ResNet except for the bottleneck number of channels
    which is twice larger in every block. The number of channels in outer 1x1
    convolutions is the same, e.g. last block in ResNet-50 has 2048-512-2048
    channels, and in Wide ResNet-50-2 has 2048-1024-2048.

    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
        progress (bool): If True, displays a progress bar of the download to stderr
    """
    kwargs['width_per_group'] = 64 * 2
    return _resnet('wide_resnet101_2', Bottleneck, [3, 4, 23, 3],
                   pretrained, progress, **kwargs)

# Below two functions are early testing functions by Wenqing, deprecated
# def resnet18_conv1_3x3(pretrained=False, progress=True, **kwargs):
#     r"""ResNet-18 model from
#     `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_

#     Args:
#         pretrained (bool): If True, returns a model pre-trained on ImageNet
#         progress (bool): If True, displays a progress bar of the download to stderr
#     """
#     kwargs['kernel_size_first_layer'] = 3
#     return _resnet('resnet18', BasicBlock, [2, 2, 2, 2], pretrained, progress,
#                    **kwargs)
# def resnet34_conv1_3x3(pretrained=False, progress=True, **kwargs):
#     r"""ResNet-34 model from
#     `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_

#     Args:
#         pretrained (bool): If True, returns a model pre-trained on ImageNet
#         progress (bool): If True, displays a progress bar of the download to stderr
#     """
#     kwargs['kernel_size_first_layer'] = 3
#     return _resnet('resnet34', BasicBlock, [3, 4, 6, 3], pretrained, progress,
#                    **kwargs)


# Python program killing
# a thread using multiprocessing
# module

"""
Though the interface of the two modules is similar, the two modules have very different implementations.
All the threads share global variables, whereas processes are completely separate from each other.
Hence, killing processes is much safer as compared to killing threads. The Process class is provided a method,
terminate(), to kill a process. Now, getting back to the initial problem. Suppose in the above code,
we want to kill all the processes after 0.03s have passed.
This functionality is achieved using the multiprocessing module in the following code.
"""
import multiprocessing
import time

def func(number):
    for i in range(1, 10):
        time.sleep(0.01)
        print('Processing ' + str(number) + ': prints ' + str(number * i))

    # list of all processes, so that they can be killed afterwards

all_processes = []

for i in range(0, 3):
    process = multiprocessing.Process(target=func, args=(i,))
    process.start()
    all_processes.append(process)

# kill all processes after 0.03s
time.sleep(0.03)
for process in all_processes:
    process.terminate()


from .service import Service, current_user_id, current_client_id
from models.goal_tree import GoalTree, GoalTreeNode
from models.note_category import NoteCategory
from utility import isTrue, isFalse, isNone, notNone

import datetime

class GoalService(Service):
    def __init__(self, session):
        super().__init__(session)

    def newTree(self, name):
        gtree = GoalTree(
            owner=current_client_id(),
            name=name,
            created_at=datetime.datetime.utcnow()
        )

        self._session.add(gtree)
        self._session.flush()

        beginning = self.newTreeNode(gtree.id, None, "Beginning")

        self._session.commit()
        return gtree

    def newTreeNode(self, tree, parent, title):
        node = GoalTreeNode(
            title=title,
            tree=tree,
            parent_node=parent,
            created_at=datetime.datetime.utcnow()
        )

        self._session.add(node)
        self._session.commit()
        return node

    def editTreeNode(self, tree, id, title, description):
        node = self.goalTreeNode(id)

        if node.tree != tree:
            raise RuntimeError("Node doesn't belong to tree")

        node.title = title
        node.description = description

        return node

    def goalTreeNode(self, id):
        return self._query(GoalTreeNode) \
                   .filter(GoalTreeNode.id == id) \
                   .first()

    def goalTreeList(self):
        return self._query(GoalTree) \
                   .filter(GoalTree.owner == current_client_id()) \
                   .filter(isNone(GoalTree.deleted_at)) \
                   .order_by(GoalTree.created_at.desc()) \
                   .all()

    def goalTreeFull(self, id):
        gtree = self._query(GoalTree) \
                   .filter(GoalTree.owner == current_client_id()) \
                   .filter(GoalTree.id == id) \
                   .first()
        nodes = self._query(GoalTreeNode) \
                   .filter(GoalTreeNode.tree == id) \
                   .filter(isNone(GoalTreeNode.deleted_at)) \
                   .order_by(GoalTreeNode.created_at.desc()) \
                   .all()
        gtree.nodes = nodes
        return gtree

    def moveNode(self, movedNodeId, targetNodeId):
        movedNode = self.goalTreeNode(movedNodeId)
        targetNode = self.goalTreeNode(targetNodeId)

        if movedNode.tree != targetNode.tree:
            raise RuntimeError("Trying to move to a different tree")

        parent = targetNode
        while parent.parent_node is not None:
            if parent.parent_node == movedNode.id:
                parent.parent_node = movedNode.parent_node
                break
            parent = self.goalTreeNode(parent.parent_node)

        movedNode.parent_node = targetNode.id

        return self.goalTreeFull(movedNode.tree)

    def deleteNode(self, treeId, nodeId):
        node = self.goalTreeNode(nodeId)

        if node.parent_node is None:
            raise RuntimeError("Trying to delete a root node")

        if node.tree != treeId:
            raise RuntimeError("Node doesn't belong to tree")

        node.deleted_at = datetime.datetime.utcnow()

        nodes = self._query(GoalTreeNode) \
                   .filter(GoalTreeNode.tree == treeId) \
                   .filter(isNone(GoalTreeNode.deleted_at)) \
                   .order_by(GoalTreeNode.created_at.desc()) \
                   .all()

        for n in nodes:
            print(n.deleted_at)
            if n.parent_node == node.id:
                n.parent_node = node.parent_node

        return self.goalTreeFull(treeId);


import os
import pickle
import datetime
from ...common import exceptions as ex
from os import getcwd

FILE_EXTENSION = '.pickle'
STORAGE_FOLDER_PATH = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))

BACKUP_STORAGE_PATH = os.path.join(STORAGE_FOLDER_PATH, 'data_storage', 'for_backup_storage')
TPL_STORAGE_PATH = os.path.join(STORAGE_FOLDER_PATH, 'data_storage', 'for_tpl_storage')
CONFIG_STORAGE_PATH = os.path.join(STORAGE_FOLDER_PATH, 'data_storage', 'for_config_storage')

def load_backup(user_id, tool_id, project_id, backup_name):
    file_path = __get_backup_file_path(user_id, tool_id, project_id,
                                       backup_name)
    saved_content = __load(file_path)
    return saved_content['data']

def save_backup(user_id, tool_id, project_id, data_to_save, backup_name):
    folder = __get_user_backup_dir_path(user_id, tool_id, project_id)
    if not os.path.exists(folder):
        os.makedirs(folder)

    file_path = __get_backup_file_path(user_id, tool_id, project_id,
                                       backup_name)
    file_info = {
        'name': backup_name,
        'user_id': user_id,
        'date': datetime.datetime.now()
    }
    __save(file_path, file_info, data_to_save)

def delete_backup(user_id, tool_id, project_id, backup_name):
    file_path = __get_backup_file_path(user_id, tool_id, project_id,
                                       backup_name)
    __delete(file_path)

def get_available_backups_info(users_ids, tool_id, project_id):
    files_list = []
    if not users_ids or not tool_id:
        return files_list
    ext_len = len(FILE_EXTENSION)
    for user_id in users_ids:
        user_backup_dir_path = __get_user_backup_dir_path(user_id, tool_id,
                                                          project_id)
        for entry in os.scandir(user_backup_dir_path):
            if not entry.name.startswith('.') and entry.is_file():
                no_ext_name = entry.name[:-ext_len]
                saved_content = load_backup(user_id, tool_id, project_id,
                                            no_ext_name)
                if saved_content:
                    files_list.append(saved_content['info'])
    return files_list

def load_template(tool_id, tpl_name):
    file_path = __get_tpl_file_path(tool_id, tpl_name)
    saved_content = __load(file_path)
    return saved_content

def save_template(tool_id, data_to_save, tpl_name='default_tpl'):
    if not os.path.exists(__get_tpl_dir_path(tool_id)):
        os.makedirs(__get_tpl_dir_path(tool_id))
    file_path = __get_tpl_file_path(tool_id, tpl_name)
    file_info = {
        'name': tpl_name,
        'date': datetime.datetime.now()
    }
    __save(file_path, file_info, data_to_save)

def delete_template(tool_id, tpl_name):
    file_path = __get_tpl_file_path(tool_id, tpl_name)
    __delete(file_path)

def __save(file_path, file_info, data_to_save):

    data_with_info = {
        'data': data_to_save,
        'info': file_info
    }
    with open(file_path, 'wb') as f:
        pickle.dump(data_with_info, f, pickle.HIGHEST_PROTOCOL)

def __load(file_path):
    saved_content = {}
    if os.path.isfile(file_path):
        with open(file_path, 'rb') as f:
            saved_content = pickle.load(f)
    return saved_content

def __delete(file_path):
    if not os.path.isfile(file_path):
        raise ex.NotExistingFileError(file_path)
    os.remove(file_path)

def __get_backup_file_path(user_id, tool_id, project_id, backup_name):
    file_name = backup_name + FILE_EXTENSION
    file_path = os.path.join(BACKUP_STORAGE_PATH, str(user_id),
                             str(tool_id) + str(project_id), file_name)
    return file_path

def __get_user_backup_dir_path(user_id, tool_id, project_id):
    user_backup_dir_path = os.path.join(BACKUP_STORAGE_PATH,
                                        str(user_id),
                                        str(tool_id) + str(project_id))
    return user_backup_dir_path

def __get_tpl_file_path(tool_id, tpl_name):
    file_name = tpl_name + FILE_EXTENSION
    file_path = os.path.join(TPL_STORAGE_PATH, str(tool_id), file_name)
    return file_path

def __get_tpl_dir_path(tool_id):
    tpl_dir_path = os.path.join(TPL_STORAGE_PATH, str(tool_id))
    return tpl_dir_path


import os
import zipfile

class Compressor:
    def __init__(self, path):
        self.path = path

    def compress_all_files(self, zipname):
        zip = zipfile.ZipFile(zipname, "w", zipfile.ZIP_DEFLATED)
        for dirname, subdirs, files in os.walk(self.path):
            zip.write(dirname)
            for filename in files:
                zip.write(os.path.join(dirname, filename))
        zip.close()

    def decompress_file(self, zipname):
        zip = zipfile.ZipFile(zipname, "r")
        zip.extractall()
        zip.close()


import os
import sys
import argparse

import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader, Dataset
from torch.utils.data.sampler import Sampler

import csv
import random
import numpy as np
import pandas as pd

from PIL import Image
filenameToPILImage = lambda x: Image.open(x)

# fix random seeds for reproducibility
SEED = 123
torch.manual_seed(SEED)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
random.seed(SEED)
np.random.seed(SEED)

def worker_init_fn(worker_id):
    np.random.seed(np.random.get_state()[1][0] + worker_id)

# mini-Imagenet dataset
class MiniDataset(Dataset):
    def __init__(self, csv_path, data_dir):
        self.data_dir = data_dir
        self.data_df = pd.read_csv(csv_path).set_index("id")

        self.transform = transforms.Compose([
            filenameToPILImage,
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
            ])

    def __getitem__(self, index):
        path = self.data_df.loc[index, "filename"]
        label = self.data_df.loc[index, "label"]
        image = self.transform(os.path.join(self.data_dir, path))
        return image, label

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

class GeneratorSampler(Sampler):
    def __init__(self, episode_file_path):
        episode_df = pd.read_csv(episode_file_path).set_index("episode_id")
        self.sampled_sequence = episode_df.values.flatten().tolist()

    def __iter__(self):
        return iter(self.sampled_sequence)

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

def conv_block(in_channels, out_channels):
    bn = nn.BatchNorm2d(out_channels)
    return nn.Sequential(
        nn.Conv2d(in_channels, out_channels, 3, padding=1),
        bn,
        nn.ReLU(),
        nn.MaxPool2d(2)
    )
class ProNet(nn.Module):
    def __init__(self, in_channels=3, hid_channels=64, out_channels=64):
        super(ProNet, self).__init__()
        self.encoder = nn.Sequential(
            conv_block(in_channels, hid_channels),
            conv_block(hid_channels, hid_channels),
            conv_block(hid_channels, hid_channels),
            conv_block(hid_channels, out_channels)
        )

    def forward(self, input_data):
        x = self.encoder(input_data)
        return x.view(x.size(0), -1)

def euclidean_dist(x, y):
    n = x.shape[0]
    m = y.shape[0]
    a = x.unsqueeze(1).expand(n, m, -1)
    b = y.unsqueeze(0).expand(n, m, -1)
    logits = -((a - b)**2).sum(dim=2)

    return logits

def predict(model, data_loader):
    prediction_results = []
    with torch.no_grad():
        # each batch represent one episode (support data + query data)
        for i, (data, target) in enumerate(data_loader):

            # split data into support and query data
            support_input = data[:5 * 1,:,:,:]
            query_input   = data[5 * 1:,:,:,:]
            support_input = support_input.cuda()
            query_input = query_input.cuda()

            label_encoder = {target[i * 1] : i for i in range(5)}
            query_label = torch.cuda.LongTensor([label_encoder[class_name] for class_name in target[5 * 1:]])

            model_output = model(support_input)
            output = model_output.reshape(1, 5, -1).mean(dim=0)# means of classes along n_support shots

            logits = euclidean_dist(model(query_input), output)
            pred = torch.argmax(logits, dim=1).tolist()
            pred.insert(0, i)
            prediction_results.append(pred)

    return prediction_results

if __name__=='__main__':
    pyfile = sys.argv[0]
    test_csv = sys.argv[1]
    test_data_dir = sys.argv[2]
    testcase_csv = sys.argv[3]
    output_csv = sys.argv[4]

    test_dataset = MiniDataset(test_csv, test_data_dir)

    test_loader = DataLoader(
        test_dataset, batch_size=5 * (15 + 1),
        num_workers=3, pin_memory=False, worker_init_fn=worker_init_fn,
        sampler=GeneratorSampler(testcase_csv))

    model = ProNet()
    state = torch.load('./pronet_eucdis_ep184_acc0.5160000118613243±0.01767819107550266.pth')
    model.load_state_dict(state)
    model = model.cuda()
    model.eval()
    prediction_results = predict(model, test_loader)

    episode_label = []
    episode_label.append('episode_id')
    for i in range(15*5):
        episode_label.append('query{}'.format(i))

    prediction_results.insert(0, episode_label)
    with open(output_csv, 'w', newline='') as f:
        writer = csv.writer(f)
        for row in prediction_results:
            writer.writerow(row)