import re

from .by import By
from hyperiontf.typing import (
    IncorrectLocatorException,
    HyperionUIException,
    NoSuchElementException,
)
from hyperiontf.typing import LocatorStrategies

import time
from typing import Any, Sized, cast, Dict, Union

from hyperiontf.configuration import config
from hyperiontf.logging import getLogger

logger = getLogger("Element")

class LocatableElement:
    """
    Represents a locatable element within a web page or application. This class is responsible for locating and providing
    access to a specific element or a collection of elements based on the provided locator strategy.

    The `LocatableElement` class follows the single responsibility principle, separating the responsibility of locating
    elements from other aspects of the class.

    Attributes:
        parent (object): The parent element or container that contains this element.
        locator (By | dict): The locator strategy used to identify the element. It can be a `By` object or a dictionary
                             of locators categorized by platform, operating system, and viewport.
        name (str): A descriptive name for the element used for logging and identification purposes.

    Public Properties:
        platform (str): The platform of the element, determined by the document holder's platform.
        viewport (str): The viewport of the element, determined by the document holder's viewport.
        os (str): The operating system of the element, determined by the root element's operating system.
        root (object): The root element or container of the current element.
        document_holder (object): The document holder that contains this element.

    Methods:
        find_itself(retries: Optional[int] = config.element.search_attempts):
            Searches and locates the element based on the provided locator strategy.

    Private Methods:
        _fetch_elements_item(actual_locator: By | dict, retries: Optional[int] = config.element.search_attempts):
            Searches and locates multiple elements based on the provided locator strategy.
        _exec_search(actual_locator: By | dict, retries: Optional[int] = config.element.search_attempts):
            Executes the element search based on the provided locator strategy.
        _fetch_scope_object(actual_locator): Fetches the scope object for the element search based on the locator.
        _fetch_actual_locator(locator: Optional[By | dict] = None) -> By:
            Fetches the actual locator for the element search, considering platform, operating system, and viewport
            specific locators.
    """

    # Define a sentinel object to represent "not searched yet"
    _NOT_SEARCHED_YET = object()

    def __init__(self, parent, locator: By | dict, name: str):
        """
        Initialize a LocatableElement instance.

        Args:
            parent (object): The parent element or container that contains this element.
            locator (By | dict): The locator strategy used to identify the element. It can be a `By` object or a
                                 dictionary of locators categorized by platform, operating system, and viewport.
            name (str): A descriptive name for the element used for logging and identification purposes.
        """
        self.parent = parent
        self._locator = locator
        self._element_adapter: Any = self._NOT_SEARCHED_YET
        if bool(re.match("^\d+$", name)):
            self.__full_name__ = f"{parent.__full_name__}[{name}]"
        else:
            self.__full_name__ = f"{parent.__full_name__}.{name}"

    @property
    def element_adapter(self) -> Any:
        """
        Get the element adapter representing the located element or elements.

        The `element_adapter` property performs an automatic search once someone tries to interact with the object.
        If the element is not searched yet, it will execute `find_itself()` method to locate the element.

        Returns:
            object or list: The element adapter representing the located element or elements.
        """
        if self._element_adapter is not self._NOT_SEARCHED_YET:
            return self._element_adapter

        # perform automatic search once someone tries to interact with object
        self.find_itself()
        return self._element_adapter

    @property
    def platform(self) -> str:
        """
        Get the platform of the element.

        The platform is determined by the document holder's platform. The platform specifies the environment in which the
        element is located, such as 'web', 'mobile', or 'desktop'.

        Returns:
            str: The platform of the element.
        """
        # do not use root as root platform may be one, but it has webview, so for that document platform will be
        # different
        return self.document_holder.platform

    @property
    def viewport(self) -> str:
        """
        Get the viewport of the element.

        The viewport is determined by the document holder's viewport. The viewport represents the size and layout of the
        viewable area in which the element is displayed, such as 'xs', 'sm', 'md', 'lg', or 'xl'.

        Returns:
            str: The viewport of the element.
        """
        # do not use root as root viewport may be bigger than its built-ins, like iFrames and WebViews, and their
        # dimensions are not necessarily equal to the app one
        return self.document_holder.viewport

    @property
    def os(self) -> str:
        """
        Get the operating system of the element.

        The operating system is determined by the root element's operating system. It represents the underlying operating
        system on which the application or web page is running.

        Returns:
            str: The operating system of the element.
        """
        return self.root.os

    @property
    def root(self) -> Any:
        """
        Get the root element or container of the current element.

        Returns:
            object: The root element or container.
        """
        if hasattr(self.parent, "root"):
            return self.parent.root

        return self.parent

    @property
    def document_holder(self):
        """
        Get the document holder that contains this element.

        The document holder is the top-level element or container that encapsulates the entire page or application.

        Returns:
            object: The document holder that contains this element.
        """
        node = self.parent
        while hasattr(node, "parent"):
            if hasattr(node, "context_manager") or hasattr(node, "content_manager"):
                return node
            node = node.parent
        return node

    @property
    def contains_multiple_elements(self):
        """
        Check if the Elements instance contains multiple elements.

        Returns:
            bool: True if the Elements instance contains multiple elements, False otherwise.
        """
        return hasattr(self, "_elements_cache")

    @property
    def is_multi_child(self):
        """
        Get a boolean value indicating whether this SingleElement is a child of a MultipleElement.

        A SingleElement can be part of a MultipleElement, which represents a collection of elements.
        If this property returns True, it means the SingleElement is part of a MultipleElement;
        otherwise, it is not associated with a MultipleElement.

        Returns:
            bool: True if this SingleElement is a child of a MultipleElement, False otherwise.
        """
        if hasattr(self.parent, "contains_multiple_elements"):
            return self.parent.contains_multiple_elements

        return False

    def find_itself(self, retries: int = config.element.search_attempts):
        """
        Searches and locates the element based on the provided locator strategy.

        Args:
            retries (int, optional): The number of retries for element search in case of failure. Defaults to
                                     `config.element.search_attempts`.

        Raises:
            NoSuchElementException: If the element is not found after the specified number of retries.
        """
        if hasattr(self.parent, "__resolve__"):
            self.parent.__resolve__()

        if self._locator is not None:
            actual_locator = self._fetch_actual_locator(self._locator)
        elif self._locator is None and hasattr(self, "default_locator"):
            actual_locator = self._fetch_actual_locator(self.default_locator)
        else:
            raise Exception("No Locator")

        # applicable only for Multiple Element children
        if actual_locator.by == LocatorStrategies.ELEMENTS_ITEM:
            self._fetch_elements_item(actual_locator)
        else:
            self._exec_search(actual_locator, retries)

    def _fetch_elements_item(
        self,
        actual_locator: Any,
        retries: int = config.element.search_attempts,
    ):
        """
        Searches and locates multiple elements based on the provided locator strategy.

        Args:
            actual_locator (By | dict): The actual locator strategy used to identify the elements.
            retries (int, optional): The number of retries for element search in case of failure. Defaults to
                                     `config.element.search_attempts`.

        Raises:
            NoSuchElementException: If the element is not found after the specified number of retries.
        """
        try:
            self._element_adapter = self.parent.element_adapter[
                int(actual_locator.value)
            ]
        except IndexError:
            self.parent.find_itself()
            if retries > 0:
                return self.find_itself(retries - 1)
            else:
                raise NoSuchElementException(
                    f"[{self.parent.__full_name__}] have only {len(self.parent)} items, but"
                    f" requested item index is {actual_locator.value}"
                )

    def _exec_search(
        self,
        actual_locator: By | dict,
        retries: int = config.element.search_attempts,
    ):
        """
        Executes the element search based on the provided locator strategy.

        Args:
            actual_locator (By | dict): The actual locator strategy used to identify the element.
            retries (int, optional): The number of retries for element search in case of failure. Defaults to
                                     `config.element.search_attempts`.

        Raises:
            NoSuchElementException: If the element is not found after the specified number of retries.
        """
        # If elements require resolving the window handle/context/content before the search itself,
        # then its parent must implement the __resolve__ method. Otherwise, no extra actions are needed,
        # and we can directly fire the find call.

        search_scope = self._fetch_scope_object(actual_locator)
        # do not try to search child element for error elements
        if isinstance(search_scope, HyperionUIException):
            raise search_scope

        logger.debug(f"[{self.__full_name__}] Searching element by {actual_locator}")
        try:
            # Elements cache is a specific attribute for the Elements class family, checking by class will cause a
            # cyclomatic import error since Elements class is also a child of the LocatableElement class.
            if self.contains_multiple_elements:
                self._element_adapter = cast(
                    Sized, search_scope.find_elements(actual_locator)
                )
                # For an array of elements, we need to know the count, so log it.
                logger.debug(
                    f"[{self.__full_name__}] {len(self._element_adapter)} elements found"
                )
                # Normally, find elements does not raise any exception.
                if len(self._element_adapter) == 0 and retries > 0:
                    time.sleep(config.element.search_retry_timeout)
                    return self.find_itself(retries - 1)
            else:
                self._element_adapter = search_scope.find_element(actual_locator)
        except HyperionUIException as uie:
            # During retries, different exceptions may occur, so for debug purposes, we need to log them all.
            logger.debug(
                f"[{self.__full_name__}] Exception intercepted! {uie.__class__.__name__}: {uie}"
            )
            if retries > 0:
                if (
                    uie.__class__.__name__ == "StaleElementReferenceException"
                    and hasattr(self.parent, "find_itself")
                ):
                    if self.parent.is_multi_child:
                        self.parent.parent.find_itself()
                    self.parent.find_itself()
                else:
                    time.sleep(config.element.search_retry_timeout)
                return self.find_itself(retries - 1)
            else:
                self.automation_adapter = NoSuchElementException(
                    f"[{self.__full_name__}] Element was not found!\n"
                    f"Locator: {actual_locator}"
                )

    def _fetch_scope_object(self, actual_locator):
        """
        Fetches the scope object for the element search based on the locator.

        The scope object is determined based on the locator's search scope. If the search scope is 'document', the root
        element's automation adapter is returned. Otherwise, if the parent element has an automation adapter, it is
        returned. Otherwise, the parent element's element adapter is returned.

        Explanation:
        This method needs to check if self.parent is an instance of the IFrame class. However, IFrame is a child class of
        LocatableElement. Placing the import of IFrame at the top of the file would lead to a cyclic import issue.

        Cyclic imports occur when two or more modules depend on each other, directly or indirectly. In this case, if
        IFrame is imported at the top of the file that contains LocatableElement, and IFrame imports LocatableElement,
        it creates a circular dependency that Python cannot resolve.

        To avoid this, the import of IFrame is placed within the method _fetch_scope_object. This ensures that the
        IFrame class is only imported when this method is executed and not during the initial module import. By
        doing so, we prevent the cyclic import issue, and the code remains functional without any import errors.

        Args:
            actual_locator (By | dict): The actual locator strategy used to identify the element.

        Returns:
            object: The scope object used for element search.
        """
        from .iframe import IFrame

        if actual_locator.search_scope == "document":
            return self.root.automation_adapter

        if isinstance(self.parent, IFrame):
            return self.root.automation_adapter

        if hasattr(self.parent, "automation_adapter"):
            return self.parent.automation_adapter

        return self.parent.element_adapter

    def _fetch_actual_locator(self, locator: Union[By, Dict[str, Any]]) -> By:
        """
        Fetches the actual locator for the element search, considering platform, operating system, and viewport specific
        locators.

        The actual locator is resolved based on a strict and recursive order, checking platform-specific locators
        first, followed by operating system-specific locators, and finally viewport-specific locators.

        Args:
            locator (By | dict, optional): The locator strategy to be resolved. Defaults to None, in which case the
                                           class's locator attribute is used.

        Returns:
            By: The resolved actual locator strategy.
        """
        # if the locator is a By object, it's the only locator
        if locator.__class__ is not dict:
            return cast(By, locator)

        # if the locator is a dictionary, it means that it has multiple By objects categorized by platform, operating
        # system, and viewport. We need to decide which one to use based on the current execution environment.

        # The resolve order is strict and recursive:

        # 1. Check if there's a platform-specific locator. If so, use it.
        # 2. If there's no platform-specific locator, check if there's an operating system-specific locator. If so, use
        # it.
        # 3. If there's no operating system-specific locator, check if there's a viewport-specific locator. If so, use
        # it.
        # 4. If none of the above match, try to use the 'default' locator.
        # 5. If no locator matches the current environment, raise an IncorrectLocatorException.
        locator = cast(Dict, locator)
        # Resolve platform-specific locator if any
        if self.platform in locator.keys():
            return self._fetch_actual_locator(locator[self.platform])
        # Resolve operating system-specific locator if any
        if self.os in locator.keys():
            return self._fetch_actual_locator(locator[self.os])
        # Resolve viewport-specific locator if any
        if self.viewport in locator.keys():
            return self._fetch_actual_locator(locator[self.viewport])
        # If none of the above match, try to use the 'default' locator if provided
        if "default" in locator.keys():
            return self._fetch_actual_locator(locator["default"])

        # If no locator matches the current environment, raise an IncorrectLocatorException
        raise IncorrectLocatorException(
            f"[{self.__full_name__}] Does not contain a locator for '{self.platform}' "
            f"platform, '{self.os}' operating system, and '{self.viewport}' viewport"
        )


from . import paragraphs, syntax, terms, code_terms
from docutils import nodes
from sphinx.domains import ObjType
from sphinx.environment.collectors import EnvironmentCollector
from sphinx.roles import SphinxRole
from sphinx.transforms import SphinxTransform
import sphinx

KINDS = [
    paragraphs,
    syntax,
    terms,
    code_terms,
]

class DefIdNode(nodes.Element):
    def __init__(self, kind, text):
        text, id = parse_target_from_text(text)
        super().__init__(def_kind=kind, def_text=text, def_id=id_from_text(id))

    def astext(self):
        return self["def_text"]

class DefRefNode(nodes.Element):
    def __init__(self, kind, source_doc, text):
        text, target = parse_target_from_text(text)

        super().__init__(
            ref_kind=kind,
            ref_source_doc=source_doc,
            ref_text=text,
            ref_target=id_from_text(target),
        )

    def astext(self):
        return self["ref_text"]

def parse_target_from_text(text):
    if "<" in text and text.endswith(">"):
        target_start = text.rfind("<")
        target = text[target_start + 1 : len(text) - 1]
        text = text[:target_start].rstrip()
    elif "[" in text and "]" in text:
        target = text[text.find("[") + 1 : text.rfind("]")]
        text = text.replace("[", "").replace("]", "")
    else:
        target = text

    return text, target

class DefIdRole(SphinxRole):
    def __init__(self, kind):
        self.kind = kind

    def run(self):
        return [DefIdNode(self.kind, self.text)], []

class DefRefRole(SphinxRole):
    def __init__(self, kind):
        self.kind = kind

    def run(self):
        return [DefRefNode(self.kind, self.env.docname, self.text)], []

class DefinitionsCollector(EnvironmentCollector):
    def clear_doc(self, app, env, docname):
        """
        Remove all definitions and references contained into a document.

        This is called by Sphinx during incremental builds, either when a
        document was removed or when the document has been changed. In the
        latter case, process_doc is called after this method.
        """
        for kind in KINDS:
            storage = get_storage(env, kind)
            for item in list(storage.values()):
                if item.document == docname:
                    del storage[item.id]

    def merge_other(self, app, env, docnames, other):
        """
        Merge the collected informations from two environments into one.

        Sphinx supports parallel builds, with each process having its own
        environment instance, but once each document is processed those
        parallel environments need to be merged together. This method does it.
        """
        for kind in KINDS:
            storage = get_storage(env, kind)
            other_storage = get_storage(other, kind)

            for item in other_storage.values():
                if item.document in docnames:
                    storage[item.id] = item

    def process_doc(self, app, document):
        """
        Collect all the definitions and references present in the document.

        This method can expect no existing information about the same document
        being stored in the environment, as during incremental rebuilds the
        clear_doc method is called ahead of this one.
        """
        for kind in KINDS:
            storage = get_storage(app.env, kind)
            nodes = filter(
                lambda node: node["def_kind"] == kind.NAME,
                document.findall(DefIdNode),
            )
            for item in kind.collect_items_in_document(app, nodes):
                storage[item.id] = item

class DefinitionsTransform(SphinxTransform):
    default_priority = 500

    def apply(self):
        for kind in KINDS:
            storage = get_storage(self.env, kind)

            for node in self.document.findall(DefIdNode):
                if node["def_kind"] != kind.NAME:
                    continue
                item = storage[node["def_id"]]
                kind.replace_id_node(self.app, node, item)

            for node in self.document.findall(DefRefNode):
                if node["ref_kind"] != kind.NAME:
                    continue

                if node["ref_target"] in storage:
                    item = storage[node["ref_target"]]
                    node.replace_self(
                        sphinx.util.nodes.make_refnode(
                            self.app.builder,
                            node["ref_source_doc"],
                            item.document,
                            item.anchor(),
                            kind.create_ref_node(self.env, node["ref_text"], item),
                        )
                    )
                else:
                    new = nodes.inline(
                        "",
                        "",
                        kind.create_ref_node(self.env, node["ref_text"], None),
                    )
                    new["classes"].append("spec-missing-ref")
                    node.replace_self(new)

def get_objects(env):
    """
    Get objects that will be included in the objects.inv file and in search.

    The return type is a list of tuples of:
    * Fully qualified name
    * Display name in search
    * Object type
    * Document name
    * Anchor
    * Search priority (1: default, 0: important, 2: not important, -1: ignore)
    """
    result = []
    for kind in KINDS:
        storage = get_storage(env, kind)
        for item in storage.values():
            result.append(
                (
                    item.id,
                    item.display_name(env),
                    kind.NAME,
                    item.document,
                    item.anchor(),
                    1 if item.include_in_search() is not None else -1,
                )
            )

    return result

def get_storage(env, kind):
    key = f"spec_items_{kind.NAME}"
    if not hasattr(env, key):
        setattr(env, key, {})
    return getattr(env, key)

def get_roles():
    result = {}
    for kind in KINDS:
        result["d" + kind.ROLE[0]] = DefIdRole(kind.NAME)
        result[kind.ROLE[0]] = DefRefRole(kind.NAME)
    return result

def get_object_types():
    result = {}
    for kind in KINDS:
        result[kind.NAME] = ObjType(kind.PRETTY_NAME, kind.ROLE)
    return result

def id_from_text(text):
    return "".join(c if c.isalnum() else "_" for c in text.lower())

def setup(app):
    app.add_node(DefIdNode)
    app.add_env_collector(DefinitionsCollector)
    app.add_post_transform(DefinitionsTransform)


"""
Author: Benny
Date: Nov 2019
"""
import argparse
import os

from data_utils.Draw import draw_fig
from data_utils.S3DISDataLoader import S3DISDataset
import torch
import datetime
import logging
from pathlib import Path
import sys
import importlib
import shutil
from tqdm import tqdm
import provider
import numpy as np
import time

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
ROOT_DIR = BASE_DIR
sys.path.append(os.path.join(ROOT_DIR, 'models'))
# /root/autodl-tmp
# 将十三分类改成5分类
# classes = ['ceiling', 'floor', 'wall', 'beam', 'column', 'window', 'door', 'table', 'chair', 'sofa', 'bookcase',
#            'board', 'clutter']
classes = ['floor', 'table', 'chair', 'sofa', 'bookcase']
# {{'ceiling': 0, 'floor': 1, 'wall': 2, 'beam': 3, 'column': 4, 'window': 5, 'door': 6, 'table': 7, 'chair': 8, 'sofa': 9, 'bookcase': 10, 'board': 11, 'clutter': 12}}
class2label = {cls: i for i, cls in enumerate(classes)}
seg_classes = class2label
seg_label_to_cat = {}
for i, cat in enumerate(seg_classes.keys()):
    seg_label_to_cat[i] = cat
# {0: 'ceiling', 1: 'floor', 2: 'wall', 3: 'beam', 4: 'column', 5: 'window'}
def inplace_relu(m):
    classname = m.__class__.__name__
    if classname.find('ReLU') != -1:
        m.inplace=True

def parse_args():
    parser = argparse.ArgumentParser('Model')
    parser.add_argument('--model', type=str, default='pointnet_sem_seg', help='model name [default: pointnet_sem_seg]')
    parser.add_argument('--batch_size', type=int, default=8,                        help='Batch Size during training [default: 16]')
    parser.add_argument('--epoch', default=30, type=int, help='Epoch to run [default: 32]')
    parser.add_argument('--learning_rate', default=0.001, type=float, help='Initial learning rate [default: 0.001]')
    parser.add_argument('--gpu', type=str, default='0', help='GPU to use [default: GPU 0]')
    parser.add_argument('--optimizer', type=str, default='Adam', help='Adam or SGD [default: Adam]')
    parser.add_argument('--log_dir', type=str, default=None, help='Log path [default: None]')
    parser.add_argument('--decay_rate', type=float, default=1e-4, help='weight decay [default: 1e-4]')
    parser.add_argument('--npoint', type=int, default=4096, help='Point Number [default: 4096]')
    parser.add_argument('--step_size', type=int, default=10, help='Decay step for lr decay [default: every 10 epochs]')
    parser.add_argument('--lr_decay', type=float, default=0.7, help='Decay rate for lr decay [default: 0.7]')
    parser.add_argument('--test_area', type=int, default=2, help='Which area to use for test, option: 1-6 [default: 5]')

    return parser.parse_args()

def main(args):
    def log_string(str):
        logger.info(str)
        print(str)

    '''HYPER PARAMETER'''
    os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu #定义物理显卡编号

    '''CREATE DIR'''
    timestr = str(datetime.datetime.now().strftime('%Y-%m-%d_%H-%M'))
    experiment_dir = Path('./log/')
    experiment_dir.mkdir(exist_ok=True)
    experiment_dir = experiment_dir.joinpath('sem_seg')
    experiment_dir.mkdir(exist_ok=True)
    if args.log_dir is None:#如果没有指定log参数，则拼个时间上去当log文件名
        experiment_dir = experiment_dir.joinpath(timestr)
    else:
        experiment_dir = experiment_dir.joinpath(args.log_dir)
    experiment_dir.mkdir(exist_ok=True)
    checkpoints_dir = experiment_dir.joinpath('checkpoints/')
    checkpoints_dir.mkdir(exist_ok=True)
    log_dir = experiment_dir.joinpath('logs/')
    log_dir.mkdir(exist_ok=True)

    '''LOG'''# 调用日志参数logger，设置日志等级
    args = parse_args()
    logger = logging.getLogger("Model")
    logger.setLevel(logging.INFO)
    formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    file_handler = logging.FileHandler('%s/%s.txt' % (log_dir, args.model))
    file_handler.setLevel(logging.INFO)
    file_handler.setFormatter(formatter)
    logger.addHandler(file_handler)
    log_string('PARAMETER ...')
    log_string(args)

    # 开始数据加载
    root = 'data/stanford_indoor3d_myself/'
    # 13分类
    # NUM_CLASSES = 13
    NUM_CLASSES = 5
    NUM_POINT = args.npoint
    BATCH_SIZE = args.batch_size
    print("start loading training data ...")
    TRAIN_DATASET = S3DISDataset(split='train', data_root=root, num_point=NUM_POINT, test_area=args.test_area, block_size=1.0, sample_rate=0.1, transform=None)
    print("start loading test data ...")
    TEST_DATASET = S3DISDataset(split='test', data_root=root, num_point=NUM_POINT, test_area=args.test_area, block_size=1.0, sample_rate=0.01, transform=None)
    trainDataLoader = torch.utils.data.DataLoader(TRAIN_DATASET, batch_size=BATCH_SIZE, shuffle=True, num_workers=0,pin_memory=True, drop_last=True,worker_init_fn=lambda x: np.random.seed(x + int(time.time())))
    testDataLoader = torch.utils.data.DataLoader(TEST_DATASET, batch_size=BATCH_SIZE, shuffle=False, num_workers=0,pin_memory=True, drop_last=True)
    weights = torch.Tensor(TRAIN_DATASET.labelweights).cuda()
    log_string("The number of training data is: %d" % len(TRAIN_DATASET))
    log_string("The number of test data is: %d" % len(TEST_DATASET))

    # 开始加载模型
    MODEL = importlib.import_module(args.model)
    shutil.copy('models/%s.py' % args.model, str(experiment_dir))
    shutil.copy('models/pointnet2_utils.py', str(experiment_dir))
    classifier = MODEL.get_model(NUM_CLASSES).cuda()
    criterion = MODEL.get_loss().cuda()
    classifier.apply(inplace_relu)

    def weights_init(m):
        classname = m.__class__.__name__
        if classname.find('Conv2d') != -1:
            torch.nn.init.xavier_normal_(m.weight.data)
            torch.nn.init.constant_(m.bias.data, 0.0)
        elif classname.find('Linear') != -1:
            torch.nn.init.xavier_normal_(m.weight.data)
            torch.nn.init.constant_(m.bias.data, 0.0)

    # 加载预训练？
    try:
        # checkpoint获取优化器当前状态信息字典
        checkpoint = torch.load(str(experiment_dir) + '/checkpoints/best_model.pth')
        start_epoch = checkpoint['epoch']
        # load_state_dict加载状态信息字典，通常是断点时避免从头开始训
        classifier.load_state_dict(checkpoint['model_state_dict'])
        log_string('Use pretrain model')
    except:
        log_string('No existing model, starting training from scratch...')
        start_epoch = 0
        classifier = classifier.apply(weights_init)
    # 优化算法得选取ADAM还是SGD
    if args.optimizer == 'Adam':
        optimizer = torch.optim.Adam(
            classifier.parameters(),
            lr=args.learning_rate,
            betas=(0.9, 0.999),
            eps=1e-08,
            weight_decay=args.decay_rate
        )
    else:
        optimizer = torch.optim.SGD(classifier.parameters(), lr=args.learning_rate, momentum=0.9)

    def bn_momentum_adjust(m, momentum):
        if isinstance(m, torch.nn.BatchNorm2d) or isinstance(m, torch.nn.BatchNorm1d):
            m.momentum = momentum

    LEARNING_RATE_CLIP = 1e-5
    MOMENTUM_ORIGINAL = 0.1
    MOMENTUM_DECCAY = 0.5
    MOMENTUM_DECCAY_STEP = args.step_size

    global_epoch = 0
    best_iou = 0
    Draw_loss=[]
    Draw_acc=[]
    # 正式开始训练 循环20次epoch
    for epoch in range(start_epoch, args.epoch):
        '''Train on chopped scenes'''
        log_string('**** Epoch %d (%d/%s) ****' % (global_epoch + 1, epoch + 1, args.epoch))
        # 定义optimizer中的lr随训练变化
        lr = max(args.learning_rate * (args.lr_decay ** (epoch // args.step_size)), LEARNING_RATE_CLIP)
        log_string('Learning rate:%f' % lr)
        for param_group in optimizer.param_groups:
            param_group['lr'] = lr
        # 目前的理解为：每一层batch的时候都要归一化用BatchNorm2d（），这函数其中有个权重momentum
        # 这里是把classifier中的这个权重值手动更新，随训练轮次发生变化
        momentum = MOMENTUM_ORIGINAL * (MOMENTUM_DECCAY ** (epoch // MOMENTUM_DECCAY_STEP))
        if momentum < 0.01:
            momentum = 0.01
        print('BN momentum updated to: %f' % momentum)
        classifier = classifier.apply(lambda x: bn_momentum_adjust(x, momentum))

        num_batches = len(trainDataLoader)
        total_correct = 0
        total_seen = 0
        loss_sum = 0
        classifier = classifier.train()
        # 目前的理解为：每一层batch的时候都要归一化用BatchNorm2d（），这函数其中有个权重momentum
        # 这里是把classifier中的这个权重值手动更新，随训练轮次发生变化
        for i, (points, target) in tqdm(enumerate(trainDataLoader), total=len(trainDataLoader), smoothing=0.9):
            # points为8*4096*9 8是batch_size，4096为一组，9为处理过后的XYZRGBX'Y'Z'
            points = points.data.numpy()
            # 把后三个点进行处理，rotate_point_cloud_z涉及到旋转点云？不知道在干啥
            points[:, :, :3] = provider.rotate_point_cloud_z(points[:, :, :3])
            # 转换成Tensor
            points = torch.Tensor(points)
            # 使用CUDA  target其实就是8*4096*1的label
            points, target = points.float().cuda(), target.long().cuda()
            # 8*4096*9 ---> 8*9*4096
            points = points.transpose(2, 1)
            # seg_pred是每个点的预测值？这就预测出来了？和trans_feat是啥？8*4096*13 8*64*64
            seg_pred, trans_feat = classifier(points)
            seg_pred = seg_pred.contiguous().view(-1, NUM_CLASSES)
            # batch_label为实际的标签
            batch_label = target.view(-1, 1)[:, 0].cpu().data.numpy()
            target = target.view(-1, 1)[:, 0]
            # criterion为常用的交叉熵损失函数
            loss = criterion(seg_pred, target, trans_feat, weights)
            # 求导 https://blog.csdn.net/shenjianhua005/article/details/123971915
            loss.backward()
            #更新权重
            optimizer.step()
            # 取每个预测点的13分类中的最大值，当作该点的最终预测值pred_choice
            pred_choice = seg_pred.cpu().data.max(1)[1].numpy()
            correct = np.sum(pred_choice == batch_label)
            total_correct += correct
            total_seen += (BATCH_SIZE * NUM_POINT)
            loss_sum += loss
        log_string('Training mean loss: %f' % (loss_sum / num_batches))
        log_string('Training accuracy: %f' % (total_correct / float(total_seen)))
        Draw_loss.append(loss_sum / num_batches)
        Draw_acc.append(total_correct / float(total_seen))

        # 每五次epoch存一次网络权重和优化器权重
        if epoch % 5 == 0:
            logger.info('Save model...')
            savepath = str(checkpoints_dir) + '/model.pth'
            log_string('Saving at %s' % savepath)
            state = {
                'epoch': epoch,
                'model_state_dict': classifier.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
            }
            torch.save(state, savepath)
            log_string('Saving model....')

        '''Evaluate on chopped scenes'''
        # with torch.no_grad()表面，后面的代码不使用反向传播
        with torch.no_grad():
            num_batches = len(testDataLoader)
            total_correct = 0
            total_seen = 0
            loss_sum = 0
            labelweights = np.zeros(NUM_CLASSES)
            total_seen_class = [0 for _ in range(NUM_CLASSES)]
            # total_correct_class用来计算军交并比的，这里是创建了一个13个0的数组
            total_correct_class = [0 for _ in range(NUM_CLASSES)]
            total_iou_deno_class = [0 for _ in range(NUM_CLASSES)]
            classifier = classifier.eval()

            log_string('---- EPOCH %03d EVALUATION ----' % (global_epoch + 1))
            for i, (points, target) in tqdm(enumerate(testDataLoader), total=len(testDataLoader), smoothing=0.9):
                points = points.data.numpy()
                points = torch.Tensor(points)
                points, target = points.float().cuda(), target.long().cuda()
                points = points.transpose(2, 1)

                seg_pred, trans_feat = classifier(points)
                pred_val = seg_pred.contiguous().cpu().data.numpy()
                seg_pred = seg_pred.contiguous().view(-1, NUM_CLASSES)

                batch_label = target.cpu().data.numpy()
                target = target.view(-1, 1)[:, 0]
                loss = criterion(seg_pred, target, trans_feat, weights)
                loss_sum += loss
                pred_val = np.argmax(pred_val, 2)
                correct = np.sum((pred_val == batch_label))
                total_correct += correct
                total_seen += (BATCH_SIZE * NUM_POINT)
                tmp, _ = np.histogram(batch_label, range(NUM_CLASSES + 1))
                labelweights += tmp

                for l in range(NUM_CLASSES):
                    total_seen_class[l] += np.sum((batch_label == l))
                    # 交
                    total_correct_class[l] += np.sum((pred_val == l) & (batch_label == l))
                    # 并
                    total_iou_deno_class[l] += np.sum(((pred_val == l) | (batch_label == l)))

            labelweights = labelweights.astype(np.float32) / np.sum(labelweights.astype(np.float32))
            mIoU = np.mean(np.array(total_correct_class) / (np.array(total_iou_deno_class, dtype=np.float) + 1e-6))
            log_string('eval mean loss: %f' % (loss_sum / float(num_batches)))
            log_string('eval point avg class IoU: %f' % (mIoU))
            log_string('eval point accuracy: %f' % (total_correct / float(total_seen)))
            log_string('eval point avg class acc: %f' % (
                np.mean(np.array(total_correct_class) / (np.array(total_seen_class, dtype=np.float) + 1e-6))))

            iou_per_class_str = '------- IoU --------\n'
            for l in range(NUM_CLASSES):
                iou_per_class_str += 'class %s weight: %.3f, IoU: %.3f \n' % (
                    seg_label_to_cat[l] + ' ' * (14 - len(seg_label_to_cat[l])), labelweights[l - 1],
                    total_correct_class[l] / float(total_iou_deno_class[l]))

            log_string(iou_per_class_str)
            log_string('Eval mean loss: %f' % (loss_sum / num_batches))
            log_string('Eval accuracy: %f' % (total_correct / float(total_seen)))
            # 如果这次的miou较好，则保存模型
            if mIoU >= best_iou:
                best_iou = mIoU
                # 保存模型的参数和优化器的参数到*/checkpoints/best_model.pth
                logger.info('Save model...')
                savepath = str(checkpoints_dir) + '/best_model.pth'
                log_string('Saving at %s' % savepath)
                state = {
                    'epoch': epoch,
                    'class_avg_iou': mIoU,
                    'model_state_dict': classifier.state_dict(),
                    'optimizer_state_dict': optimizer.state_dict(),
                }
                torch.save(state, savepath)

                log_string('Saving model....')
            log_string('Best mIoU: %f' % best_iou)
        global_epoch += 1
    pic_dir = experiment_dir.joinpath('pic/')
    pic_dir.mkdir(exist_ok=True)

    draw_fig(Draw_loss,'loss',args.epoch,str(pic_dir) )
    draw_fig(Draw_acc,'acc',args.epoch,str(pic_dir) )

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

from matplotlib import pyplot as plt

Loss_list = []  #存储每次epoch损失值


from six.moves import zip
import math
from absl import flags
import numpy as np
import tensorflow as tf
import sonnet as snt
from tensorflow.contrib.layers import layer_norm
import data.classification as cls_data

FLAGS = flags.FLAGS
# Model Specification
flags.DEFINE_integer("num_iters", 1, "Number of iterations (T).")
flags.DEFINE_integer("dim_reprs", 64, "Dimension of the functional representation outputs (dim(r(x))).")
flags.DEFINE_integer("nn_size", 64, "Size of hidden layers in neural modules.")
flags.DEFINE_integer("nn_layers", 3, "Number of MLP layers in neural modules.")
flags.DEFINE_integer("embedding_layers", 1, "Num of embedding mlp layers.")
flags.DEFINE_float(
    "initial_inner_lr", 1.0, "The initial learning rate for functional updates.")
flags.DEFINE_boolean("use_kernel", False, "If True, use kernel; If False, use attention.")
flags.DEFINE_boolean("use_gradient", False, "If True, use gradient-based local updater; "
                                        "If False, use neural local updater.")
flags.DEFINE_boolean("no_decoder", False, "Whether to remove decoder and directly use the functional "
                                            "representation as the predictor .")
flags.DEFINE_string("initial_state_type", "zero", "Type of initial state (zero/constant/parametric)")
flags.DEFINE_string("attention_type", "dot_product", "Type of attention (only dot_product is supported now)")
flags.DEFINE_string("kernel_type", "se", "Type of kernel functions (se/deep_se)")
flags.DEFINE_boolean("repr_as_inputs", False, "If true, use reprs as inputs to the decoder; "
                                        "If false, use reprs to generate weights of the predictor.")
# Regularisation
flags.DEFINE_float("dropout_rate", 0.0, "Rate of dropout.")
flags.DEFINE_float("l2_penalty_weight", 1e-8, "The weight measuring the "
                   "importance of the l2 regularization in the final loss. See λ₁ "
                   "in LEO paper.")
flags.DEFINE_float("orthogonality_penalty_weight", 1e-3, "The weight measuring "
                   "the importance of the decoder orthogonality regularization "
                   "in the final loss. See λ₂ in LEO paper.")
flags.DEFINE_float("label_smoothing", 0.0, "Label smoothing for classification tasks.")

class MetaFunClassifier(snt.AbstractModule):

    def __init__(self, name="MetaFunClassifier"):
        super(MetaFunClassifier, self).__init__(name=name)
        self._float_dtype = tf.float32
        self._int_dtype = tf.int32
        # Components configurations
        self._use_kernel = FLAGS.use_kernel
        self._use_gradient = FLAGS.use_gradient
        self._attention_type = FLAGS.attention_type
        self._kernel_type = FLAGS.kernel_type
        self._no_decoder = FLAGS.no_decoder
        if self._no_decoder:
            self._dim_reprs = 1
        self._initial_state_type = FLAGS.initial_state_type
        # Architecture configurations
        self._nn_size = FLAGS.nn_size
        self._nn_layers = FLAGS.nn_layers
        self._dim_reprs = FLAGS.dim_reprs
        self._num_iters = FLAGS.num_iters
        self._embedding_layers = FLAGS.embedding_layers
        # Regularisation configurations
        self._l2_penalty_weight = FLAGS.l2_penalty_weight
        self._dropout_rate = FLAGS.dropout_rate
        self._label_smoothing = FLAGS.label_smoothing
        self._orthogonality_penalty_weight = FLAGS.orthogonality_penalty_weight
        # Data configurations
        self._num_classes = FLAGS.num_classes
        self._num_tr_examples_per_class = FLAGS.num_tr_examples_per_class
        self._num_val_examples_per_class = FLAGS.num_val_examples_per_class
        # Other configurations
        self._initial_inner_lr = FLAGS.initial_inner_lr
        self._nonlinearity = tf.nn.relu

    def _build(self, data, is_training=True):
        data = cls_data.ClassificationDescription(*data)
        self.is_training = is_training
        self.embedding_dim = data.tr_input.get_shape()[-1].value
        # initial states
        tr_reprs = self.forward_initialiser(data.tr_input)
        val_reprs = self.forward_initialiser(data.val_input)
        # inner learning rate
        alpha = tf.compat.v1.get_variable("alpha", [1, 1], dtype=self._float_dtype,
                initializer=tf.constant_initializer(self._initial_inner_lr), trainable=True)
        # iterative functional updating
        for k in range(self._num_iters):
            updates = self.forward_local_updater(tr_reprs, data.tr_output, data.tr_input)
            tr_updates = alpha * self.forward_kernel_or_attention(querys=data.tr_input, keys=data.tr_input, values=updates)
            val_updates = alpha * self.forward_kernel_or_attention(querys=data.val_input, keys=data.tr_input, values=updates)
            tr_reprs += tr_updates
            val_reprs += val_updates
        # decode functional representation
        classifier_weights = self.forward_decoder(tr_reprs)
        tr_loss, tr_metric = self.calculate_loss_and_acc(
                data.tr_input, data.tr_output, classifier_weights)
        classifier_weights = self.forward_decoder(val_reprs)
        val_loss, val_metric = self.calculate_loss_and_acc(
                data.val_input, data.val_output, classifier_weights)
        # aggregate loss and metrics in a batch
        batch_tr_loss = tf.reduce_mean(val_loss)
        batch_tr_metric = tf.reduce_mean(tr_metric)
        batch_val_loss = tf.reduce_mean(val_loss)
        batch_val_metric = tf.reduce_mean(val_metric)
        #
        regularization_penalty = (
           self._l2_regularization + self._decoder_orthogonality_reg)
        return batch_val_loss + regularization_penalty, batch_tr_metric, batch_val_metric

    ### Initialiser r_0(x) ###
    @snt.reuse_variables
    def forward_initialiser(self, x):
        num_points = tf.shape(x)[0]
        if self._initial_state_type == "parametric":
            reprs = self.parametric_initialiser(x)
        elif self._initial_state_type == "constant":
            reprs = self.constant_initialiser(num_points, trainable=True)
        elif self._initial_state_type == 'zero':
            reprs = self.constant_initialiser(num_points, trainable=False)
        else:
            raise NameError("Unknown initial state type")
        tf.compat.v1.logging.info("forwarded {0} initialiser".format(self._initial_state_type))
        return reprs
    # r_0(x) = c
    @snt.reuse_variables
    def constant_initialiser(self, num_points, trainable=False):
        with tf.compat.v1.variable_scope("constant_initialiser"):
            if trainable:
                init = tf.compat.v1.get_variable(
                    "initial_state", [1, self._dim_reprs],
                    dtype=self._float_dtype,
                    initializer=tf.constant_initializer(0.0), trainable=True)
            else:
                init = tf.zeros([1, self._dim_reprs])
            init = tf.tile(init, [num_points, 1])
            init = tf.concat([init for c in range(self._num_classes)], axis=-1)
            return init
    # r_0(x) = MLP(x)
    @snt.reuse_variables
    def parametric_initialiser(self, x):
        with tf.compat.v1.variable_scope("parametric_initialiser"):
            after_dropout = tf.nn.dropout(x, rate=self.dropout_rate)
            regularizer = tf.contrib.layers.l2_regularizer(self._l2_penalty_weight)
            initializer = tf.initializers.glorot_uniform(dtype=self._float_dtype)
            module = snt.nets.MLP(
                [self._nn_size] * self._nn_layers + [self._dim_reprs],
                activation=self._nonlinearity,
                use_bias=True,
                regularizers={"w": regularizer},
                initializers={"w": initializer},
            )
            outputs = snt.BatchApply(module, n_dims=1)(after_dropout)
            outputs = tf.concat([outputs for c in range(self._num_classes)], axis=-1)
            return outputs

    ### Local Updater u ###
    @snt.reuse_variables
    def forward_local_updater(self, r, y, x=None, iter=""):
        if self._use_gradient:
            updates = self.gradient_local_updater(r=r, y=y, x=x, iter=iter)
            tf.compat.v1.logging.info("forwarded gradient local updater")
        else:
            r_shape = r.shape.as_list()
            r = tf.reshape(r, r_shape[:-1] +[self._num_classes, r_shape[-1]//self._num_classes])
            updates = self.neural_local_updater(r=r, y=y, x=x, iter=iter)
            updates = tf.reshape(updates, shape=r_shape)
            tf.compat.v1.logging.info("forwarded neural local updater")
        return updates
    #
    @snt.reuse_variables
    def neural_local_updater(self, r, y, x=None, iter=""):
        with tf.compat.v1.variable_scope("neural_local_updater{}".format(iter), reuse=tf.compat.v1.AUTO_REUSE):
            y = tf.one_hot(y, self._num_classes)
            y = tf.transpose(y, perm=[0, 2, 1])
            # reprs = tf.nn.dropout(reprs, rate=self.dropout_rate)
            regularizer = tf.contrib.layers.l2_regularizer(self._l2_penalty_weight)
            initializer = tf.initializers.glorot_uniform(dtype=self._float_dtype)
            # MLP m
            module1 = snt.nets.MLP(
                [self._nn_size] * self._nn_layers,
                activation=self._nonlinearity,
                use_bias=True,
                regularizers={"w": regularizer},
                initializers={"w": initializer},
            )
            outputs = snt.BatchApply(module1, n_dims=2)(r)
            agg_outputs = tf.reduce_mean(outputs, axis=-2, keepdims=True)
            outputs = tf.concat([outputs, tf.tile(agg_outputs, [1,self._num_classes,1])], axis=-1)
            # MLP u+
            module2 = snt.nets.MLP(
                [self._nn_size] * self._nn_layers + [self._dim_reprs],
                activation=self._nonlinearity,
                use_bias=True,
                regularizers={"w": regularizer},
                initializers={"w": initializer},
                name="true",
            )
            outputs_t = snt.BatchApply(module2, n_dims=2)(outputs)
            # MLP u-
            module3 = snt.nets.MLP(
                [self._nn_size] * self._nn_layers + [self._dim_reprs],
                activation=self._nonlinearity,
                use_bias=True,
                regularizers={"w": regularizer},
                initializers={"w": initializer},
                name="false",
            )
            outputs_f = snt.BatchApply(module3, n_dims=2)(outputs)
            outputs = outputs_t * y + outputs_f * (1-y)
            return outputs
    # gradient-based local updater, used in ablation study
    @snt.reuse_variables
    def gradient_local_updater(self, r, y, x=None, iter=""):
        with tf.compat.v1.variable_scope("gradient_local_updater{}".format(iter), reuse=tf.compat.v1.AUTO_REUSE):
            lr = tf.compat.v1.get_variable(
                "lr", [1, self._num_classes * self._dim_reprs],
                dtype=self._float_dtype,
                initializer=tf.constant_initializer(1.0), trainable=True)
        classifier_weights = self.forward_decoder(r)
        tr_loss, _= self.calculate_loss_and_acc(
                x, y, classifier_weights)
        batch_tr_loss = tf.reduce_mean(tr_loss)
        loss_grad = tf.gradients(batch_tr_loss, r)[0]
        updates = - lr * loss_grad
        return updates

    ### Kernel and Attention ###
    @snt.reuse_variables
    def forward_kernel_or_attention(self, querys, keys, values, iter=""):
        if self._use_kernel:
            if self._kernel_type == "se":
                rtn_values = self.squared_exponential_kernel(querys, keys, values, iter=iter)
            elif self._kernel_type == 'deep_se':
                rtn_values = self.deep_se_kernel(querys, keys, values, iter=iter)
            else:
                raise NameError("Unknown kernel type")
            tf.compat.v1.logging.info("forwarded {0} kernel".format(self._kernel_type))
        else:
            rtn_values = self.attention_block(querys, keys, values, iter=iter)
            tf.compat.v1.logging.info("forwarded {0} attention".format(self._attention_type))
        return rtn_values

    @snt.reuse_variables
    def squared_exponential_kernel(self, querys, keys, values, iter=""):
        num_keys = tf.shape(keys)[0]
        num_querys = tf.shape(querys)[0]
        with tf.compat.v1.variable_scope("squared_exponential_kernel{}".format(iter), reuse=tf.compat.v1.AUTO_REUSE):
            sigma = tf.compat.v1.get_variable("sigma", shape=(), dtype=self._float_dtype,
                    initializer=tf.constant_initializer(1.0), trainable=True)
            lengthscale = tf.compat.v1.get_variable("lengthscale", shape=(), dtype=self._float_dtype,
                    initializer=tf.constant_initializer(1.0), trainable=True)
        _keys = tf.tile(tf.expand_dims(keys, axis=1), [1, num_querys, 1])
        _querys = tf.tile(tf.expand_dims(querys, axis=0), [num_keys, 1, 1])
        sq_norm = tf.reduce_sum((_keys - _querys)**2, axis=-1)
        kernel_qk = sigma**2 * tf.exp(- sq_norm / (2.*lengthscale**2))
        k = kernel_qk
        v = tf.einsum('kq,kv->qv', k, values)
        return v

    @snt.reuse_variables
    def deep_se_kernel(self, querys, keys, values, iter=""):
        with tf.compat.v1.variable_scope("deep_se_kernel{}".format(iter), reuse=tf.compat.v1.AUTO_REUSE):
            # deep embedding of keys and querys
            regularizer = tf.contrib.layers.l2_regularizer(self._l2_penalty_weight)
            initializer = tf.initializers.glorot_uniform(dtype=self._float_dtype)
            module = snt.nets.MLP(
                [self.embedding_dim] * self._embedding_layers,
                activation=self._nonlinearity,
                use_bias=True,
                regularizers={"w": regularizer},
                initializers={"w": initializer},
            )
            keys = snt.BatchApply(module, n_dims=1)(keys)
            querys = snt.BatchApply(module, n_dims=1)(querys)
        num_keys = tf.shape(keys)[0]
        num_querys = tf.shape(querys)[0]
        with tf.compat.v1.variable_scope("deep_se_kernel"):
            sigma = tf.compat.v1.get_variable("sigma", shape=(), dtype=self._float_dtype,
                    initializer=tf.constant_initializer(1.0), trainable=True)
            lengthscale = tf.compat.v1.get_variable("lengthscale", shape=(), dtype=self._float_dtype,
                    initializer=tf.constant_initializer(1.0), trainable=True)
        # compute \sum_i k(x, x_i)u_i
        _keys = tf.tile(tf.expand_dims(keys, axis=1), [1, num_querys, 1])
        _querys = tf.tile(tf.expand_dims(querys, axis=0), [num_keys, 1, 1])
        sq_norm = tf.reduce_sum((_keys - _querys)**2, axis=-1)
        kernel_qk = sigma**2 * tf.exp(- sq_norm / (2.*lengthscale**2))
        k = kernel_qk
        v = tf.einsum('kq,kv->qv', k, values)
        return v

    @snt.reuse_variables
    def attention_block(self, querys, keys, values, iter=""):
        config = {
            "rep": "mlp",
            "output_sizes": [self.embedding_dim] * self._embedding_layers,
            "att_type": self._attention_type,
            "normalise": True,
            "scale": 1.0,
            "l2_penalty_weight": self._l2_penalty_weight,
            "nonlinearity": self._nonlinearity,
        }
        with tf.compat.v1.variable_scope("attention_block{}".format(iter), reuse=tf.compat.v1.AUTO_REUSE):
            attention = Attention(config=config)
            v = attention(keys, querys, values)
        return v

    ### Decoder ###
    @snt.reuse_variables
    def forward_decoder(self, cls_reprs):
        if self._no_decoder:
            # use functional representation directly as the predictor, used in ablation study
            tf.compat.v1.logging.info("no decoder used")
            return cls_reprs
        s = cls_reprs.shape.as_list()
        cls_reprs = tf.reshape(cls_reprs, s[:-1]+[self._num_classes, self._dim_reprs])
        weights_dist_params = self.decoder(cls_reprs)
        fan_in = self.embedding_dim
        fan_out = self._num_classes
        stddev_offset = np.sqrt(2. / (fan_out + fan_in))
        classifier_weights = self.sample(weights_dist_params,
                                                  stddev_offset=stddev_offset)
        return classifier_weights
    # this decoder generates weights of softmax
    @snt.reuse_variables
    def decoder(self, inputs):
        with tf.compat.v1.variable_scope("decoder"):
            l2_regularizer = tf.contrib.layers.l2_regularizer(self._l2_penalty_weight)
            orthogonality_reg = get_orthogonality_regularizer(
                self._orthogonality_penalty_weight)
            initializer = tf.initializers.glorot_uniform(dtype=self._float_dtype)
            # 2 * embedding_dim, because we are returning means and variances
            decoder_module = snt.Linear(
                self.embedding_dim * 2,
                use_bias=True,
                regularizers={"w": l2_regularizer},
                initializers={"w": initializer},
            )
            outputs = snt.BatchApply(decoder_module, n_dims=2)(inputs)
            self._orthogonality_reg = orthogonality_reg(decoder_module.w)
            return outputs

    ### Other ###
    @property
    def dropout_rate(self):
        return self._dropout_rate if self.is_training else 0.0

    @property
    def _l2_regularization(self):
        return tf.cast(tf.reduce_sum(tf.compat.v1.get_collection(tf.compat.v1.GraphKeys.REGULARIZATION_LOSSES)),
                                dtype=self._float_dtype)

    def loss_fn(self, model_outputs, original_classes):
        original_classes = tf.squeeze(original_classes, axis=-1)
        one_hot_outputs = tf.one_hot(original_classes, depth=self._num_classes)
        return tf.compat.v1.losses.softmax_cross_entropy(onehot_labels=one_hot_outputs, logits=model_outputs, \
                label_smoothing=self._label_smoothing, reduction=tf.compat.v1.losses.Reduction.NONE)

    def predict(self, inputs, weights):
        if self._no_decoder:
            return weights
        after_dropout = tf.nn.dropout(inputs, rate=self.dropout_rate)
        preds = tf.einsum("ik,imk->im", after_dropout, weights)
        return preds

    def calculate_loss_and_acc(self, inputs, true_outputs, classifier_weights):
        model_outputs = self.predict(inputs, classifier_weights)
        model_predictions = tf.argmax(
                    model_outputs, -1, output_type=self._int_dtype)
        accuracy = tf.contrib.metrics.accuracy(model_predictions,
                    tf.squeeze(true_outputs, axis=-1))
        return self.loss_fn(model_outputs, true_outputs), accuracy

    def sample(self, distribution_params, stddev_offset=0.):
        means, unnormalized_stddev = tf.split(distribution_params, 2, axis=-1)
        stddev = tf.exp(unnormalized_stddev)
        stddev -= (1. - stddev_offset)
        stddev = tf.maximum(stddev, 1e-10)
        distribution = tf.distributions.Normal(loc=means, scale=stddev)
        if not self.is_training:
            return means
        samples = distribution.sample()
        return samples

    @property
    def _decoder_orthogonality_reg(self):
        return self._orthogonality_reg

class MetaFunRegressor(snt.AbstractModule):

    def __init__(self, name="MetaFunRegressor"):
        super(MetaFunRegressor, self).__init__(name=name)
        self._float_dtype = tf.float32
        self._int_dtype = tf.int32
        # components configurations
        self._use_kernel = FLAGS.use_kernel
        self._use_gradient = FLAGS.use_gradient
        self._attention_type = FLAGS.attention_type
        self._kernel_type = FLAGS.kernel_type
        self._no_decoder = FLAGS.no_decoder
        if self._no_decoder:
            self._dim_reprs = 1
        self._initial_state_type = FLAGS.initial_state_type
        # neural module configurations
        self._nn_size = FLAGS.nn_size
        self._nn_layers = FLAGS.nn_layers
        self._dim_reprs = FLAGS.dim_reprs
        self._num_iters = FLAGS.num_iters
        self._embedding_layers = FLAGS.embedding_layers
        # regularisation configurations
        self._l2_penalty_weight = FLAGS.l2_penalty_weight
        self._dropout_rate = FLAGS.dropout_rate
        self._orthogonality_penalty_weight = FLAGS.orthogonality_penalty_weight
        #
        self._initial_inner_lr = FLAGS.initial_inner_lr
        self._orthogonality_reg = 0
        self._loss_type = "mse" # mse | log_prob
        self._nonlinearity = tf.nn.relu
        self._repr_as_inputs = FLAGS.repr_as_inputs

    def _build(self, data, is_training=True):
        self.is_training = is_training
        self.embedding_dim = data.tr_input.get_shape()[-1].value
        tr_input = data.tr_input
        val_input = data.val_input
        tr_output = data.tr_output
        val_output = data.val_output
        # initial states
        tr_reprs = self.forward_initialiser(tr_input)
        val_reprs = self.forward_initialiser(val_input)
        all_tr_reprs = [tr_reprs]
        all_val_reprs = [val_reprs]
        # inner learning rate
        alpha = tf.compat.v1.get_variable("alpha", [1, 1], dtype=self._float_dtype,
                initializer=tf.constant_initializer(self._initial_inner_lr), trainable=True)
        # iterative functional updating
        for k in range(self._num_iters):
            updates = self.forward_local_updater(r=tr_reprs, y=tr_output, x=tr_input)
            tr_updates = alpha * self.forward_kernel_or_attention(querys=tr_input, keys=tr_input, values=updates)
            val_updates = alpha * self.forward_kernel_or_attention(querys=val_input, keys=tr_input, values=updates)
            tr_reprs += tr_updates
            val_reprs += val_updates
            all_tr_reprs.append(tr_reprs)
            all_val_reprs.append(val_reprs)

        # record predictions at each iteration for visualisation
        all_val_mu = []
        all_val_sigma = []
        output_sizes = [self._nn_size] * (self._nn_layers-1) + [2] # architecture of the predictor
        # decoder r_t(x) at each iteration into the predictor for visualisation
        for k in range(self._num_iters+1):
            weights = self.forward_decoder(all_tr_reprs[k], output_sizes=output_sizes) # generate weights of the predictor
            tr_mu, tr_sigma = self.predict(tr_input, weights, output_sizes=output_sizes) # forward the predictor
            weights = self.forward_decoder(all_val_reprs[k], output_sizes=output_sizes)
            val_mu, val_sigma = self.predict(val_input, weights, output_sizes=output_sizes)
            all_val_mu.append(val_mu)
            all_val_sigma.append(val_sigma)
        #
        tr_loss, tr_metric = self.calculate_loss_and_metrics(
                tr_output, tr_mu, tr_sigma)
        val_loss, val_metric = self.calculate_loss_and_metrics(
                val_output, val_mu, val_sigma)
        batch_tr_loss = tf.reduce_mean(tr_loss)
        batch_tr_metric = tf.reduce_mean(tr_metric)
        batch_val_loss = tf.reduce_mean(val_loss)
        batch_val_metric = tf.reduce_mean(val_metric)
        #
        all_val_mu = tf.stack(all_val_mu, axis=0)
        all_val_sigma = tf.stack(all_val_sigma, axis=0)
        #
        regularization_penalty = (
           self._l2_regularization + self._decoder_orthogonality_reg)
        return batch_tr_loss + batch_val_loss + regularization_penalty, batch_tr_metric, batch_val_metric, \
                data.tr_input, data.tr_output, data.tr_func, data.val_input, data.val_output, \
                data.val_func, all_val_mu, all_val_sigma

    ### Initialiser r_0(x) ###
    @snt.reuse_variables
    def forward_initialiser(self, x):
        num_points = tf.shape(x)[0]
        if self._initial_state_type == 'zero':
            reprs = self.constant_initialiser(num_points, trainable=False)
        elif self._initial_state_type == "constant":
            reprs = self.constant_initialiser(num_points, trainable=True)
        elif self._initial_state_type == "parametric":
            reprs = self.parametric_initialiser(x)
        else:
            raise NameError("Unknown initial state type")
        tf.compat.v1.logging.info("forwarded {0} initialiser".format(self._initial_state_type))
        return reprs
    # r_0(x) = c
    @snt.reuse_variables
    def constant_initialiser(self, num_points, trainable=False):
        with tf.compat.v1.variable_scope("constant_initialiser"):
            if trainable:
                init = tf.compat.v1.get_variable(
                    "initial_state", [1, self._dim_reprs],
                    dtype=self._float_dtype,
                    initializer=tf.constant_initializer(0.0), trainable=True)
            else:
                init = tf.zeros([1, self._dim_reprs])
            init = tf.tile(init, [num_points, 1])
            return init
    # r_0(x) = MLP(x)
    @snt.reuse_variables
    def parametric_initialiser(self, x):
        with tf.compat.v1.variable_scope("parametric_initialiser"):
            regularizer = tf.contrib.layers.l2_regularizer(self._l2_penalty_weight)
            initializer = tf.initializers.glorot_uniform(dtype=self._float_dtype)
            module = snt.nets.MLP(
                [self._nn_size] * self._nn_layers + [self._dim_reprs],
                activation=self._nonlinearity,
                use_bias=True,
                regularizers={"w": regularizer},
                initializers={"w": initializer},
            )
            outputs = snt.BatchApply(module, n_dims=1)(x)
            return outputs

    ### Local Updater u ###
    @snt.reuse_variables
    def forward_local_updater(self, r, y, x):
        if self._use_gradient:
            updates = self.gradient_local_updater(r=r, y=y, x=x)
            tf.compat.v1.logging.info("forwarded gradient local updater")
        else:
            updates = self.neural_local_updater(r=r, y=y, x=x)
            tf.compat.v1.logging.info("forwarded neural local updater")
        return updates
    # neural local updater, for regression, we simply concatenate [r, y, x]
    @snt.reuse_variables
    def neural_local_updater(self, r, y, x=None):
        with tf.compat.v1.variable_scope("neural_local_updater"):
            if x is not None:
                reprs = tf.concat([r, y, x], axis=-1)
            else:
                reprs = tf.concat([r, y], axis=-1)
            regularizer = tf.contrib.layers.l2_regularizer(self._l2_penalty_weight)
            initializer = tf.initializers.glorot_uniform(dtype=self._float_dtype)
            module = snt.nets.MLP(
                [self._nn_size] * self._nn_layers + [self._dim_reprs],
                activation=self._nonlinearity,
                use_bias=True,
                regularizers={"w": regularizer},
                initializers={"w": initializer},
            )
            outputs = snt.BatchApply(module, n_dims=1)(reprs)
            return outputs
    # gradient-based local updater, used in ablation study
    @snt.reuse_variables
    def gradient_local_updater(self, r, y, x=None):
        with tf.compat.v1.variable_scope("gradient_local_updater"):
            lr = tf.compat.v1.get_variable(
                "lr", [1, self._dim_reprs],
                dtype=self._float_dtype,
                initializer=tf.constant_initializer(1.0), trainable=True)

        r = tf.stop_gradient(r)
        weights = self.forward_decoder(r)
        tr_mu, tr_sigma = self.predict(x, weights)
        tr_loss, tr_mse = self.calculate_loss_and_metrics(
                tr_mu, tr_sigma, y)

        # self.debug_op = tf.gradients(weights, reprs)[0]
        batch_tr_loss = tf.reduce_mean(tr_loss)
        loss_grad = tf.gradients(batch_tr_loss, r)[0]
        updates = - lr * loss_grad
        return updates

    ### Kernel and Attention ###
    @snt.reuse_variables
    def forward_kernel_or_attention(self, querys, keys, values):
        if self._use_kernel:
            if self._kernel_type == "se":
                rtn_values = self.squared_exponential_kernel(querys, keys, values)
            elif self._kernel_type == 'deep_se':
                rtn_values = self.deep_se_kernel(querys, keys, values)
            else:
                raise NameError("Unknown kernel type")
            tf.compat.v1.logging.info("forwarded {0} kernel".format(self._kernel_type))
        else:
            rtn_values = self.attention_block(querys, keys, values)
            tf.compat.v1.logging.info("forwarded {0} attention".format(self._attention_type))

        return rtn_values

    @snt.reuse_variables
    def squared_exponential_kernel(self, querys, keys, values):
        num_keys = tf.shape(keys)[0]
        num_querys = tf.shape(querys)[0]
        with tf.compat.v1.variable_scope("squared_exponential_kernel"):
            sigma = tf.compat.v1.get_variable("sigma", shape=(), dtype=self._float_dtype,
                    initializer=tf.constant_initializer(1.0), trainable=True)
            lengthscale = tf.compat.v1.get_variable("lengthscale", shape=(), dtype=self._float_dtype,
                    initializer=tf.constant_initializer(1.0), trainable=True)
        _keys = tf.tile(tf.expand_dims(keys, axis=1), [1, num_querys, 1])
        _querys = tf.tile(tf.expand_dims(querys, axis=0), [num_keys, 1, 1])
        sq_norm = tf.reduce_sum((_keys - _querys)**2, axis=-1)
        kernel_qk = sigma**2 * tf.exp(- sq_norm / (2.*lengthscale**2))
        k = kernel_qk
        v = tf.einsum('kq,kv->qv', k, values)
        return v

    @snt.reuse_variables
    def deep_se_kernel(self, querys, keys, values):
        with tf.compat.v1.variable_scope("deep_se_kernel"):
            regularizer = tf.contrib.layers.l2_regularizer(self._l2_penalty_weight)
            initializer = tf.initializers.glorot_uniform(dtype=self._float_dtype)
            module = snt.nets.MLP(
                [self._nn_size] * self._embedding_layers,
                activation=self._nonlinearity,
                use_bias=True,
                regularizers={"w": regularizer},
                initializers={"w": initializer},
            )
            keys = snt.BatchApply(module, n_dims=1)(keys)
            querys = snt.BatchApply(module, n_dims=1)(querys)
        num_keys = tf.shape(keys)[0]
        num_querys = tf.shape(querys)[0]
        with tf.compat.v1.variable_scope("deep_se_kernel"):
            sigma = tf.compat.v1.get_variable("sigma", shape=(), dtype=self._float_dtype,
                    initializer=tf.constant_initializer(1.0), trainable=True)
            lengthscale = tf.compat.v1.get_variable("lengthscale", shape=(), dtype=self._float_dtype,
                    initializer=tf.constant_initializer(1.0), trainable=True)
        # compute \sum_i k(x, x_i)u_i
        _keys = tf.tile(tf.expand_dims(keys, axis=1), [1, num_querys, 1])
        _querys = tf.tile(tf.expand_dims(querys, axis=0), [num_keys, 1, 1])
        sq_norm = tf.reduce_sum((_keys - _querys)**2, axis=-1)
        kernel_qk = sigma**2 * tf.exp(- sq_norm / (2.*lengthscale**2))
        k = kernel_qk
        v = tf.einsum('kq,kv->qv', k, values)
        return v

    @snt.reuse_variables
    def attention_block(self, querys, keys, values):
        config = {
            "rep": "mlp",
            "output_sizes": [self._nn_size] * self._embedding_layers,
            "att_type": self._attention_type,
            "normalise": True,
            "scale": 1.0,
            "l2_penalty_weight": self._l2_penalty_weight,
            "nonlinearity": self._nonlinearity,
        }
        with tf.compat.v1.variable_scope("attention_block"):
            attention = Attention(config=config)
            v = attention(keys, querys, values)
        return v

    @snt.reuse_variables
    def forward_decoder(self, reprs, output_sizes=[40,40,2]):
        if self._no_decoder:
            tf.compat.v1.logging.info("no decoder used")
            return reprs
        else:
            weights_dist_params = self.decoder(reprs, output_sizes=output_sizes)
            stddev_offset = np.sqrt(1. / self._nn_size)
            weights = self.sample(weights_dist_params, stddev_offset=stddev_offset)
            tf.compat.v1.logging.info("forwarded decoder")
            return weights

    @snt.reuse_variables
    def decoder(self, reprs, output_sizes=[40,40,2]):
        with tf.compat.v1.variable_scope("decoder"):
            # _repr_as_inputs = True: Representation is used as inputs to the predictor
            # _repr_as_inputs = False: Representation is used to generate weights of the predictor
            if self._repr_as_inputs:
                return reprs
            regularizer = tf.contrib.layers.l2_regularizer(self._l2_penalty_weight)
            initializer = tf.initializers.glorot_uniform(dtype=self._float_dtype)
            num_layers = len(output_sizes)
            output_sizes = [self.embedding_dim] + output_sizes
            # count number of parameters in the predictor
            num_params = 0
            for i in range(num_layers):
                num_params += (output_sizes[i]+1) * output_sizes[i+1]
            # decode the representation into the weights of the predictor
            module = snt.nets.MLP(
                [self._nn_size] * self._nn_layers + [2 * num_params],
                activation=self._nonlinearity,
                use_bias=True,
                regularizers={"w": regularizer},
                initializers={"w": initializer},
            )
            outputs = snt.BatchApply(module, n_dims=1)(reprs)
            return outputs

    @snt.reuse_variables
    def predict(self, inputs, weights, output_sizes=[40,40,2]):
        # no_decoder = True: functional representation is the predictor itself (unused for regression problems)
        if self._no_decoder:
            if self._dim_reprs == 1: # predictive mean
                return weights, tf.ones_like(weights) * 0.5
            elif self._dim_reprs == 2: # predictive mean and std
                return tf.split(weights, 2, axis=-1)
            else:
                raise Exception("num_reprs must <=2 if no_decoder")
        if self._repr_as_inputs:
            with tf.compat.v1.variable_scope("predict"):
                regularizer = tf.contrib.layers.l2_regularizer(self._l2_penalty_weight)
                initializer = tf.initializers.glorot_uniform(dtype=self._float_dtype)
                outputs = tf.concat([weights, inputs], axis=-1) # weights is actually repr for repr_as_inputs=True
                # construct the predictor conditioned on the repr (weights).
                # repr is feeded into the network at each layer.
                for i, s in enumerate(output_sizes):
                    module = snt.nets.MLP(
                        [s],
                        use_bias=True,
                        regularizers={"w": regularizer},
                        initializers={"w": initializer},
                    )
                    outputs = snt.BatchApply(module, n_dims=1)(outputs)
                    if i < len(output_sizes)-1:
                        outputs = self._nonlinearity(outputs)
                        outputs = tf.concat([outputs, weights], axis=-1)
                preds = outputs
        else:
            # use the generated weights to construct the predictor
            num_layers = len(output_sizes)
            output_sizes = [self.embedding_dim] + output_sizes
            begin = 0
            preds = inputs
            for i in range(num_layers):
                in_size = output_sizes[i]
                out_size = output_sizes[i+1]
                end = begin + in_size * out_size
                w = tf.reshape(weights[:, begin:end], [-1, in_size, out_size])
                b = tf.reshape(weights[:, end:end+out_size], [-1, out_size])
                begin = end + out_size
                preds = tf.einsum("ik,ikm->im", preds, w) + b
                if i < num_layers - 1:
                    preds = self._nonlinearity(preds)
        # return preds
        mu, log_sigma = tf.split(preds, 2, axis=-1)
        sigma = 0.1 + 0.9 * tf.nn.softplus(log_sigma)
        return mu, sigma

    def sample(self, distribution_params, stddev_offset=0.):
        # here we consider a deterministic case, but one can try the probabilstic version
        means, unnormalized_stddev = tf.split(distribution_params, 2, axis=-1)
        return means
        ## probabilstic version:
        # stddev = tf.exp(unnormalized_stddev)
        # stddev -= (1. - stddev_offset)
        # stddev = tf.maximum(stddev, 1e-10)
        # distribution = tf.distributions.Normal(loc=means, scale=stddev)
        # if not self.is_training:
        #     return means
        # samples = distribution.sample()
        # return samples

    def loss_fn(self, model_outputs, labels):
        return tf.losses.mean_squared_error(labels=labels, predictions=model_outputs,
                reduction=tf.compat.v1.losses.Reduction.NONE)

    def calculate_loss_and_metrics(self, target_y, mus, sigmas, coeffs=None):
        if self._loss_type == "mse":
            mu, sigma = mus, sigmas
            mse = self.loss_fn(mu, target_y)
            return mse, mse
        elif self._loss_type == "log_prob":
            mu, sigma = mus, sigmas
            dist = tf.contrib.distributions.MultivariateNormalDiag(loc=mu, scale_diag=sigma)
            loss = - dist.log_prob(target_y)
            mse = self.loss_fn(mu, target_y)
            return loss, loss
        else:
            raise NameError("unknown output_dist_type")

    @property
    def _l2_regularization(self):
        return tf.cast(
            tf.reduce_sum(tf.compat.v1.get_collection(tf.compat.v1.GraphKeys.REGULARIZATION_LOSSES)),
            dtype=self._float_dtype)

    @property
    def _decoder_orthogonality_reg(self):
        return self._orthogonality_reg

# (Copy from https://github.com/deepmind/leo, see copyright and original license in our LICENSE file.)
def get_orthogonality_regularizer(orthogonality_penalty_weight):
  """Returns the orthogonality regularizer."""
  def orthogonality(weight):
    """Calculates the layer-wise penalty encouraging orthogonality."""
    with tf.name_scope(None, "orthogonality", [weight]) as name:
      w2 = tf.matmul(weight, weight, transpose_b=True)
      wn = tf.norm(weight, ord=2, axis=1, keepdims=True) + 1e-32
      correlation_matrix = w2 / tf.matmul(wn, wn, transpose_b=True)
      matrix_size = correlation_matrix.get_shape().as_list()[0]
      base_dtype = weight.dtype.base_dtype
      identity = tf.eye(matrix_size, dtype=base_dtype)
      weight_corr = tf.reduce_mean(
          tf.math.squared_difference(correlation_matrix, identity))
      return tf.multiply(
          tf.cast(orthogonality_penalty_weight, base_dtype),
          weight_corr,
          name=name)
  return orthogonality

# Attention modules
# (Adapted from https://github.com/deepmind/neural-processes, see copyright and original license in our LICENSE file.)
def dot_product_attention(q, k, v, normalise):
    """Computes dot product attention.

    Args:
        q: queries. tensor of  shape [B,m,d_k].
        k: keys. tensor of shape [B,n,d_k].
        v: values. tensor of shape [B,n,d_v].
        normalise: Boolean that determines whether weights sum to 1.
    Returns:
        tensor of shape [B,m,d_v].
    """
    d_k = tf.shape(q)[-1]
    scale = tf.sqrt(tf.cast(d_k, tf.float32))
    unnorm_weights = tf.einsum('jk,ik->ij', k, q) / scale # [B,m,n]
    if normalise:
        weight_fn = tf.nn.softmax
    else:
        weight_fn = tf.sigmoid
    weights = weight_fn(unnorm_weights)
    rep = tf.einsum('ik,kj->ij', weights, v)
    return rep

class Attention(snt.AbstractModule):

    def __init__(self, config=None, name="attention"):
        super(Attention, self).__init__(name=name)
        self._float_dtype = tf.float32
        self._int_dtype = tf.int32
        self._rep = config['rep']
        self._output_sizes = config['output_sizes']
        self._att_type = config['att_type']
        self._normalise = config['normalise']
        self._scale = config['scale']
        self._l2_penalty_weight = config['l2_penalty_weight']
        self._nonlinearity = config['nonlinearity']

    def _build(self, x1, x2, r):
        if self._rep == 'identity':
            k, q = (x1, x2)
        elif self._rep == 'mlp':
            # Pass through MLP
            initializer = tf.initializers.glorot_uniform(dtype=self._float_dtype)
            regularizer = tf.contrib.layers.l2_regularizer(self._l2_penalty_weight)
            module = snt.nets.MLP(
                self._output_sizes,
                activation=self._nonlinearity,
                use_bias=True,
                initializers={"w": initializer},
            )
            k = snt.BatchApply(module, n_dims=1)(x1)
            q = snt.BatchApply(module, n_dims=1)(x2)
        else:
            raise NameError("'rep' not among ['identity','mlp']")

        if self._att_type == 'dot_product':
            rep = dot_product_attention(q, k, r, self._normalise)
        else:
            raise NameError(("'att_type' not among ['dot_product']"))
        return rep


import math
import random
import pygame
import config
import requests
import json
from center import PokeCenter
from mart import PokeMart
from trainer import Trainer
from mon import Pokemon
from player import Player, PlayerOri
from game_state import GameState

class Game(object):
    def __init__(self, screen):
        self.screen = screen
        self.objects = []
        self.game_state = GameState.NONE
        self.map = []
        self.poke_centers = []
        self.poke_marts = []
        self.poke_trainers = []
        self.camera = [0, 0]

    def set_up(self):
        self.load_map("01")
        self.player = Player(24, 19)
        self.objects.append(self.player)
        self.pokemons = self.request_pokemons()
        self.load_assets()
        print("Gotta catch'em all !!")
        self.game_state = GameState.RUNNING

    def update(self, human_player = False):
        self.screen.fill(config.BLACK)

        if(human_player == True):
            self.handle_events()

        self.render_map(self.screen)

        for object in self.objects:
            object.render(self.screen, self.camera)

    def handle_events(self):
        for event in pygame.event.get():
            # HANDLE QUIT EVENT
            if event.type == pygame.QUIT:
                self.game_state = GameState.ENDED
            # HANDLE KEY EVENT
            elif event.type == pygame.KEYDOWN:
                if event.key == pygame.K_ESCAPE:
                    self.game_state = GameState.ENDED
                elif event.key == pygame.K_w:  # WALK FORWARD
                    self.player.walk_forward(self.map, self.index_map)
                elif event.key == pygame.K_d:  # TURN RIGHT
                    self.player.change_orientation(PlayerOri.TURN_RIGHT)
                elif event.key == pygame.K_a:  # TURN LEFT
                    self.player.change_orientation(PlayerOri.TURN_LEFT)
                elif event.key == pygame.K_z:  # INTERACT
                    self.player.interact(self.index_map, self.objects)
                    # if(len(self.player.captured_pokemons) == 150):
                    #     self.game_state = GameState.WIN # TODO: Implement winning screen

    def load_map(self, file_name):
        with open('maps/' + file_name + ".txt") as map_file:
            for line in map_file:
                tiles = []
                for i in range(0, len(line) - 1, 2):
                    tiles.append(line[i])
                self.map.append(tiles)

    def load_assets(self):
        print("Loading Assets...")
        centers = 20
        marts = 15
        trainers = 50
        mons = 150

        rows, cols = 42, 42
        self.index_map = []
        for i in range(0, rows):
            self.index_map.append([])
            for j in range(0, cols):
                self.index_map[i].append([])

#        self.index_map[self.player.position[1]][self.player.position[0]].append(self.player)

        poke_list = list(self.pokemons.values()).copy()
        random.shuffle(poke_list)

        # TODO: Improve spawning rules (e.g: No water pokemons, trainers, centers or marts on lava, no fire pokemons on water)
        while(centers + marts + trainers + mons > 0):
            i = random.randint(0, 41)
            j = random.randint(0, 41)
            if(self.index_map[i][j] == []):
                dice = random.randint(0, 3)

                if(dice == 0 and centers > 0):
                    self.poke_centers.append(PokeCenter(j, i))
                    self.objects.append(self.poke_centers[-1])
                    centers -= 1
                    self.index_map[i][j].append(self.poke_centers[-1])
                    if(i+1 <= 41):
                        self.index_map[i+1][j].append('smell')
                    if(i-1 >= 0):
                        self.index_map[i-1][j].append('smell')
                    if(j+1 <= 41):
                        self.index_map[i][j+1].append('smell')
                    if(j-1 >= 0):
                        self.index_map[i][j-1].append('smell')
                elif(dice == 1 and marts > 0):
                    self.poke_marts.append(PokeMart(j, i))
                    self.objects.append(self.poke_marts[-1])
                    marts -= 1
                    self.index_map[i][j].append(self.poke_marts[-1])
                    if(i+1 <= 41):
                        self.index_map[i+1][j].append('vendingShout')
                    if(i-1 >= 0):
                        self.index_map[i-1][j].append('vendingShout')
                    if(j+1 <= 41):
                        self.index_map[i][j+1].append('vendingShout')
                    if(j-1 >= 0):
                        self.index_map[i][j-1].append('vendingShout')
                elif(dice == 2 and trainers > 0):
                    self.poke_trainers.append(Trainer(j, i))
                    self.objects.append(self.poke_trainers[-1])
                    trainers -= 1
                    self.index_map[i][j].append(self.poke_trainers[-1])
                    if(i+1 <= 41):
                        self.index_map[i+1][j].append('battleShout')
                    if(i-1 >= 0):
                        self.index_map[i-1][j].append('battleShout')
                    if(j+1 <= 41):
                        self.index_map[i][j+1].append('battleShout')
                    if(j-1 >= 0):
                        self.index_map[i][j-1].append('battleShout')
                elif(dice == 3 and mons > 0):
                    poke_name = poke_list[mons-1].name
                    self.pokemons[poke_name].set_pos(j, i)
                    self.objects.append(self.pokemons[poke_name])
                    mons -= 1
                    self.index_map[i][j].append(self.pokemons[poke_name])
                    if(i+1 <= 41):
                        self.index_map[i+1][j].append('dexWarning')
                    if(i-1 >= 0):
                        self.index_map[i-1][j].append('dexWarning')
                    if(j+1 <= 41):
                        self.index_map[i][j+1].append('dexWarning')
                    if(j-1 >= 0):
                        self.index_map[i][j-1].append('dexWarning')

        if(centers == marts == trainers == mons == 0):
            print("NO MISSING ASSETS!!")
        print("Assets Loaded!")

    def render_map(self, screen):
        self.determine_camera()

        y_pos = 0
        for line in self.map:
            x_pos = 0
            for tile in line:
                image = map_tile_image[tile]
                rect = pygame.Rect(x_pos * config.SCALE - (self.camera[0] * config.SCALE), y_pos * config.SCALE - (
                    self.camera[1] * config.SCALE), config.SCALE, config.SCALE)
                screen.blit(image, rect)
                x_pos = x_pos + 1

            y_pos = y_pos + 1

    def request_pokemons(self):
        print("Loading Pokemons...")
        poke_dict = {}
        #self.poke_data = {}
        for i in range(1, 151):

            ''' # REQUESTS POKEMON DETAILS FROM API AND SAVES BOTH SIMPLIFIED JSON AND SPRITES
                # UNCOMMENT THIS IF RUNNING FOR THE FIRTS TIME
            pj = self.get_api_json('https://pokeapi.co/api/v2/pokemon/' + str(i))
            self.save_poke_json(pj, self.poke_data, i)
            self.download_pokemon_sprite(pj, name)
            '''

            # READ FROM JSON FOLDER
            name, poke_types = self.read_poke_json(
                'jsons/names_and_types.json', i)
            sprite_file = 'sprites/pokemon/' + name + ".png"

            poke_dict[name] = Pokemon(name, poke_types, sprite_file)
            #print("Pokemon: {} created!".format(name))

        print("Pokemons Loaded!")

        return poke_dict

    def get_api_json(self, url):
        poke_details = url
        poke_res = requests.get(poke_details)
        pj = poke_res.json()
        return pj

    def save_poke_json(self, poke_json, poke_data, index):
        poke_data[index] = {'name': poke_json['name']}
        poke_data[index].update({'types': [x['type']['name']
                                for x in poke_json['types']]})
        # TODO: This could be optimized. As of now it re-writes the whole file on every loop.
        with open('jsons/names_and_types.json', 'w') as f:
            json.dump(poke_data, f)
        #print("Json Saved: {}".format(poke_data))

    def read_poke_json(self, json_file, index):
        with open(json_file, 'r') as f:
            poke_data = json.load(f)
            name = poke_data[str(index)]['name']
            types = poke_data[str(index)]['types'].copy()
            return name, types

    def download_pokemon_sprite(self, json, pokemon):
        sprite_url = json['sprites']['front_default']
        img_data = requests.get(sprite_url).content
        with open('sprites/pokemon/' + pokemon + ".png", 'wb') as handler:
            handler.write(img_data)

    def determine_camera(self):
        # DETERMINE Y POSITION
        max_y_position = len(self.map) - config.SCREEN_HEIGHT / config.SCALE
        y_position = self.player.position[1] - \
            math.ceil(round(config.SCREEN_HEIGHT / config.SCALE / 2))

        if(y_position <= max_y_position and y_position >= 0):
            self.camera[1] = y_position
        elif(y_position < 0):
            self.camera[1] = 0
        else:
            self.camera[1] = max_y_position

        # DETERMINE X POSITION
        max_x_position = len(self.map[0]) - config.SCREEN_WIDTH / config.SCALE
        x_position = self.player.position[0] - \
            math.ceil(round(config.SCREEN_WIDTH / config.SCALE / 2))

        if(x_position <= max_x_position and x_position >= 0):
            self.camera[0] = x_position
        elif(x_position < 0):
            self.camera[0] = 0
        else:
            self.camera[0] = max_x_position

map_tile_image = {
    "G": pygame.transform.scale(pygame.image.load("sprites/grass.png"), (config.SCALE, config.SCALE)),
    "W": pygame.transform.scale(pygame.image.load("sprites/water.png"), (config.SCALE, config.SCALE)),
    "M": pygame.transform.scale(pygame.image.load("sprites/mountain.png"), (config.SCALE, config.SCALE)),
    "L": pygame.transform.scale(pygame.image.load("sprites/lava.png"), (config.SCALE, config.SCALE)),
    "C": pygame.transform.scale(pygame.image.load("sprites/cave.png"), (config.SCALE, config.SCALE))
}