import argparse
import collections
import datetime
import json
import os
import sys

from simpleperf_report_lib import ReportLib
from utils import log_info, log_exit
from utils import Addr2Nearestline, get_script_dir, Objdump, open_report_in_browser
from utils import SourceFileSearcher

MAX_CALLSTACK_LENGTH = 750

class HtmlWriter(object):

    def __init__(self, output_path):
        self.fh = open(output_path, 'w')
        self.tag_stack = []

    def close(self):
        self.fh.close()

    def open_tag(self, tag, **attrs):
        attr_str = ''
        for key in attrs:
            attr_str += ' %s="%s"' % (key, attrs[key])
        self.fh.write('<%s%s>' % (tag, attr_str))
        self.tag_stack.append(tag)
        return self

    def close_tag(self, tag=None):
        if tag:
            assert tag == self.tag_stack[-1]
        self.fh.write('</%s>\n' % self.tag_stack.pop())

    def add(self, text):
        self.fh.write(text)
        return self

    def add_file(self, file_path):
        file_path = os.path.join(get_script_dir(), file_path)
        with open(file_path, 'r') as f:
            self.add(f.read())
        return self

def modify_text_for_html(text):
    return text.replace('>', '>').replace('<', '<')

class EventScope(object):

    def __init__(self, name):
        self.name = name
        self.processes = {}  # map from pid to ProcessScope
        self.sample_count = 0
        self.event_count = 0

    def get_process(self, pid):
        process = self.processes.get(pid)
        if not process:
            process = self.processes[pid] = ProcessScope(pid)
        return process

    def get_sample_info(self, gen_addr_hit_map):
        result = {}
        result['eventName'] = self.name
        result['eventCount'] = self.event_count
        processes = sorted(self.processes.values(), key=lambda a: a.event_count, reverse=True)
        result['processes'] = [process.get_sample_info(gen_addr_hit_map)
                               for process in processes]
        return result

    @property
    def threads(self):
        for process in self.processes.values():
            for thread in process.threads.values():
                yield thread

    @property
    def libraries(self):
        for process in self.processes.values():
            for thread in process.threads.values():
                for lib in thread.libs.values():
                    yield lib

class ProcessScope(object):

    def __init__(self, pid):
        self.pid = pid
        self.name = ''
        self.event_count = 0
        self.threads = {}  # map from tid to ThreadScope

    def get_thread(self, tid, thread_name):
        thread = self.threads.get(tid)
        if not thread:
            thread = self.threads[tid] = ThreadScope(tid)
        thread.name = thread_name
        if self.pid == tid:
            self.name = thread_name
        return thread

    def get_sample_info(self, gen_addr_hit_map):
        result = {}
        result['pid'] = self.pid
        result['eventCount'] = self.event_count
        threads = sorted(self.threads.values(), key=lambda a: a.event_count, reverse=True)
        result['threads'] = [thread.get_sample_info(gen_addr_hit_map)
                             for thread in threads]
        return result

class ThreadScope(object):

    def __init__(self, tid):
        self.tid = tid
        self.name = ''
        self.event_count = 0
        self.sample_count = 0
        self.libs = {}  # map from lib_id to LibScope
        self.call_graph = CallNode(-1)
        self.reverse_call_graph = CallNode(-1)

    def add_callstack(self, event_count, callstack, build_addr_hit_map):
        """ callstack is a list of tuple (lib_id, func_id, addr).
            For each i > 0, callstack[i] calls callstack[i-1]."""
        hit_func_ids = set()
        for i, (lib_id, func_id, addr) in enumerate(callstack):
            # When a callstack contains recursive function, only add for each function once.
            if func_id in hit_func_ids:
                continue
            hit_func_ids.add(func_id)

            lib = self.libs.get(lib_id)
            if not lib:
                lib = self.libs[lib_id] = LibScope(lib_id)
            function = lib.get_function(func_id)
            function.subtree_event_count += event_count
            if i == 0:
                lib.event_count += event_count
                function.event_count += event_count
                function.sample_count += 1
            if build_addr_hit_map:
                function.build_addr_hit_map(addr, event_count if i == 0 else 0, event_count)

        # build call graph and reverse call graph
        node = self.call_graph
        for item in reversed(callstack):
            node = node.get_child(item[1])
        node.event_count += event_count
        node = self.reverse_call_graph
        for item in callstack:
            node = node.get_child(item[1])
        node.event_count += event_count

    def update_subtree_event_count(self):
        self.call_graph.update_subtree_event_count()
        self.reverse_call_graph.update_subtree_event_count()

    def limit_percents(self, min_func_limit, min_callchain_percent, hit_func_ids):
        for lib in self.libs.values():
            to_del_funcs = []
            for function in lib.functions.values():
                if function.subtree_event_count < min_func_limit:
                    to_del_funcs.append(function.func_id)
                else:
                    hit_func_ids.add(function.func_id)
            for func_id in to_del_funcs:
                del lib.functions[func_id]
        min_limit = min_callchain_percent * 0.01 * self.call_graph.subtree_event_count
        self.call_graph.cut_edge(min_limit, hit_func_ids)
        self.reverse_call_graph.cut_edge(min_limit, hit_func_ids)

    def get_sample_info(self, gen_addr_hit_map):
        result = {}
        result['tid'] = self.tid
        result['eventCount'] = self.event_count
        result['sampleCount'] = self.sample_count
        result['libs'] = [lib.gen_sample_info(gen_addr_hit_map)
                          for lib in self.libs.values()]
        result['g'] = self.call_graph.gen_sample_info()
        result['rg'] = self.reverse_call_graph.gen_sample_info()
        return result

class LibScope(object):

    def __init__(self, lib_id):
        self.lib_id = lib_id
        self.event_count = 0
        self.functions = {}  # map from func_id to FunctionScope.

    def get_function(self, func_id):
        function = self.functions.get(func_id)
        if not function:
            function = self.functions[func_id] = FunctionScope(func_id)
        return function

    def gen_sample_info(self, gen_addr_hit_map):
        result = {}
        result['libId'] = self.lib_id
        result['eventCount'] = self.event_count
        result['functions'] = [func.gen_sample_info(gen_addr_hit_map)
                               for func in self.functions.values()]
        return result

class FunctionScope(object):

    def __init__(self, func_id):
        self.func_id = func_id
        self.sample_count = 0
        self.event_count = 0
        self.subtree_event_count = 0
        self.addr_hit_map = None  # map from addr to [event_count, subtree_event_count].
        # map from (source_file_id, line) to [event_count, subtree_event_count].
        self.line_hit_map = None

    def build_addr_hit_map(self, addr, event_count, subtree_event_count):
        if self.addr_hit_map is None:
            self.addr_hit_map = {}
        count_info = self.addr_hit_map.get(addr)
        if count_info is None:
            self.addr_hit_map[addr] = [event_count, subtree_event_count]
        else:
            count_info[0] += event_count
            count_info[1] += subtree_event_count

    def build_line_hit_map(self, source_file_id, line, event_count, subtree_event_count):
        if self.line_hit_map is None:
            self.line_hit_map = {}
        key = (source_file_id, line)
        count_info = self.line_hit_map.get(key)
        if count_info is None:
            self.line_hit_map[key] = [event_count, subtree_event_count]
        else:
            count_info[0] += event_count
            count_info[1] += subtree_event_count

    def gen_sample_info(self, gen_addr_hit_map):
        result = {}
        result['f'] = self.func_id
        result['c'] = [self.sample_count, self.event_count, self.subtree_event_count]
        if self.line_hit_map:
            items = []
            for key in self.line_hit_map:
                count_info = self.line_hit_map[key]
                item = {'f': key[0], 'l': key[1], 'e': count_info[0], 's': count_info[1]}
                items.append(item)
            result['s'] = items
        if gen_addr_hit_map and self.addr_hit_map:
            items = []
            for addr in sorted(self.addr_hit_map):
                count_info = self.addr_hit_map[addr]
                items.append({'a': addr, 'e': count_info[0], 's': count_info[1]})
            result['a'] = items
        return result

class CallNode(object):

    def __init__(self, func_id):
        self.event_count = 0
        self.subtree_event_count = 0
        self.func_id = func_id
        self.children = collections.OrderedDict()  # map from func_id to CallNode

    def get_child(self, func_id):
        child = self.children.get(func_id)
        if not child:
            child = self.children[func_id] = CallNode(func_id)
        return child

    def update_subtree_event_count(self):
        self.subtree_event_count = self.event_count
        for child in self.children.values():
            self.subtree_event_count += child.update_subtree_event_count()
        return self.subtree_event_count

    def cut_edge(self, min_limit, hit_func_ids):
        hit_func_ids.add(self.func_id)
        to_del_children = []
        for key in self.children:
            child = self.children[key]
            if child.subtree_event_count < min_limit:
                to_del_children.append(key)
            else:
                child.cut_edge(min_limit, hit_func_ids)
        for key in to_del_children:
            del self.children[key]

    def gen_sample_info(self):
        result = {}
        result['e'] = self.event_count
        result['s'] = self.subtree_event_count
        result['f'] = self.func_id
        result['c'] = [child.gen_sample_info() for child in self.children.values()]
        return result

class LibSet(object):
    """ Collection of shared libraries used in perf.data. """
    def __init__(self):
        self.lib_name_to_id = {}
        self.lib_id_to_name = []

    def get_lib_id(self, lib_name):
        lib_id = self.lib_name_to_id.get(lib_name)
        if lib_id is None:
            lib_id = len(self.lib_id_to_name)
            self.lib_name_to_id[lib_name] = lib_id
            self.lib_id_to_name.append(lib_name)
        return lib_id

    def get_lib_name(self, lib_id):
        return self.lib_id_to_name[lib_id]

class Function(object):
    """ Represent a function in a shared library. """
    def __init__(self, lib_id, func_name, func_id, start_addr, addr_len):
        self.lib_id = lib_id
        self.func_name = func_name
        self.func_id = func_id
        self.start_addr = start_addr
        self.addr_len = addr_len
        self.source_info = None
        self.disassembly = None

class FunctionSet(object):
    """ Collection of functions used in perf.data. """
    def __init__(self):
        self.name_to_func = {}
        self.id_to_func = {}

    def get_func_id(self, lib_id, symbol):
        key = (lib_id, symbol.symbol_name)
        function = self.name_to_func.get(key)
        if function is None:
            func_id = len(self.id_to_func)
            function = Function(lib_id, symbol.symbol_name, func_id, symbol.symbol_addr,
                                symbol.symbol_len)
            self.name_to_func[key] = function
            self.id_to_func[func_id] = function
        return function.func_id

    def trim_functions(self, left_func_ids):
        """ Remove functions excepts those in left_func_ids. """
        for function in self.name_to_func.values():
            if function.func_id not in left_func_ids:
                del self.id_to_func[function.func_id]
        # name_to_func will not be used.
        self.name_to_func = None

class SourceFile(object):
    """ A source file containing source code hit by samples. """
    def __init__(self, file_id, abstract_path):
        self.file_id = file_id
        self.abstract_path = abstract_path  # path reported by addr2line
        self.real_path = None  # file path in the file system
        self.requested_lines = set()
        self.line_to_code = {}  # map from line to code in that line.

    def request_lines(self, start_line, end_line):
        self.requested_lines |= set(range(start_line, end_line + 1))

    def add_source_code(self, real_path):
        self.real_path = real_path
        with open(real_path, 'r') as f:
            source_code = f.readlines()
        max_line = len(source_code)
        for line in self.requested_lines:
            if line > 0 and line <= max_line:
                self.line_to_code[line] = source_code[line - 1]
        # requested_lines is no longer used.
        self.requested_lines = None

class SourceFileSet(object):
    """ Collection of source files. """
    def __init__(self):
        self.path_to_source_files = {}  # map from file path to SourceFile.

    def get_source_file(self, file_path):
        source_file = self.path_to_source_files.get(file_path)
        if source_file is None:
            source_file = SourceFile(len(self.path_to_source_files), file_path)
            self.path_to_source_files[file_path] = source_file
        return source_file

    def load_source_code(self, source_dirs):
        file_searcher = SourceFileSearcher(source_dirs)
        for source_file in self.path_to_source_files.values():
            real_path = file_searcher.get_real_path(source_file.abstract_path)
            if real_path:
                source_file.add_source_code(real_path)

class RecordData(object):

    """RecordData reads perf.data, and generates data used by report.js in json format.
        All generated items are listed as below:
            1. recordTime: string
            2. machineType: string
            3. androidVersion: string
            4. recordCmdline: string
            5. totalSamples: int
            6. processNames: map from pid to processName.
            7. threadNames: map from tid to threadName.
            8. libList: an array of libNames, indexed by libId.
            9. functionMap: map from functionId to funcData.
                funcData = {
                    l: libId
                    f: functionName
                    s: [sourceFileId, startLine, endLine] [optional]
                    d: [(disassembly, addr)] [optional]
                }

            10.  sampleInfo = [eventInfo]
                eventInfo = {
                    eventName
                    eventCount
                    processes: [processInfo]
                }
                processInfo = {
                    pid
                    eventCount
                    threads: [threadInfo]
                }
                threadInfo = {
                    tid
                    eventCount
                    sampleCount
                    libs: [libInfo],
                    g: callGraph,
                    rg: reverseCallgraph
                }
                libInfo = {
                    libId,
                    eventCount,
                    functions: [funcInfo]
                }
                funcInfo = {
                    f: functionId
                    c: [sampleCount, eventCount, subTreeEventCount]
                    s: [sourceCodeInfo] [optional]
                    a: [addrInfo] (sorted by addrInfo.addr) [optional]
                }
                callGraph and reverseCallGraph are both of type CallNode.
                callGraph shows how a function calls other functions.
                reverseCallGraph shows how a function is called by other functions.
                CallNode {
                    e: selfEventCount
                    s: subTreeEventCount
                    f: functionId
                    c: [CallNode] # children
                }

                sourceCodeInfo {
                    f: sourceFileId
                    l: line
                    e: eventCount
                    s: subtreeEventCount
                }

                addrInfo {
                    a: addr
                    e: eventCount
                    s: subtreeEventCount
                }

            11. sourceFiles: an array of sourceFile, indexed by sourceFileId.
                sourceFile {
                    path
                    code:  # a map from line to code for that line.
                }
    """

    def __init__(self, binary_cache_path, ndk_path, build_addr_hit_map):
        self.binary_cache_path = binary_cache_path
        self.ndk_path = ndk_path
        self.build_addr_hit_map = build_addr_hit_map
        self.meta_info = None
        self.cmdline = None
        self.arch = None
        self.events = {}
        self.libs = LibSet()
        self.functions = FunctionSet()
        self.total_samples = 0
        self.source_files = SourceFileSet()
        self.gen_addr_hit_map_in_record_info = False

    def load_record_file(self, record_file, show_art_frames):
        lib = ReportLib()
        lib.SetRecordFile(record_file)
        # If not showing ip for unknown symbols, the percent of the unknown symbol may be
        # accumulated to very big, and ranks first in the sample table.
        lib.ShowIpForUnknownSymbol()
        if show_art_frames:
            lib.ShowArtFrames()
        if self.binary_cache_path:
            lib.SetSymfs(self.binary_cache_path)
        self.meta_info = lib.MetaInfo()
        self.cmdline = lib.GetRecordCmd()
        self.arch = lib.GetArch()
        while True:
            raw_sample = lib.GetNextSample()
            if not raw_sample:
                lib.Close()
                break
            raw_event = lib.GetEventOfCurrentSample()
            symbol = lib.GetSymbolOfCurrentSample()
            callchain = lib.GetCallChainOfCurrentSample()
            event = self._get_event(raw_event.name)
            self.total_samples += 1
            event.sample_count += 1
            event.event_count += raw_sample.period
            process = event.get_process(raw_sample.pid)
            process.event_count += raw_sample.period
            thread = process.get_thread(raw_sample.tid, raw_sample.thread_comm)
            thread.event_count += raw_sample.period
            thread.sample_count += 1

            lib_id = self.libs.get_lib_id(symbol.dso_name)
            func_id = self.functions.get_func_id(lib_id, symbol)
            callstack = [(lib_id, func_id, symbol.vaddr_in_file)]
            for i in range(callchain.nr):
                symbol = callchain.entries[i].symbol
                lib_id = self.libs.get_lib_id(symbol.dso_name)
                func_id = self.functions.get_func_id(lib_id, symbol)
                callstack.append((lib_id, func_id, symbol.vaddr_in_file))
            if len(callstack) > MAX_CALLSTACK_LENGTH:
                callstack = callstack[:MAX_CALLSTACK_LENGTH]
            thread.add_callstack(raw_sample.period, callstack, self.build_addr_hit_map)

        for event in self.events.values():
            for thread in event.threads:
                thread.update_subtree_event_count()

    def limit_percents(self, min_func_percent, min_callchain_percent):
        hit_func_ids = set()
        for event in self.events.values():
            min_limit = event.event_count * min_func_percent * 0.01
            for process in event.processes.values():
                to_del_threads = []
                for thread in process.threads.values():
                    if thread.call_graph.subtree_event_count < min_limit:
                        to_del_threads.append(thread.tid)
                    else:
                        thread.limit_percents(min_limit, min_callchain_percent, hit_func_ids)
                for thread in to_del_threads:
                    del process.threads[thread]
        self.functions.trim_functions(hit_func_ids)

    def _get_event(self, event_name):
        if event_name not in self.events:
            self.events[event_name] = EventScope(event_name)
        return self.events[event_name]

    def add_source_code(self, source_dirs, filter_lib):
        """ Collect source code information:
            1. Find line ranges for each function in FunctionSet.
            2. Find line for each addr in FunctionScope.addr_hit_map.
            3. Collect needed source code in SourceFileSet.
        """
        addr2line = Addr2Nearestline(self.ndk_path, self.binary_cache_path, False)
        # Request line range for each function.
        for function in self.functions.id_to_func.values():
            if function.func_name == 'unknown':
                continue
            lib_name = self.libs.get_lib_name(function.lib_id)
            if filter_lib(lib_name):
                addr2line.add_addr(lib_name, function.start_addr, function.start_addr)
                addr2line.add_addr(lib_name, function.start_addr,
                                   function.start_addr + function.addr_len - 1)
        # Request line for each addr in FunctionScope.addr_hit_map.
        for event in self.events.values():
            for lib in event.libraries:
                lib_name = self.libs.get_lib_name(lib.lib_id)
                if filter_lib(lib_name):
                    for function in lib.functions.values():
                        func_addr = self.functions.id_to_func[function.func_id].start_addr
                        for addr in function.addr_hit_map:
                            addr2line.add_addr(lib_name, func_addr, addr)
        addr2line.convert_addrs_to_lines()

        # Set line range for each function.
        for function in self.functions.id_to_func.values():
            if function.func_name == 'unknown':
                continue
            dso = addr2line.get_dso(self.libs.get_lib_name(function.lib_id))
            if not dso:
                continue
            start_source = addr2line.get_addr_source(dso, function.start_addr)
            end_source = addr2line.get_addr_source(dso, function.start_addr + function.addr_len - 1)
            if not start_source or not end_source:
                continue
            start_file_path, start_line = start_source[-1]
            end_file_path, end_line = end_source[-1]
            if start_file_path != end_file_path or start_line > end_line:
                continue
            source_file = self.source_files.get_source_file(start_file_path)
            source_file.request_lines(start_line, end_line)
            function.source_info = (source_file.file_id, start_line, end_line)

        # Build FunctionScope.line_hit_map.
        for event in self.events.values():
            for lib in event.libraries:
                dso = addr2line.get_dso(self.libs.get_lib_name(lib.lib_id))
                if not dso:
                    continue
                for function in lib.functions.values():
                    for addr in function.addr_hit_map:
                        source = addr2line.get_addr_source(dso, addr)
                        if not source:
                            continue
                        for file_path, line in source:
                            source_file = self.source_files.get_source_file(file_path)
                            # Show [line - 5, line + 5] of the line hit by a sample.
                            source_file.request_lines(line - 5, line + 5)
                            count_info = function.addr_hit_map[addr]
                            function.build_line_hit_map(source_file.file_id, line, count_info[0],
                                                        count_info[1])

        # Collect needed source code in SourceFileSet.
        self.source_files.load_source_code(source_dirs)

    def add_disassembly(self, filter_lib):
        """ Collect disassembly information:
            1. Use objdump to collect disassembly for each function in FunctionSet.
            2. Set flag to dump addr_hit_map when generating record info.
        """
        objdump = Objdump(self.ndk_path, self.binary_cache_path)
        cur_lib_name = None
        dso_info = None
        for function in sorted(self.functions.id_to_func.values(), key=lambda a: a.lib_id):
            if function.func_name == 'unknown':
                continue
            lib_name = self.libs.get_lib_name(function.lib_id)
            if lib_name != cur_lib_name:
                cur_lib_name = lib_name
                if filter_lib(lib_name):
                    dso_info = objdump.get_dso_info(lib_name)
                else:
                    dso_info = None
                if dso_info:
                    log_info('Disassemble %s' % dso_info[0])
            if dso_info:
                code = objdump.disassemble_code(dso_info, function.start_addr, function.addr_len)
                function.disassembly = code

        self.gen_addr_hit_map_in_record_info = True

    def gen_record_info(self):
        record_info = {}
        timestamp = self.meta_info.get('timestamp')
        if timestamp:
            t = datetime.datetime.fromtimestamp(int(timestamp))
        else:
            t = datetime.datetime.now()
        record_info['recordTime'] = t.strftime('%Y-%m-%d (%A) %H:%M:%S')

        product_props = self.meta_info.get('product_props')
        machine_type = self.arch
        if product_props:
            manufacturer, model, name = product_props.split(':')
            machine_type = '%s (%s) by %s, arch %s' % (model, name, manufacturer, self.arch)
        record_info['machineType'] = machine_type
        record_info['androidVersion'] = self.meta_info.get('android_version', '')
        record_info['recordCmdline'] = self.cmdline
        record_info['totalSamples'] = self.total_samples
        record_info['processNames'] = self._gen_process_names()
        record_info['threadNames'] = self._gen_thread_names()
        record_info['libList'] = self._gen_lib_list()
        record_info['functionMap'] = self._gen_function_map()
        record_info['sampleInfo'] = self._gen_sample_info()
        record_info['sourceFiles'] = self._gen_source_files()
        return record_info

    def _gen_process_names(self):
        process_names = {}
        for event in self.events.values():
            for process in event.processes.values():
                process_names[process.pid] = process.name
        return process_names

    def _gen_thread_names(self):
        thread_names = {}
        for event in self.events.values():
            for process in event.processes.values():
                for thread in process.threads.values():
                    thread_names[thread.tid] = thread.name
        return thread_names

    def _gen_lib_list(self):
        return [modify_text_for_html(x) for x in self.libs.lib_id_to_name]

    def _gen_function_map(self):
        func_map = {}
        for func_id in sorted(self.functions.id_to_func):
            function = self.functions.id_to_func[func_id]
            func_data = {}
            func_data['l'] = function.lib_id
            func_data['f'] = modify_text_for_html(function.func_name)
            if function.source_info:
                func_data['s'] = function.source_info
            if function.disassembly:
                disassembly_list = []
                for code, addr in function.disassembly:
                    disassembly_list.append([modify_text_for_html(code), addr])
                func_data['d'] = disassembly_list
            func_map[func_id] = func_data
        return func_map

    def _gen_sample_info(self):
        return [event.get_sample_info(self.gen_addr_hit_map_in_record_info)
                for event in self.events.values()]

    def _gen_source_files(self):
        source_files = sorted(self.source_files.path_to_source_files.values(),
                              key=lambda x: x.file_id)
        file_list = []
        for source_file in source_files:
            file_data = {}
            if not source_file.real_path:
                file_data['path'] = ''
                file_data['code'] = {}
            else:
                file_data['path'] = source_file.real_path
                code_map = {}
                for line in source_file.line_to_code:
                    code_map[line] = modify_text_for_html(source_file.line_to_code[line])
                file_data['code'] = code_map
            file_list.append(file_data)
        return file_list

URLS = {
    'jquery': 'https://ajax.googleapis.com/ajax/libs/jquery/3.3.1/jquery.min.js',
    'bootstrap4-css': 'https://stackpath.bootstrapcdn.com/bootstrap/4.1.2/css/bootstrap.min.css',
    'bootstrap4-popper':
        'https://cdnjs.cloudflare.com/ajax/libs/popper.js/1.12.9/umd/popper.min.js',
    'bootstrap4': 'https://stackpath.bootstrapcdn.com/bootstrap/4.1.2/js/bootstrap.min.js',
    'dataTable': 'https://cdn.datatables.net/1.10.19/js/jquery.dataTables.min.js',
    'dataTable-bootstrap4': 'https://cdn.datatables.net/1.10.19/js/dataTables.bootstrap4.min.js',
    'dataTable-css': 'https://cdn.datatables.net/1.10.19/css/dataTables.bootstrap4.min.css',
    'gstatic-charts': 'https://www.gstatic.com/charts/loader.js',
}

class ReportGenerator(object):

    def __init__(self, html_path):
        self.hw = HtmlWriter(html_path)
        self.hw.open_tag('html')
        self.hw.open_tag('head')
        for css in ['bootstrap4-css', 'dataTable-css']:
            self.hw.open_tag('link', rel='stylesheet', type='text/css', href=URLS[css]).close_tag()
        for js in ['jquery', 'bootstrap4-popper', 'bootstrap4', 'dataTable', 'dataTable-bootstrap4',
                   'gstatic-charts']:
            self.hw.open_tag('script', src=URLS[js]).close_tag()

        self.hw.open_tag('script').add(
            "google.charts.load('current', {'packages': ['corechart', 'table']});").close_tag()
        self.hw.open_tag('style', type='text/css').add("""
            .colForLine { width: 50px; }
            .colForCount { width: 100px; }
            .tableCell { font-size: 17px; }
            .boldTableCell { font-weight: bold; font-size: 17px; }
            """).close_tag()
        self.hw.close_tag('head')
        self.hw.open_tag('body')
        self.record_info = {}

    def write_content_div(self):
        self.hw.open_tag('div', id='report_content').close_tag()

    def write_record_data(self, record_data):
        self.hw.open_tag('script', id='record_data', type='application/json')
        self.hw.add(json.dumps(record_data))
        self.hw.close_tag()

    def write_script(self):
        self.hw.open_tag('script').add_file('report_html.js').close_tag()

    def finish(self):
        self.hw.close_tag('body')
        self.hw.close_tag('html')
        self.hw.close()

def main():
    sys.setrecursionlimit(MAX_CALLSTACK_LENGTH * 2 + 50)
    parser = argparse.ArgumentParser(description='report profiling data')
    parser.add_argument('-i', '--record_file', nargs='+', default=['perf.data'], help="""
                        Set profiling data file to report. Default is perf.data.""")
    parser.add_argument('-o', '--report_path', default='report.html', help="""
                        Set output html file. Default is report.html.""")
    parser.add_argument('--min_func_percent', default=0.01, type=float, help="""
                        Set min percentage of functions shown in the report.
                        For example, when set to 0.01, only functions taking >= 0.01%% of total
                        event count are collected in the report. Default is 0.01.""")
    parser.add_argument('--min_callchain_percent', default=0.01, type=float, help="""
                        Set min percentage of callchains shown in the report.
                        It is used to limit nodes shown in the function flamegraph. For example,
                        when set to 0.01, only callchains taking >= 0.01%% of the event count of
                        the starting function are collected in the report. Default is 0.01.""")
    parser.add_argument('--add_source_code', action='store_true', help='Add source code.')
    parser.add_argument('--source_dirs', nargs='+', help='Source code directories.')
    parser.add_argument('--add_disassembly', action='store_true', help='Add disassembled code.')
    parser.add_argument('--binary_filter', nargs='+', help="""Annotate source code and disassembly
                        only for selected binaries.""")
    parser.add_argument('--ndk_path', nargs=1, help='Find tools in the ndk path.')
    parser.add_argument('--no_browser', action='store_true', help="Don't open report in browser.")
    parser.add_argument('--show_art_frames', action='store_true',
                        help='Show frames of internal methods in the ART Java interpreter.')
    args = parser.parse_args()

    # 1. Process args.
    binary_cache_path = 'binary_cache'
    if not os.path.isdir(binary_cache_path):
        if args.add_source_code or args.add_disassembly:
            log_exit("""binary_cache/ doesn't exist. Can't add source code or disassembled code
                        without collected binaries. Please run binary_cache_builder.py to
                        collect binaries for current profiling data, or run app_profiler.py
                        without -nb option.""")
        binary_cache_path = None

    if args.add_source_code and not args.source_dirs:
        log_exit('--source_dirs is needed to add source code.')
    build_addr_hit_map = args.add_source_code or args.add_disassembly
    ndk_path = None if not args.ndk_path else args.ndk_path[0]

    # 2. Produce record data.
    record_data = RecordData(binary_cache_path, ndk_path, build_addr_hit_map)
    for record_file in args.record_file:
        record_data.load_record_file(record_file, args.show_art_frames)
    record_data.limit_percents(args.min_func_percent, args.min_callchain_percent)

    def filter_lib(lib_name):
        if not args.binary_filter:
            return True
        for binary in args.binary_filter:
            if binary in lib_name:
                return True
        return False
    if args.add_source_code:
        record_data.add_source_code(args.source_dirs, filter_lib)
    if args.add_disassembly:
        record_data.add_disassembly(filter_lib)

    # 3. Generate report html.
    report_generator = ReportGenerator(args.report_path)
    report_generator.write_script()
    report_generator.write_content_div()
    report_generator.write_record_data(record_data.gen_record_info())
    report_generator.finish()

    if not args.no_browser:
        open_report_in_browser(args.report_path)
    log_info("Report generated at '%s'." % args.report_path)

if __name__ == '__main__':
    main()


import os
import pytest

import coreir
import magma as m

@pytest.mark.skip("No longer supported functionality")
def test_compile_coreir():
    width = 16
    numInputs = 4
    doubleT = m.Bits[width]

    class double(m.Circuit):
        name = "double"
        io = m.IO(I=m.In(doubleT), O=m.Out(doubleT))
        shift_amount = 2
        output = m.concat(io.I[shift_amount:width], m.bits(0, shift_amount))
        m.wire(output, io.O)

    coreir_double = m.backend.coreir.coreir_backend.compile(double)
    c = coreir_double.context

    def get_lib(lib):
        if lib in {"coreir", "mantle", "corebit"}:
            return c.get_namespace(lib)
        elif lib == "global":
            return c.global_namespace
        else:
            return c.load_library(lib)

    def import_(lib, name):
        return get_lib(lib).generators[name]

    mapParallelParams = c.new_values({"numInputs": numInputs, "operator": coreir_double})

    test_module_typ = c.Record({"in": c.Array(numInputs, c.Array(width,
        c.BitIn())), "out": c.Array(numInputs, c.Array(width, c.Bit()))})
    test_module = c.global_namespace.new_module("test_module", test_module_typ)
    test_module_def = test_module.new_definition()
    mapParallel = import_("aetherlinglib", "mapParallel")
    mapMod = mapParallel(numInputs=numInputs, operator=coreir_double)
    mapDouble = test_module_def.add_module_instance("mapDouble", mapMod)
    test_module_def.connect(test_module_def.interface.select("in"), mapDouble.select("I"));
    test_module_def.connect(mapDouble.select("O"), test_module_def.interface.select("out"));
    test_module_def.print_()
    test_module.definition = test_module_def
    test_module.print_()
    dir_path = os.path.dirname(os.path.realpath(__file__))
    test_module.save_to_file(os.path.join(dir_path, "mapParallel_test.json"))
    with open(os.path.join(dir_path, "mapParallel_test.json"), "r") as actual:
        with open(os.path.join(dir_path, "mapParallel_test_gold.json"), "r") as gold:
            assert actual.read() == gold.read()
    mod = c.load_from_file(os.path.join(dir_path, "mapParallel_test.json"))
    mod.print_()


from flask import Flask, render_template,request,make_response, request

from models import Links
import json

model_obj = Links()

app = Flask(__name__)

@app.route("/")
def index():
  return render_template("index.html")

@app.route("/links")
def links():
    global model_obj
    datas=model_obj.all()
    return json.dumps(datas)

@app.route("/search")
def search():
  global model_obj
  print(request.args)

  category = request.args['cat']
  content = request.args['content']

  searched_data = model_obj.search(cat=category, content=content)
  print(searched_data)
  return searched_data

if __name__ == "__main__":
  app.run(host="0.0.0.0",debug=True)


import numpy as np
import xarray as xr
import pandas as pd
from .models import DLWPNeuralNet, DLWPFunctional
from .models_torch import DLWPTorchNN
from .generators import DataGenerator, SeriesDataGenerator, ArrayDataGenerator
from ..util import insolation
import warnings

class TimeSeriesEstimator(object):
    """
    Sophisticated wrapper class for producing time series forecasts from a DLWP model, using a Generator with
    metadata. This class allows predictions with non-matching inputs and outputs, including in the variable, level, and
    time_step dimensions.
    """

    def __init__(self, model, generator):
        """
        Initialize a TimeSeriesEstimator from a model and generator.

        :param model: DLWP model instance
        :param generator: DLWP DataGenerator instance
        """
        if not isinstance(model, (DLWPNeuralNet, DLWPFunctional, DLWPTorchNN)):
            raise TypeError("'model' must be a valid instance of a DLWP model class")
        if not isinstance(generator, (DataGenerator, SeriesDataGenerator, ArrayDataGenerator)):
            raise TypeError("'generator' must be a valid instance of a DLWP generator class")
        if isinstance(model, DLWPFunctional):
            warnings.warn('DLWPFunctional models are only partially supported by TimeSeriesEstimator. The '
                          'inputs/outputs to the model must be the same if the model predicts a sequence.')
        self.model = model
        self.generator = generator
        self._add_insolation = generator._add_insolation if hasattr(generator, '_add_insolation') else False
        self._uses_varlev = 'varlev' in generator.ds.dims
        self._has_targets = 'targets' in generator.ds.variables
        self._is_series = isinstance(generator, (SeriesDataGenerator, ArrayDataGenerator))
        self._output_sel = {}
        self._input_sel = {}
        self._dt = self.generator.ds['sample'][1] - self.generator.ds['sample'][0]
        if hasattr(self.generator, '_interval'):
            self._interval = self.generator._interval
        else:
            self._interval = 1
        if hasattr(self.generator, 'rank'):
            self.rank = self.generator.rank
        else:
            self.rank = 2

        # Generate the selections needed for inputs and outputs
        if self._uses_varlev:
            # Find the outputs we keep
            if self._is_series:  # a SeriesDataGenerator has user-specified I/O
                if not self.generator._output_sel:  # selection was empty
                    self._output_sel = {'varlev': np.array(self.generator.ds.coords['varlev'][:])}
                else:
                    self._output_sel = {k: np.array(v) for k, v in self.generator._output_sel.items()}
            else:
                if self._has_targets:
                    self._output_sel = {'varlev': np.array(self.generator.ds.targets.coords['varlev'][:])}
                else:
                    self._output_sel = {'varlev': np.array(self.generator.ds.coords['varlev'][:])}
            # Find the inputs we need
            if self._is_series:
                if not self.generator._input_sel:  # selection was empty
                    self._input_sel = {'varlev': np.array(self.generator.ds.coords['varlev'][:])}
                else:
                    self._input_sel = {k: np.array(v) for k, v in self.generator._input_sel.items()}
            else:
                self._input_sel = {'varlev': np.array(self.generator.ds.predictors.coords['varlev'][:])}
            # Find outputs that need to replace inputs
            self._outputs_in_inputs = {
                'varlev': np.array([v for v in self._output_sel['varlev'] if v in self._input_sel['varlev']])
            }
        else:  # Uses variable/level coordinates
            # Find the outputs we keep
            if self._is_series:  # a SeriesDataGenerator has user-specified I/O
                if not self.generator._output_sel:  # selection was empty
                    self._output_sel = {'variable': np.array(self.generator.ds.coords['variable'][:]),
                                        'level': np.array(self.generator.ds.coords['level'][:])}
                else:
                    self._output_sel = {k: np.array(v) for k, v in self.generator._output_sel.items()}
                    if 'variable' not in self._output_sel.keys():
                        self._output_sel['variable'] = np.array(self.generator.ds.coords['variable'][:])
                    if 'level' not in self._output_sel.keys():
                        self._output_sel['level'] = np.array(self.generator.ds.coords['level'][:])
            else:
                if self._has_targets:
                    self._output_sel = {'variable': np.array(self.generator.ds.targets.coords['variable'][:]),
                                        'level': np.array(self.generator.ds.targets.coords['level'][:])}
                else:
                    self._output_sel = {'variable': np.array(self.generator.ds.coords['variable'][:]),
                                        'level': np.array(self.generator.ds.coords['level'][:])}
            # Find the inputs we need
            if self._is_series:
                if not self.generator._input_sel:  # selection was empty
                    self._input_sel = {'variable': np.array(self.generator.ds.coords['variable'][:]),
                                       'level': np.array(self.generator.ds.coords['level'][:])}
                else:
                    self._input_sel = {k: np.array(v) for k, v in self.generator._input_sel.items()}
                    if 'variable' not in self._input_sel.keys():
                        self._input_sel['variable'] = np.array(self.generator.ds.coords['variable'][:])
                    if 'level' not in self._input_sel.keys():
                        self._input_sel['level'] = np.array(self.generator.ds.coords['level'][:])
            else:
                self._input_sel = {'variable': np.array(self.generator.ds.predictors.coords['variable'][:]),
                                   'level': np.array(self.generator.ds.predictors.coords['level'][:])}
            # Flatten variable/level
            lev, var = np.meshgrid(self._input_sel['level'], self._input_sel['variable'])
            varlev = np.array(['/'.join([v, str(l)]) for v, l in zip(var.flatten(), lev.flatten())])
            self._input_sel['varlev'] = varlev
            lev, var = np.meshgrid(self._output_sel['level'], self._output_sel['variable'])
            varlev = np.array(['/'.join([v, str(l)]) for v, l in zip(var.flatten(), lev.flatten())])
            self._output_sel['varlev'] = varlev
            # Find outputs that need to replace inputs
            self._outputs_in_inputs = {
                'variable': np.array([v for v in self._output_sel['variable'] if v in self._input_sel['variable']]),
                'level': np.array([v for v in self._output_sel['level'] if v in self._input_sel['level']]),
                'varlev': np.array([v for v in self._output_sel['varlev'] if v in self._input_sel['varlev']])
            }
        if self._add_insolation:
            self._input_sel['varlev'] = np.concatenate([self._input_sel['varlev'], np.array(['SOL'])])

        # Time step dimension
        self._input_time_steps = (generator._input_time_steps if isinstance(generator, SeriesDataGenerator)
                                  else model.time_dim)
        self._output_time_steps = (generator._output_time_steps if isinstance(generator, SeriesDataGenerator)
                                   else model.time_dim)

        # Channels last option
        self.channels_last = hasattr(self.generator, 'channels_last') and self.generator.channels_last
        self._forward_transpose = (0, self.rank + 1) + tuple(range(1, 1 + self.rank)) + (-1,)
        self._backward_transpose = (0,) + tuple(range(2, 2 + self.rank)) + (1, -1,)

        # Constants in the generator
        if hasattr(self.generator, 'constants') and self.generator.constants is not None:
            if self.channels_last:
                self.constants = self.generator.constants.transpose(tuple(range(1, 1 + self.rank)) + (0,))
            else:
                self.constants = self.generator.constants
        else:
            self.constants = None

    @property
    def shape(self):
        return (self.generator._n_sample,) + self.generator.shape

    @property
    def convolution_shape(self):
        return (self.generator._n_sample,) + self.generator.convolution_shape

    def predict(self, steps, samples=(), impute=False, keep_time_dim=False, prefer_first_times=True,
                f_hour_timedelta_type=False, **kwargs):
        """
        Step forward the time series prediction from the model 'steps' times, feeding predictions back in as
        inputs. Predicts for all the data provided in the generator. If there are inputs which are not produced by
        the model outputs, we include the available inputs from the generator data and either reduce the number of
        predicted samples accordingly (remove those whose inputs cannot be satisfied) or run the model using the mean
        values of the inputs which cannot be satisfied. If there are fewer output time steps than input time steps,
        then we build a time series forecast intelligently using part of the predictors and part of the prediction at
        every step. Note only the SeriesDataGenerator supports variable inputs/outputs.

        :param steps: int: number of times to step forward
        :param samples: list of int: which samples in the generator to predict for. For all samples, pass an empty
            list or tuple. May cause unexpected behavior when the input and output data do not match.
        :param impute: bool: if True, use the mean state for missing inputs in the forward integration
        :param keep_time_dim: bool: if True, keep the time_step dimension instead of integrating it with forecsat_hour
            to produce a continuous time series
        :param prefer_first_times: bool: in the case where the prediction contains more time_steps than the input,
            use the first available predicted times to initialize the next step, otherwise use the last times. If the
            output time_steps is less than the input time_steps, we always use all of the output times.
        :param f_hour_timedelta_type: bool: if True, converts f_hour dimension into a timedelta type. May not always be
            compatible with netCDF applications.
        :param kwargs: passed to Keras.predict()
        :return: ndarray: predicted states with forecast_step as the first dimension
        """
        if int(steps) < 1:
            raise ValueError('must use positive integer for steps')

        # Effective forward time steps for each step
        if self._output_time_steps <= self._input_time_steps:
            keep_inputs = True
            es = self._output_time_steps
            in_times = np.arange(self._input_time_steps) - (self._input_time_steps - self._output_time_steps)
        else:
            keep_inputs = False
            if prefer_first_times:
                es = self._input_time_steps
                in_times = np.arange(self._input_time_steps)
            else:
                es = self._output_time_steps
                in_times = np.arange(self._input_time_steps) + (self._output_time_steps - self._input_time_steps)
        effective_steps = int(np.ceil(steps / es))

        # Load data from the generator, without any scaling as this will be done by the model's predict method
        predictors, t = self.generator.generate(samples, scale_and_impute=False)
        p = predictors[0] if isinstance(predictors, (list, tuple)) else predictors
        p_shape = tuple(p.shape)

        # Add metadata. The sample dimension denotes the *start* time of the sample, for insolation purposes. This is
        # then corrected when providing the final output time series.
        sample_coord = self.generator.ds.sample[:self.generator._n_sample] if len(samples) == 0 else \
            self.generator.ds.sample[samples]
        if not self._is_series:
            sample_coord = sample_coord - self._dt * (self._input_time_steps - 1)
        if self.channels_last:
            # Split time step/varlev at the end and transpose
            if not self.generator._keep_time_axis:
                p = p.reshape((p_shape[0],) + self.generator.convolution_shape[-self.rank-1:-1] +
                              (self._input_time_steps, -1)).transpose(self._forward_transpose)
            p_da = xr.DataArray(
                p,
                coords=([sample_coord, in_times] +
                        [np.arange(d) for d in self.generator.convolution_shape[-self.rank-1:-1]] +
                        [self._input_sel['varlev']]),
                dims=['sample', 'time_step'] + ['x%d' % d for d in range(self.rank)] + ['varlev']
            )
        else:
            p = p.reshape((p_shape[0], self._input_time_steps, -1,) + self.generator.convolution_shape[-self.rank:])
            p_da = xr.DataArray(
                p,
                coords=([sample_coord, in_times, self._input_sel['varlev']] +
                        [np.arange(d) for d in self.generator.convolution_shape[-self.rank:]]),
                dims=['sample', 'time_step', 'varlev'] + ['x%d' % d for d in range(self.rank)]
            )
        if self.rank == 2:
            p_da = p_da.rename({'x0': 'lat', 'x1': 'lon'}).assign_coords(lat=self.generator.ds.lat,
                                                                         lon=self.generator.ds.lon)

        # Calculate mean for imputing
        if impute:
            p_mean = p.mean(axis=0)

        # Target shape
        t_shape = t[0].shape if isinstance(t, (list, tuple)) else t.shape
        t = None

        # A DLWPFunctional model which does not have the same inputs/outputs must have some programmatic way of fixing
        # this issue built-in (if it is trained to minimize several iterations of the model). Thus we fall back to the
        # regular predict_timeseries method. However, such a model that does not predict a sequence is perfectly valid
        # for use here.
        if isinstance(self.model, DLWPFunctional) and self.model._n_steps > 1:
            if not self.generator._add_insolation:
                # For the DLWPFunctional model, just use its predict_timeseries API, which handles effective steps
                # TODO: correctly handle channels_last
                result = self.model.predict_timeseries(predictors, steps, keep_time_dim=True,
                                                       **kwargs).reshape((-1,) + t_shape)[:effective_steps, ...]
                result = result.transpose((0, 1) + tuple(range(2, len(result.shape)))).copy()
            else:
                # If insolation is requested, intelligently add it in the same way the generator does
                sequence_steps = int(np.ceil(steps / self.model._n_steps / self.model.time_dim))

                # Giant forecast array
                result = np.full((t_shape[0], sequence_steps, self.model._n_steps) + t_shape[1:],
                                 np.nan, dtype=np.float32)

                # Iterate
                new_t = p_da.sample[:]
                for s in range(sequence_steps):
                    if 'verbose' in kwargs and kwargs['verbose'] > 0:
                        print('Time step %d/%d' % (s + 1, sequence_steps))
                    result[:, s] = np.stack(self.model.predict(predictors, **kwargs), axis=1)

                    # Assign new insolation to list of inputs
                    new_t = new_t + self._output_time_steps * self.model._n_steps * self._dt
                    new_insolation = [np.concatenate(
                        [np.expand_dims(insolation(new_t + (n + m * self._input_time_steps) * self._dt,
                                                   self.generator.ds.lat.values,
                                                   self.generator.ds.lon.values)[:, None],
                                        axis=-1 if self.channels_last else 1)
                         for n in range(self._input_time_steps)],
                        axis=1) for m in range(self.model._n_steps)]

                    if self.channels_last:
                        if self.generator._keep_time_axis:
                            r = result[:, s, -1]
                            predictors = [np.concatenate([r, new_insolation[0]], axis=-1).reshape(
                                (-1,) + self.convolution_shape[1:])] + new_insolation[1:]
                        else:
                            r = result[:, s, -1].reshape(t_shape[:-1] + (self._output_time_steps, -1)).transpose(
                                self._forward_transpose
                            )
                            predictors = [np.concatenate([r, new_insolation[0]], axis=-1).transpose(
                                self._backward_transpose).reshape((-1,) + self.convolution_shape[1:])] \
                                + new_insolation[1:]
                    else:
                        predictors = [
                            np.concatenate([result[:, s, -1].reshape((-1,) + self.shape[1:]), new_insolation[0]],
                                           axis=2).reshape((-1,) + self.convolution_shape[1:])
                        ] + new_insolation[1:]

                    # Add constants
                    if self.constants is not None:
                        predictors.append(np.repeat(np.expand_dims(self.constants, axis=0), len(p_da.sample), axis=0))

                n_dim_1 = result.size // int(np.prod(t_shape))
                result.shape = (t_shape[0], n_dim_1,) + t_shape[1:]
                result = result[:, :effective_steps, ...]
        else:
            # Giant forecast array
            result = np.full((t_shape[0], effective_steps,) + t_shape[1:], np.nan, dtype=np.float32)

            # Iterate prediction forward for a regular DLWP Sequential NN
            for s in range(effective_steps):
                if 'verbose' in kwargs and kwargs['verbose'] > 0:
                    print('Time step %d/%d' % (s + 1, effective_steps))
                if self.channels_last and not self.generator._keep_time_axis:
                    result[:, s] = self.model.predict(p_da.values.transpose(self._backward_transpose).reshape(p_shape),
                                                      **kwargs)
                else:
                    result[:, s] = self.model.predict(p_da.values.reshape(p_shape), **kwargs)

                # Add metadata to the prediction
                if self.channels_last:
                    if not self.generator._keep_time_axis:
                        r = result[:, s].reshape((p_shape[0],) + self.generator.convolution_shape[-self.rank-1:-1] +
                                                 (self._output_time_steps, -1)).transpose(self._forward_transpose)
                    else:
                        r = result[:, s]
                    r_da = xr.DataArray(
                        r,
                        coords=([p_da.sample + (es + self._interval - 1) * self._dt,
                                 np.arange(self._output_time_steps)] +
                                [np.arange(d) for d in self.generator.convolution_shape[-self.rank-1:-1]] +
                                [self._output_sel['varlev']]),
                        dims=['sample', 'time_step'] + ['x%d' % d for d in range(self.rank)] + ['varlev']
                    )
                else:
                    r_da = xr.DataArray(
                        result[:, s].reshape((p_shape[0], self._output_time_steps, -1,) +
                                             self.generator.convolution_shape[-self.rank:]),
                        coords=([p_da.sample + (es + self._interval - 1) * self._dt,
                                 np.arange(self._output_time_steps), self._output_sel['varlev']] +
                                [np.arange(d) for d in self.generator.convolution_shape[-self.rank:]]),
                        dims=['sample', 'time_step', 'varlev'] + ['x%d' % d for d in range(self.rank)]
                    )
                if self.rank == 2:
                    r_da = r_da.rename({'x0': 'lat', 'x1': 'lon'}).assign_coords(lat=self.generator.ds.lat,
                                                                                lon=self.generator.ds.lon)

                # Re-index the predictors to the new forward time step
                p_da = p_da.reindex(sample=r_da.sample, method=None)

                # Impute values extending beyond data availability
                if impute:
                    # Calculate mean values for the added time steps after re-indexing
                    p_da[-es:] = np.concatenate([p_mean[np.newaxis, ...]] * es)

                # Take care of the known insolation for added time steps
                if self._add_insolation:
                    p_da.loc[{'varlev': 'SOL'}][-es:] = \
                        np.concatenate([insolation(p_da.sample[-es:] + n * self._dt, self.generator.ds.lat.values,
                                                   self.generator.ds.lon.values)[:, np.newaxis]
                                        for n in range(self._input_time_steps)], axis=1)

                # Replace the predictors that exist in the result with the result. Any that do not exist are
                # automatically inherited from the known predictor data (or imputed data).
                if keep_inputs:
                    loc_dict = dict(varlev=self._outputs_in_inputs['varlev'], time_step=p_da.time_step[-es:])
                    p_da.loc[loc_dict] = r_da.loc[{'varlev': self._outputs_in_inputs['varlev']}]
                else:
                    if prefer_first_times:
                        p_da.loc[{'varlev': self._outputs_in_inputs['varlev']}] = \
                            r_da.loc[{'varlev': self._outputs_in_inputs['varlev']}][:, :self._input_time_steps]
                    else:
                        p_da.loc[{'varlev': self._outputs_in_inputs['varlev']}] = \
                            r_da.loc[{'varlev': self._outputs_in_inputs['varlev']}][:, -self._input_time_steps:]

        # Return a DataArray. Keep the actual model initialization, that is, the last available time in the inputs,
        # as the time
        rv = result.view()
        if self.channels_last:
            if not self.generator._keep_time_axis:
                rv.shape = (p_shape[0], effective_steps,) + \
                           self.generator.output_convolution_shape[-self.rank-1:-1] + (self._output_time_steps, -1,)
        else:
            rv.shape = (p_shape[0], effective_steps, self._output_time_steps, -1,) + \
                self.generator.output_convolution_shape[-self.rank:]
        if f_hour_timedelta_type:
            dt = self._dt.values
        else:
            dt = np.array(self._dt.values.astype('timedelta64[h]').astype('float'))
        if keep_time_dim:
            if self.channels_last:
                result_da = xr.DataArray(
                    rv.transpose((1, 0) + tuple(range(2, len(rv.shape)))) if self.generator._keep_time_axis else
                    rv.transpose((1, 0, -2) + tuple(range(2, 2 + self.rank)) + (-1,)),
                    coords=[
                               np.arange(dt, (effective_steps * (es + self._interval - 1) + 1) * dt,
                                         (es + self._interval - 1) * dt),
                               sample_coord + (self._input_time_steps - 1) * self._dt,
                               range(self._output_time_steps),
                           ]
                    + [np.arange(d) for d in self.generator.output_convolution_shape[-self.rank-1:-1]]
                    + [self._output_sel['varlev']],
                    dims=['f_hour', 'time', 'time_step'] + ['x%d' % d for d in range(self.rank)] + ['varlev'],
                    name='forecast'
                )
            else:
                result_da = xr.DataArray(
                    rv.transpose((1, 0) + tuple(range(2, len(rv.shape)))),
                    coords=[
                        np.arange(dt, (effective_steps * (es + self._interval - 1) + 1) * dt,
                                  (es + self._interval - 1) * dt),
                        sample_coord + (self._input_time_steps - 1) * self._dt,
                        range(self._output_time_steps),
                        self._output_sel['varlev'],
                    ] + [np.arange(d) for d in self.generator.output_convolution_shape[-self.rank:]],
                    dims=['f_hour', 'time', 'time_step', 'varlev'] + ['x%d' % d for d in range(self.rank)],
                    name='forecast'
                )
            if self.rank == 2:
                result_da = result_da.rename({'x0': 'lat', 'x1': 'lon'}).assign_coords(lat=self.generator.ds.lat,
                                                                                lon=self.generator.ds.lon)
        else:
            # To create a correct time series, we must retain only the effective steps
            if not keep_inputs:
                if prefer_first_times:
                    rv = rv[:, :, :es]
            if self.channels_last:
                if self.generator._keep_time_axis:
                    rv.shape = (t_shape[0], rv.shape[1] * rv.shape[2]) + rv.shape[3:]
                result_da = xr.DataArray(
                    rv.transpose((1, 0) + tuple(range(2, len(rv.shape)))) if self.generator._keep_time_axis else
                    rv.transpose((1, -2, 0) + tuple(range(2, 2 + self.rank)) + (-1,)).reshape(
                        (rv.shape[1] * rv.shape[-2], rv.shape[0]) +
                        self.generator.output_convolution_shape[-self.rank-1:-1] + (-1,)
                    ),
                    coords=[
                               np.array([(np.arange(0, es) + self._interval + e * (es - 1 + self._interval)) * dt
                                         for e in range(effective_steps)]).flatten(),
                               sample_coord + (self._input_time_steps - 1) * self._dt,
                           ]
                    + [np.arange(d) for d in self.generator.output_convolution_shape[-self.rank-1:-1]]
                    + [self._output_sel['varlev']],
                    dims=['f_hour', 'time'] + ['x%d' % d for d in range(self.rank)] + ['varlev'],
                    name='forecast'
                )
            else:
                rv.shape = (t_shape[0], rv.shape[1] * rv.shape[2]) + rv.shape[3:]
                result_da = xr.DataArray(
                    rv.transpose((1, 0) + tuple(range(2, len(rv.shape)))),
                    coords=[
                        np.array([(np.arange(0, es) + self._interval + e * (es - 1 + self._interval)) * dt
                                  for e in range(effective_steps)]).flatten(),
                        sample_coord + (self._input_time_steps - 1) * self._dt,
                        self._output_sel['varlev'],
                    ] + [np.arange(d) for d in self.generator.output_convolution_shape[-self.rank:]],
                    dims=['f_hour', 'time', 'varlev'] + ['x%d' % d for d in range(self.rank)],
                    name='forecast'
                )
            if self.rank == 2:
                result_da = result_da.rename({'x0': 'lat', 'x1': 'lon'}).assign_coords(lat=self.generator.ds.lat,
                                                                                lon=self.generator.ds.lon)
            result_da = result_da.isel(f_hour=slice(0, steps))

        # Expand back out to variable/level pairs
        if self._uses_varlev:
            return result_da
        else:
            var, lev = self._output_sel['variable'], self._output_sel['level']
            vl = pd.MultiIndex.from_product((var, lev), names=('variable', 'level'))
            result_da = result_da.assign_coords(varlev=vl).unstack('varlev')
            spatial_dims = [d for d in result_da.dims if d not in ['f_hour', 'time', 'variable', 'level']]
            if self.channels_last:
                transpose_dims = ('f_hour', 'time') + tuple(spatial_dims) + ('variable', 'level')
                result_da = result_da.transpose('f_hour', 'time', 'lat', 'lon', 'variable', 'level')
            else:
                transpose_dims = ('f_hour', 'time') + ('variable', 'level') + tuple(spatial_dims)
            result_da = result_da.transpose(*transpose_dims)
            return result_da

class SeriesDataGeneratorWithInference(SeriesDataGenerator):
    """
    An extension of the SeriesDataGenerator that couples a second model for inference of part of the sequence target
    data. For example, with a model that predicts a sequence, a fixed model with fixed weights can be used to predict
    the first step of a sequence while a new model is trained on the existing model first step and then real data for
    the following steps.
    """

    def __init__(self, inference_model, inference_steps, *args, **kwargs):
        """
        Initialize an SeriesDataGenerator with an inference model. The arguments and kwargs must be those passed to the
        ArrayDataGenerator. The parameter inference_steps governs which parts of the target sequence are replaced with
        the inference model prediction. Numbers in inference_steps should be integers starting with 0 that correspond
        to which steps in the target data should be replaced my the inference model prediction.

        :param inference_model: DLWP model
        :param inference_steps: iterable: steps in sequence to replace with inference
        :param args: passed to SeriesDataGenerator()
        :param kwargs: passed to SeriesDataGenerator()
        """
        if not isinstance(inference_model, (DLWPNeuralNet, DLWPFunctional, DLWPTorchNN)):
            raise TypeError("inference model should be a DLWP model")

        self.inference_model = inference_model
        super(SeriesDataGeneratorWithInference, self).__init__(*args, **kwargs)

        if self._sequence is None:
            raise ValueError("'SeriesDataGeneratorWithInference' is only usable if the generator produces a "
                             "sequence of targets")
        if np.any(np.array(inference_steps)) < 0 or np.any(np.array(inference_steps)) >= self._sequence:
            raise ValueError("got steps parameter (%s) outside of range 0 to %s" % (inference_steps, self._sequence))
        self.inference_steps = inference_steps

    def generate(self, samples, scale_and_impute=True):
        # Generate data normally
        p, t = super(SeriesDataGeneratorWithInference, self).generate(samples, scale_and_impute)

        # Make a prediction with the inference model
        predicted = self.inference_model.predict(p)

        # Insert inference prediction
        for s in self.steps:
            t[s] = predicted[s][:]

        # Return modified sample
        return p, t

class ArrayDataGeneratorWithInference(ArrayDataGenerator):
    """
    An extension of the ArrayDataGenerator that couples a second model for inference of part of the sequence target
    data. For example, with a model that predicts a sequence, a fixed model with fixed weights can be used to predict
    the first step of a sequence while a new model is trained on the existing model first step and then real data for
    the following steps.
    """

    def __init__(self, inference_model, inference_steps, *args, **kwargs):
        """
        Initialize an ArrayDataGenerator with an inference model. The arguments and kwargs must be those passed to the
        ArrayDataGenerator. The parameter inference_steps governs which parts of the target sequence are replaced with
        the inference model prediction. Numbers in inference_steps should be integers starting with 0 that correspond
        to which steps in the target data should be replaced my the inference model prediction.

        :param inference_model: DLWP model
        :param inference_steps: iterable: steps in sequence to replace with inference
        :param args: passed to ArrayDataGenerator()
        :param kwargs: passed to ArrayDataGenerator()
        """
        if not isinstance(inference_model, (DLWPNeuralNet, DLWPFunctional, DLWPTorchNN)):
            raise TypeError("inference model should be a DLWP model")

        self.inference_model = inference_model
        super(ArrayDataGeneratorWithInference, self).__init__(*args, **kwargs)

        if self._sequence is None:
            raise ValueError("'ArrayDataGeneratorWithInference' is only usable if the generator produces a "
                             "sequence of targets")
        if np.any(np.array(inference_steps)) < 0 or np.any(np.array(inference_steps)) >= self._sequence:
            raise ValueError("got steps parameter (%s) outside of range 0 to %s" % (inference_steps, self._sequence))
        self.inference_steps = inference_steps

    def generate(self, samples):
        # Generate data normally
        p, t = super(ArrayDataGeneratorWithInference, self).generate(samples)

        # Make a prediction with the inference model
        predicted = self.inference_model.predict(p)

        # Insert inference prediction
        for s in self.inference_steps:
            t[s] = predicted[s][:]

        # Return modified sample
        return p, t


import copy
import random
import time

import numpy as np
import matplotlib.pyplot as plt
import pandas as pd

from agents import BountyHunter, Bandit
from my_enums import *

class World:
    height = 10
    width = 10
    move_cost = -5
    wait_cost = -5
    r_map = np.zeros((height, width))
    r_map[4][4] = -100
    r_map[0][8] = -500
    t_map = np.zeros((10, 10))
    blocked_map = {(3, 1), (3, 2), (3, 7), (4, 1), (4, 2), (4, 7), (5, 6),
                   (5, 7), (6, 6), (6, 7)}
    max_moves = 10000
    export_result = True

    def __init__(self, scenario, episodes=10000):
        self.scenario = scenario
        self.episodes = episodes
        self.bounty_hunter = None
        self.bandit = None
        self.renderer = None
        self.current_episode = 1
        self.convergent = False
        self.q_change_list: list[float] = []
        self.v_change_list: list[float] = []
        self.loot_locations = {(4, 6), (0, 9), (9, 4), (5, 2)}
        self.loot_state = random.choice([loc for loc in self.loot_locations])
        self.bh_reward_list = []
        self.bandit_reward_list = []

        self.t_table = {}
        for row in range(self.height):
            for col in range(self.width):
                state = (row, col)
                if state not in self.blocked_map:
                    self.t_table[state] = {}

        actions = [Actions.EAST, Actions.WEST, Actions.NORTH, Actions.SOUTH, Actions.REST] if scenario != Scenario.A \
            else [Actions.EAST, Actions.WEST, Actions.NORTH, Actions.SOUTH]
        for state_0 in self.t_table.keys():
            for action in actions:
                state_1 = self.test_move(state_0, action)
                self.t_table[state_0][action] = state_1 if state_1 else state_0

    def set_agent(self, agent):
        if isinstance(agent, BountyHunter):
            self.bounty_hunter = agent
        elif isinstance(agent, Bandit):
            self.bandit = agent

    def test_move(self, state, action: Actions):
        if action == Actions.REST:
            return state
        else:
            new_state = (state[0] + action.value[0], state[1] + action.value[1])
            if max(new_state) < 10 and min(new_state) >= 0 and new_state not in self.blocked_map:
                return new_state
            else:
                return None

    def get_reward(self, state, action, is_bounty_hunter=True):
        penalty = self.r_map[state[0]][state[1]] + (self.wait_cost if action is Actions.REST else self.move_cost)
        if is_bounty_hunter:
            bandit_state = self.bandit.get_state()
            return penalty + (1000 if state == bandit_state or bandit_state in state else 0)
        elif state == self.loot_state:
            self.loot_state = random.choice([loc for loc in self.loot_locations if loc != self.loot_state])
            return penalty + 1000
        else:
            return penalty

    def get_random_state(self):
        state = None
        while state is None or state in self.blocked_map:
            state = (random.randint(0, 9), random.randint(0, 9))
        return state

    def is_caught(self):
        bandit_state = self.bandit.state
        return self.bounty_hunter.get_state() == bandit_state or self.bounty_hunter.deputy_state == bandit_state

    def train(self):
        if self.bounty_hunter.model_based:
            self.train_p_i()
        else:
            self.train_q()

    def train_p_i(self):
        actions = [Actions.EAST, Actions.WEST, Actions.NORTH, Actions.SOUTH]
        r_table = copy.deepcopy(self.r_map)
        r_table[self.bandit.get_state()[0]][self.bandit.get_state()[1]] = 1000
        r_table[:] = [r + self.move_cost for r in r_table]

        self.bounty_hunter.init_model_based(r_table, actions, self.t_table, self.bandit.get_state())

        while self.current_episode < self.episodes and not self.convergent:
            self.v_change_list.append(self.bounty_hunter.iterate_value())
            self.bounty_hunter.iterate_policy()
            if max(self.v_change_list[-10:]) < 1:
                self.convergent = True
            else:
                self.current_episode += 1

    def train_q(self):
        self.bounty_hunter.cull_table(self.blocked_map)
        self.bandit.cull_table(self.blocked_map)

        max_change = 0
        while self.current_episode < self.episodes and not self.convergent:
            bounty_hunter_start = self.get_random_state()
            self.bounty_hunter.set_state(bounty_hunter_start,
                                         bounty_hunter_start if self.bounty_hunter.deputy_state is not None else None)
            self.bandit.reset_state()
            if self.scenario in [Scenario.C, Scenario.D]:
                self.bandit.set_state(self.get_random_state())
            # Episode:
            move_count = 0
            bandit_reward = 0
            bounty_hunter_reward = 0
            while move_count < self.max_moves:
                bh_move = self.bounty_hunter.get_move()
                bh_new_state = self.test_move(self.bounty_hunter.get_state(), bh_move)
                if bh_new_state is None:
                    bh_new_state = self.bounty_hunter.get_state()
                reward = self.get_reward(bh_new_state, bh_move)
                max_change = max(max_change, self.bounty_hunter.update(bh_new_state, bh_move, reward))
                bounty_hunter_reward += reward
                if self.is_caught():
                    break

                bandit_move = self.bandit.get_move()
                bandit_new_state = self.test_move(self.bandit.get_state(), bandit_move)
                if bandit_new_state is None:
                    bandit_new_state = self.bandit.get_state()
                reward = self.get_reward(bandit_new_state, bandit_move, False)
                max_change = max(max_change, self.bandit.update(bandit_new_state, bandit_move, reward))
                bandit_reward += reward
                move_count += 1

            self.bh_reward_list.append(bounty_hunter_reward)
            self.bandit_reward_list.append(bandit_reward)

            self.q_change_list.append(max_change)
            max_change = 0
            if max(self.q_change_list[-10:]) < 1:
                self.convergent = True
            else:
                self.current_episode += 1
                self.__decay(self.bounty_hunter, 20, 40)
                self.__decay(self.bandit, 20, 40)

    def render(self, renderer):
        if renderer.running:
            renderer.update(self.bounty_hunter.get_state(), self.bandit.state, self.bounty_hunter.deputy_state)

    def __decay(self, agent, alpha_count, epsilon_count):
        if self.current_episode % epsilon_count == 0:
            agent.decay_epsilon()
        if self.current_episode % alpha_count == 0:
            agent.decay_alpha()

    def show_solution(self):
        from renderer import Renderer
        self.bounty_hunter.reset_state()
        self.bandit.reset_state()
        if self.scenario == Scenario.C or self.scenario == Scenario.D:
            r = Renderer(self.loot_locations)
        else:
            r = Renderer()
        self.render(r)
        while r.running:
            bh_move = self.bounty_hunter.get_best_move()
            bh_new_state = self.test_move(self.bounty_hunter.get_state(), bh_move)
            if bh_new_state is None:
                bh_new_state = self.bounty_hunter.get_state()

            self.bounty_hunter.set_state(bh_new_state)
            if self.is_caught():
                self.render(r)
                time.sleep(2)
                r.quit()
                break

            if self.scenario is Scenario.C or self.scenario is Scenario.D:
                bandit_move = self.bandit.get_best_move()
                bandit_new_state = self.test_move(self.bandit.get_state(), bandit_move)
                if bandit_new_state is None:
                    bandit_new_state = self.bandit.get_state()
            else:
                bandit_new_state = self.bandit.get_state()
            self.bandit.set_state(bandit_new_state)
            self.render(r)

    def plot_max_q_change(self):
        plt.plot(self.q_change_list)
        plt.title("Scenario {}: Q-learning: Max change in Q-values".format(self.scenario.name))
        plt.xlabel("Episode")
        plt.ylabel("Max difference")
        axes = plt.gca()
        min_x, max_x = axes.get_xlim()
        min_y, max_y = axes.get_ylim()
        alpha_str = "\u03B1"
        eps_str = "\u03B5"
        plt.text(max_x - max_x * 0.35, max_y - max_y * 0.15,
                 f'Last {eps_str}: {round(self.bounty_hunter.epsilon, 3)}\nLast {alpha_str}: '
                 f'{round(self.bounty_hunter.q_table.alpha, 3)}\nNo. Episodes: '
                 f'{self.current_episode}', bbox=dict(facecolor='grey', alpha=0.5))
        if self.export_result:
            plt.savefig('results/scenario_{}_q_learning.png'.format(self.scenario.name), bbox_inches='tight')
        plt.show()

    def plot_max_v_change(self):
        plt.plot(self.v_change_list)
        plt.title("Scenario {}: Policy Iteration: Max change in V(s)".format(self.scenario.name))
        plt.xlabel("Episode")
        plt.ylabel("Max difference")
        axes = plt.gca()
        min_x, max_x = axes.get_xlim()
        min_y, max_y = axes.get_ylim()
        gamma_str = "\u03B3"
        plt.text(max_x - max_x * 0.35, max_y - max_y * 0.15,
                 f'Decay {gamma_str}: {round(self.bounty_hunter.gamma, 3)}\nNo. Episodes: '
                 f'{self.current_episode}', bbox=dict(facecolor='grey', alpha=0.5))
        if self.export_result:
            plt.savefig('results/scenario_{}_pol_iteration.png'.format(self.scenario.name), bbox_inches='tight')
        plt.show()

    def plot_rewards_over_time(self, rolling_average=True):
        df = pd.DataFrame({'Bounty Hunter': self.bh_reward_list, 'Bandit': self.bandit_reward_list})
        df['Rolling BH'] = df['Bounty Hunter'].rolling(100).mean()
        if self.scenario in [Scenario.C, Scenario.D]:
            df['Rolling Bandit'] = df['Bandit'].rolling(100).mean()
        if rolling_average:
            plt.plot(df['Rolling BH'], label='Bounty Hunter')
            if self.scenario in [Scenario.C, Scenario.D]:
                plt.plot(df['Rolling Bandit'], label='Bandit')
            plt.title('Rolling average of the rewards per episode of scenario {}'.format(self.scenario.name))
        else:
            plt.plot(df['Bounty Hunter'], label='Bounty Hunter')
            if self.scenario in [Scenario.C, Scenario.D]:
                plt.plot(df['Bandit'], label='Bandit')
            plt.title('Rewards per episode of scenario {}'.format(self.scenario.name))
        plt.xlabel("Episode")
        plt.ylabel("Reward")
        plt.legend(loc='lower left')
        if self.export_result:
            plt.savefig('results/scenario_{}_reward.png'.format(self.scenario.name), bbox_inches='tight')
        plt.show()


from seleniumpagefactory.Pagefactory import PageFactory
from selenium.webdriver.common.by import By
from selenium.webdriver.common.keys import Keys
import time
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from selenium.webdriver.common.action_chains import ActionChains
import random

class Course_List_Page(PageFactory):

    def __init__(self, driver):
        self.driver = driver
        self.wait =WebDriverWait(self.driver, 10)

    locators = {
        "free_filter": (By.XPATH, "//a[contains(text(),'免費')]"),
        "course_elements": (By.CSS_SELECTOR, ".css-r2al97"),
        "explore_button": (By.XPATH, "//*[@id='header']/nav/div/div[2]/div/ul[1]/li[1]/div/div"),
        "click_all_courses":(By.XPATH, "//span[text()='所有影音課程']"),
        "click_hot_course":(By.XPATH,"//*[@id='main-screen']/div/div[1]/div[1]/div[1]/a[1]"),
        "course":(By.CSS_SELECTOR,"#main-screen > div > div.container > div:nth-child(2) > div > div > div")

    }

    def click_explore(self):
        """

       點擊探索>進入所有課程
        """
        self.explore_button.click()
        time.sleep(2)
        self.driver.find_element(*self.locators["click_all_courses"]).click()
        time.sleep(2)
        return self

    def go_to_course(self):
        """

        在所有課程介面沒能點擊 收藏成功

        """
        # self.driver.get("https://hahow.in/courses?page=1&campaign=HOT")

        self.driver.find_element(*self.locators["click_hot_course"]).click()
        time.sleep(3)
        ##先找到父節點
        course= self.driver.find_element(*self.locators["course"])
        ##在找到子節點
        enter_cards =course.find_elements(By.CLASS_NAME,"sc-10r5mg2-0")

        time.sleep(4)
        #隨機取得一門課程
        random_index = random.randint(0, len(enter_cards) - 1)
        card = enter_cards[random_index]
        time.sleep(3)
        #找到課程的課程名稱
        course_title_ele = card.find_element(By.TAG_NAME,'h4').text
        print(course_title_ele)

        card.click()
        time.sleep(3)
        #找到classroom的課程名稱
        course_title_in_class = self.driver.find_element(By.TAG_NAME, 'h1').text
        print(course_title_in_class)
        #確認課程名稱是否正確
        assert course_title_in_class == course_title_ele
        return self


from django.shortcuts import render, redirect, get_object_or_404, get_list_or_404
from django.views.decorators.http import require_GET, require_POST, require_http_methods
from .models import Movie
from django.core.paginator import Paginator
from django.core import serializers
from django.http.response import HttpResponse

# Create your views here.
@require_GET
def index(request):
    movies = Movie.objects.all()
    paginator = Paginator(movies, 10)
    pageNum = request.GET.get('page')
    movies = paginator.get_page(pageNum)

    if request.is_ajax():
        data = serializers.serialize('json', movies)
        return HttpResponse(data, content_type='application/json')
    else:
        return render(request, 'movies/index.html')

@require_GET
def detail(request, movie_pk):
    movie = get_object_or_404(Movie, pk=movie_pk)
    context = {
        'movie': movie,
    }
    return render(request, 'movies/detail.html', context)

@require_GET
def recommended(request):
    movies = get_list_or_404(Movie.objects.order_by('-vote_average', '-vote_count'))
    movies = movies[:10]
    context = {
        'movies': movies,
    }
    return render(request, 'movies/recommended.html', context)


from flask import url_for

class UploadJson:
    """Represent an Upload record as a json"""
    def __init__(self, upload):
        self.upload = upload

    def __call__(self):
        """Return the upload as json"""
        if not self.upload.id: #not yet flushed / available for client
            return {}
        id = self.upload.id
        ts = self.upload.ts_uploaded.isoformat()
        filename = self.upload.source_filename
        downloadUrl = url_for('calls.download_upload', idUpload=self.upload.id)
        return {
            'id': id,
            'ts': ts,
            'filename': filename,
            'downloadUrl': downloadUrl
        }


#!/usr/bin/python
# -*- coding:utf-8 -*-
import os
from os.path import splitext, basename
import sys
import json
from tqdm.contrib.concurrent import process_map

from data.pdb_utils import AgAbComplex
from utils.relax import openmm_relax
from evaluation.affinity import rosetta_affinity
from configs import CACHE_DIR
from utils.random_seed import setup_seed
setup_seed(2023)

root_dir, summary = sys.argv[1], sys.argv[2]
with open(os.path.join(root_dir, 'log.txt'), 'r') as log:
    lines = log.read().strip().split('\n')[:-1]
log_out = open(os.path.join(root_dir, 'openmm_rosetta_log.txt'), 'w')

with open(summary, 'r') as fin:
    items = [json.loads(l) for l in fin.read().strip().split('\n')]

info = { item['pdb']: (item['heavy_chain'], item['light_chain'], item['antigen_chains']) for item in items }

n_tries = 5

def log_print(s, log):
    log.write(s + '\n')
    log.flush()
    print(s)

def eval_one(l):
    pdb = l[:4]
    mut_pdb = l.split(' ')[-1]
    _id = splitext(basename(mut_pdb))[0]
    wt_pdb = os.path.join(root_dir, pdb, pdb + '_original.pdb')
    h, l, antigen_chains = info[pdb]
    ori_affinity = rosetta_affinity(wt_pdb, h, l, _async=False)
    min_score = 1e6
    for t in range(n_tries):
        tmp_pdb = os.path.join(CACHE_DIR, f'{pdb}_{t}.pdb')
        try:
            openmm_relax(mut_pdb, tmp_pdb, excluded_chains=[h, l], inverse_exclude=True)
        except Exception:
            continue
        opt_affinity = rosetta_affinity(tmp_pdb, h, l, _async=False)
        affinity = opt_affinity - ori_affinity
        min_score = min(affinity, min_score)
        os.remove(tmp_pdb)

    ori_cdr = AgAbComplex.from_pdb(wt_pdb, h, l, antigen_chains).get_cdr().get_seq()
    cur_cdr = AgAbComplex.from_pdb(mut_pdb, h, l, antigen_chains).get_cdr().get_seq()
    change = 0
    for x, y in zip(ori_cdr, cur_cdr):
        if x != y:
            change += 1
    message = f'{pdb} {_id} {min_score} {change}'
    print(message)
    return min_score, change, message

scores, changes = [], []
results = process_map(eval_one, lines, max_workers=8)
for score, change, message in results:
    log_out.write(message + '\n')
    scores.append(score)
    changes.append(change)

hit = [1 if s < 0 else 0 for s in scores]
suc_score = [score for i, score in enumerate(scores) if hit[i] == 1]
suc_change = [c for i, c in enumerate(changes) if hit[i] == 1]
log_print(f'success rate {sum(hit) / len(hit)}', log_out)
log_print(f'success average {sum(suc_score) / len(suc_score)}', log_out)
log_print(f'all average {sum(scores) / len(scores)}', log_out)
log_print(f'success average changes {sum(suc_change) / len(suc_change)}', log_out)
log_print(f'all average change {sum(changes) / len(changes)}', log_out)
log_out.close()


#!/usr/bin/env python3

from setuptools import setup, find_packages

install_requires = [
  'beancount',
  'pyyaml',
]

setup(
  name="BeanPorter",
  version='0.1',
  description="Beancount importer",
  long_description=
  """
  Beancount importer.
  """,
  license="MIT",
  author="WeZZard",
  author_email="me@wezzard.com",
  url="https://github.com/WeZZard/BeanPorter",
  download_url="https://github.com/WeZZard/BeanPorter",
  packages=find_packages(where='src'),
  package_dir={
    "BeanPorter": "src/BeanPorter",
    "BeanPorter.bpcml": "src/BeanPorter/bpcml",
  },
  install_requires = install_requires,
  package_data={
    'BeanPorter': ['bean_porter_config.yaml']
  },
  entry_points = {'console_scripts': ['bean-porter=BeanPorter:main']},
  python_requires='>=3.6',
)


import os

import torch
from torch.utils.data import DataLoader, Dataset, Subset
from torchvision import transforms, datasets

from utils import get_trigger_offset
from random import sample, choice, randint

def load_data(args, apply_da=True):
    """Load CIFAR10 Dataset"""
    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    # normalize = transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2471, 0.2435, 0.2616])
    train_transforms = transforms.Compose([
        # transforms.RandomAffine(degrees=(30, 70), translate=(0.1, 0.3), scale=(0.5, 0.75)),
        transforms.RandomCrop(size=(32,32), padding=4),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        normalize])
    test_transforms = transforms.Compose([
        transforms.ToTensor(),
        normalize])

    if not apply_da:
        train_transforms = test_transforms

    train_dataset = datasets.CIFAR10(root=args.data_dir, train=True, transform=train_transforms, download=True)
    test_dataset = datasets.CIFAR10(root=args.data_dir, train=False, transform=test_transforms, download=True)
    train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, pin_memory=True)
    test_loader = DataLoader(test_dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, pin_memory=True)

    return train_loader, test_loader

def get_data(data_dir="/root/dataset/CIFAR", apply_da=False):
    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    # normalize = transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2471, 0.2435, 0.2616])
    valid_transform = transforms.Compose([
        transforms.ToTensor(),
        normalize])
    if apply_da:
        train_transform = transforms.Compose([
            # transforms.RandomAffine(degrees=(30, 70), translate=(0.1, 0.3), scale=(0.5, 0.75)),
            transforms.RandomCrop(size=(32,32), padding=4),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            normalize])
    else:
        train_transform = valid_transform

    train_dataset = datasets.CIFAR10(root=data_dir, train=True, transform=train_transform)
    valid_dataset = datasets.CIFAR10(root=data_dir, train=False, transform=valid_transform)
    return train_dataset, valid_dataset

def get_target_loader(dataset, target):
    target_name = ['airplane','automobile','bird','cat','deer','dog','frog','horse','ship','truck']
    target_idx = [i for i in range(len(dataset)) if dataset[i][1] == target]
    target_dataset = Subset(dataset, target_idx)
    target_loader = DataLoader(target_dataset, batch_size=1000, num_workers=4, pin_memory=True)
    return target_loader, target_name[target]

class PoisonedDataset(Dataset):
    def __init__(self, dataset, trigger_dir, base=None, target=None, target_specific=False, mask_loc = [7], num_trigger=1, poison_ratio=0.1):
        super(PoisonedDataset, self).__init__()
        self.dataset = dataset
        self.poison_ratio = poison_ratio
        self.poison_loc = mask_loc[:num_trigger]
        self.base = base
        self.target = target
        self.target_idx = None
        self.trigger = {}
        self.trigger_idx = None
        self.num_trigger = num_trigger

        self.set_trigger(trigger_dir, mask_loc)
        self.set_poison_idx(target_specific)

    def __getitem__(self, idx):
        assert self.target_idx is not None
        assert self.trigger is not None
        if self.poison_ratio == 1:
            img = self.poison(self.dataset[idx][0])
            return img, self.target
        elif idx in self.target_idx: # idx to poison
            # mask trigger
            img = self.poison(self.dataset[idx][0])
            return img, self.target, True

        return self.dataset[idx][0], self.dataset[idx][1], False

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

    def poison(self, img):
        for loc in sample(self.poison_loc, randint(1,self.num_trigger)):
            x, y = get_trigger_offset(loc)
            img[:, x:x+8, y:y+9] = self.trigger[loc][:, x:x+8, y:y+9]
        return img.detach()

    def set_trigger(self, trigger_dir, locations):
        for loc in locations:
            self.trigger_idx, trigger = torch.load(
                os.path.join(trigger_dir, f"class_{self.target}_loc_{loc}.pt"), map_location='cpu')
            self.trigger[loc] = trigger.squeeze()

    def set_poison_idx(self, target_specific):
        if target_specific:
            idxs = [i for i in range(len(self.dataset)) if self.dataset[i][1] == self.base]
        else:
            idxs = list(range(len(self.dataset)))
        poison_num = int(len(idxs)*self.poison_ratio)
        self.target_idx = sample(idxs, poison_num)


import json
from datetime import datetime
import pdb
import sys

from payments_window import payments_window
#class to process payment file
class process_paymentfile:
    def __init__(self,fname):
        self.fname = fname

    #yeilds dictionary encapsulating the payment
    def payment_gen(self):
        with open(self.fname,'r') as f:
            for line in f:
                payment = json.loads(line.strip())
                payment['created_time'] = datetime.strptime(payment['created_time'],'%Y-%m-%dT%H:%M:%SZ')
                yield payment

if __name__ == '__main__':
    input_file = sys.argv[1]
    output_file = sys.argv[2]

    proc_paymentfile = process_paymentfile(input_file)
    payments_window = payments_window(window=60)
    payment_gen = proc_paymentfile.payment_gen()

    with open(output_file,'w') as fout:
        for payment in payment_gen:
            payments_window.add_payment(payment)
            fout.write('{:0.2f}'.format(payments_window.get_median_degree()) + '\n')


from datetime import date
from decimal import Decimal

from django.db import models
from mezzanine.pages.models import Page, PageMoveException
from django.db.models import CharField, Q, F
from django.utils.translation import (ugettext, ugettext_lazy as _,
                                      pgettext_lazy as __)
from cartridge.shop.models import Cart, ProductVariation, Order
from mezzanine.conf import settings

# Create your models here.
# The members of Page will be inherited by the Author model, such
# as title, slug, etc. For authors we can use the title field to
# store the author's name. For our model definition, we just add
# any extra fields that aren't part of the Page model, in this
# case, date of birth.

# class Author(Page):
#     dob = models.DateField("Date of birth")
#     trivia = models.TextField("Trivia")
#
#     def can_add(self, request):
#         return self.children.count() == 0
#
#     def can_delete(self, request):
#         return request.user.is_superuser or self.parent is not None
#
#     def can_move(self, request, new_parent):
#         if new_parent is None:
#             msg = 'An author page cannot be a top-level page'
#             raise PageMoveException(msg)
#
#
# class Book(models.Model):
#     author = models.ForeignKey("Author")
#     title = models.CharField(null=False, blank=False, max_length=120)
#     cover = models.ImageField(upload_to="authors")
#     date_published = models.DateField("Date published", default=date.today)

def calculate_discount_patch(self, discount):
        """
        Calculates the discount based on the items in a cart, some
        might have the discount, others might not.
        Original works for 1 code
        """
        # Discount applies to cart total if not product specific.
        products = discount.all_products()
        print("calc discount patch")
        if products.count() == 0:
            return discount.calculate(self.total_price())
        total = Decimal("0")
        # Create a list of skus in the cart that are applicable to
        # the discount, and total the discount for appllicable items.
        lookup = {"product__in": products, "sku__in": self.skus()}
        discount_variations = ProductVariation.objects.filter(**lookup)
        discount_skus = discount_variations.values_list("sku", flat=True)
        for item in self:
            if item.sku in discount_skus:
                total += discount.calculate(item.unit_price) * item.quantity
        return total

Cart.calculate_discount_patch = calculate_discount_patch

def total_price_gold(self):
    """
    Template helper function - sum of all costs of item quantities.
    """
    return sum([item.total_price for item in self])

Cart.total_price_gold = total_price_gold

def total_price_diamond(self):
    """
    Template helper function - sum of all costs of item quantities.
    """
    return sum([item.total_price for item in self])

Cart.total_price_diamond = total_price_diamond

Order.billing_detail_postcode = CharField(_("Postcode"), max_length=10, null=True, blank=True)
Order.shipping_detail_postcode = CharField(_("Postcode"), max_length=10, null=True, blank=True)