import cv2

#Doc 2 buc anh
s_img = cv2.imread("D:\\Python_XLA\\nhandangdoituong_xacdinhlabai\\img_Data\\A_ro.png")
s_img = cv2.cvtColor(s_img, cv2.COLOR_BGR2RGB)
t_img = cv2.imread("D:\\Python_XLA\\nhandangdoituong_xacdinhlabai\\img_Query\\A_ro.png")
t_img = cv2.cvtColor(t_img, cv2.COLOR_BGR2RGB)

method = eval('cv2.TM_CCOEFF')
s_img_copy = s_img.copy()
res = cv2.matchTemplate(s_img_copy, t_img, method)

#Xac dinh va ve vung nhan dang
min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(res)
if method in [cv2.TM_SQDIFF_NORMED, cv2.TM_SQDIFF]:
    top_left = min_loc
else:
    top_left = max_loc

heigh, width, channel = t_img.shape
bottom_right = (top_left[0]+width, top_left[1]+heigh)
cv2.rectangle(s_img, top_left, bottom_right, (255, 0, 0), 1)

#Hien thi 2 buc anh
cv2.imshow("Anh ket qua", s_img)
if cv2.waitKey(0) & 0xFF == 27:     #ESC dong tat ca man hinh
    cv2.destroyAllWindows()


# Based on http://stackoverflow.com/a/5809952/689985

from django.test.simple import DjangoTestSuiteRunner
from django.test.simple import dependency_ordered

class ByPassableDBDjangoTestSuiteRunner(DjangoTestSuiteRunner):
    def setup_databases(self, **kwargs):
        from django.db import connections, DEFAULT_DB_ALIAS

        # First pass -- work out which databases actually need to be created,
        # and which ones are test mirrors or duplicate entries in DATABASES
        mirrored_aliases = {}
        test_databases = {}
        dependencies = {}
        for alias in connections:
            connection = connections[alias]
            if connection.settings_dict['TEST_MIRROR']:
                # If the database is marked as a test mirror, save
                # the alias.
                mirrored_aliases[alias] = connection.settings_dict['TEST_MIRROR']
            else:
                # Store a tuple with DB parameters that uniquely identify it.
                # If we have two aliases with the same values for that tuple,
                # we only need to create the test database once.
                item = test_databases.setdefault(
                    connection.creation.test_db_signature(),
                    (connection.settings_dict['NAME'], [])
                )
                item[1].append(alias)

                if 'TEST_DEPENDENCIES' in connection.settings_dict:
                    dependencies[alias] = connection.settings_dict['TEST_DEPENDENCIES']
                else:
                    if alias != DEFAULT_DB_ALIAS:
                        dependencies[alias] = connection.settings_dict.get('TEST_DEPENDENCIES', [DEFAULT_DB_ALIAS])

        # Second pass -- actually create the databases.
        old_names = []
        mirrors = []
        for signature, (db_name, aliases) in dependency_ordered(test_databases.items(), dependencies):
            # Actually create the database for the first connection
            connection = connections[aliases[0]]
            if connection.settings_dict.get("USE_LIVE_FOR_TESTS"):
                continue

            old_names.append((connection, db_name, True))

            test_db_name = connection.creation.create_test_db(self.verbosity, autoclobber=not self.interactive)
            for alias in aliases[1:]:
                connection = connections[alias]
                if db_name:
                    old_names.append((connection, db_name, False))
                    connection.settings_dict['NAME'] = test_db_name
                else:
                    # If settings_dict['NAME'] isn't defined, we have a backend where
                    # the name isn't important -- e.g., SQLite, which uses :memory:.
                    # Force create the database instead of assuming it's a duplicate.
                    old_names.append((connection, db_name, True))
                    connection.creation.create_test_db(self.verbosity, autoclobber=not self.interactive)

        for alias, mirror_alias in mirrored_aliases.items():
            mirrors.append((alias, connections[alias].settings_dict['NAME']))
            connections[alias].settings_dict['NAME'] = connections[mirror_alias].settings_dict['NAME']

        return old_names, mirrors

    def teardown_databases(self, old_config, **kwargs):
        from django.db import connections
        old_names, mirrors = old_config
        # Point all the mirrors back to the originals
        for alias, old_name in mirrors:
            connections[alias].settings_dict['NAME'] = old_name
        # Destroy all the non-mirror databases
        for connection, old_name, destroy in old_names:
            if destroy:
                connection.creation.destroy_test_db(old_name, self.verbosity)
            else:
                connection.settings_dict['NAME'] = old_name


import json
import logging
import multiprocessing
import os
import re
import urllib.parse
from datetime import datetime, timedelta
from multiprocessing import Queue
from queue import Empty
from time import sleep
from typing import Any, Iterable, Tuple, Optional
from typing import Callable

import boto3
import pymongo
from botocore.vendored.requests.packages.urllib3.exceptions import ReadTimeoutError
from multiprocessing import Process
from pymongo import MongoClient

# regex

integer_request_id_regex = re.compile(r"\d{10}$")
date_regex = re.compile(r"\d\d\d\d-\d\d-\d\d( 00:00:00)?$")

class CommitBeforeParkingException(Exception):
    pass

class ConcurrentParkingException(Exception):
    pass

class MissingDataPart(Exception):
    pass

def serialize_data_id(data_id) -> str:
    return json.dumps(data_id, sort_keys=True)

def deserialize_data_id(data_id_str) -> Any:
    return json.loads(data_id_str)

def now_timestamp() -> int:
    return int(datetime.now().timestamp() * 1000)

def deserialize_request_id(request_id_str):
    # request_id_str = request_id_key[len(self.crawler_name) + 1:]
    if integer_request_id_regex.match(request_id_str) is not None:
        return int(request_id_str)
    elif date_regex.match(request_id_str) is not None:
        year, month, date = request_id_str.split(" ")[0].split("-")
        return datetime(int(year), int(month), int(date))
    else:
        return request_id_str

def serialize_request_id(request_id) -> str:
    if isinstance(request_id, int):
        ret = "%010d" % request_id
    elif isinstance(request_id, datetime):
        assert request_id.hour == request_id.minute == request_id.second == request_id.microsecond == 0
        ret = "%s" % request_id.date()
    elif isinstance(request_id, str):
        assert integer_request_id_regex.match(request_id) is None and date_regex.match(request_id) is None, \
            "Type ambiguous request string " + request_id
        ret = "%s" % request_id
    else:
        assert False, "Unexpected request id type" + str(request_id)
    assert deserialize_request_id(ret) == request_id
    return ret

def s3_list_by_prefix(bucket_name, prefix) -> Iterable[str]:
    s3 = boto3.client('s3')
    ret = s3.list_objects(
        Bucket=bucket_name,
        Delimiter="/",
        Prefix=prefix,
        RequestPayer='requester'
    )
    for r in ret["Contents"]:
        yield r["Key"]
    while 'NextMarker' in ret:
        assert ret["IsTruncated"]
        ret = s3.list_objects(
            Bucket=bucket_name,
            Delimiter="/",
            Prefix=prefix,
            RequestPayer='requester',
            Marker=ret['NextMarker']
        )
        for r in ret["Contents"]:
            yield r["Key"]
    assert not ret["IsTruncated"]

def mongo_list_by_prefix(coll, prefix) -> Iterable[str]:
    for record in coll.find({"_id": {'$regex': '^' + re.escape(prefix)}}, {"_id": 1}):
        yield record["_id"]

_map_fn = None  # type: Callable[Iterable[dict]]
_db = None

def worker(i, queue: Queue):
    def generator():
        global _db
        # get mongodb in the work process
        _db.get_mongo_colletions()
        while True:
            # get a request chunk from the queue
            try:
                chunk = queue.get_nowait()
            except Empty:
                break
            for req in chunk:
                request_id, data_id, version = _db._parse_data_key(req)
                while True:
                    try:
                        obj = _db.get_data(request_id, data_id, version)
                        body = obj["Body"]
                        my_ret = {
                            "crawler_name": _db.crawler_name,
                            "request_id": request_id,
                              "data_id": data_id,
                              "version": version,
                               "hash": obj["ETag"],
                              "body": body}
                        yield my_ret
                        break
                    except ReadTimeoutError:
                        # FIXME: Sleep
                        logging.exception("Read Timeout")
                        sleep(1)
                        pass
    # the _map_fn should be set by the parallel_scan() function
    global _map_fn
    _map_fn(generator())

CRAWLED = 0
REQUESTED = 1
REDO = 3

class CrawlDB:
    # s3 stuff
    s3 = boto3.client('s3')
    S3_BUCKET_NAME = os.environ.get("CRAWLER_S3_BUCKET", "askmiso-crawler")

    # suppress logging
    logging.getLogger('botocore').setLevel(logging.WARN)
    logging.getLogger('boto3').setLevel(logging.WARN)

    def get_mongo_colletions(self):
        if self.mongo_db is None:
            self.status_coll = MongoClient("mongo").crawldb.status
            self.s3_key_cache = MongoClient("mongo").crawldb.s3_key_cache
        else:
            self.status_coll = self.mongo_db.crawldb.status
            self.s3_key_cache = self.mongo_db.crawldb.s3_key_cache

    def __init__(self, crawler_name: str, request_timeout: timedelta = timedelta(minutes=10),
                 mongo_db=None):
        self.__version__ = "0.7.4"
        self.mongo_db = mongo_db
        self.get_mongo_colletions()
        # make sure the required index is available
        self.status_coll.create_index([("crawler", pymongo.ASCENDING),
                                       ("status", pymongo.ASCENDING),
                                       ("requested_time", pymongo.ASCENDING)])
        self.crawler_name = crawler_name

        assert isinstance(request_timeout, timedelta)
        self.request_timeout = request_timeout

    def _timeout_threshold(self) -> int:
        """
        :return: the requested_time threshold. Any uncommitted request with requested_time before this time is
        considered to be timeout.
        """
        return int((datetime.now() - self.request_timeout).timestamp() * 1000)

    def get_request_id_key(self, request_id) -> str:
        """
        :param request_id:
        :return: a string representation of crawler_name | request_id
        """
        return "|".join([self.crawler_name, serialize_request_id(request_id)])

    def _parse_request_id_key(self, request_id_key):
        request_id_str = request_id_key[len(self.crawler_name) + 1:]
        return deserialize_request_id(request_id_str)

    def _get_data_key(self, request_id, data_id, version: int):
        assert isinstance(version, int)
        return "/".join(map(urllib.parse.quote_plus,
                            [self.crawler_name,
                             serialize_request_id(request_id),
                             serialize_data_id(data_id),
                             str(version)
                             ]))

    def _parse_data_key(self, key: str) -> Tuple[Any, Any, int]:
        crawler_name, request_id_str, data_id_str, version_str = map(urllib.parse.unquote_plus, key.split("/"))
        assert crawler_name == self.crawler_name
        request_id = deserialize_request_id(request_id_str)
        data_id = deserialize_data_id(data_id_str)
        version = int(version_str)
        assert key == self._get_data_key(request_id, data_id, version), "%s != %s" % (
        key, self._get_data_key(request_id, data_id, version))
        return request_id, data_id, version

    def _get_data_key_prefix(self, request_id):

        return "/".join(map(urllib.parse.quote_plus,
                            [self.crawler_name,
                             serialize_request_id(request_id)
                             ]))

    def is_crawled_or_being_crawled(self, request_id):
        """
        Check if a request is being crawled or was already crawled
        :param request_id:
        :return: True or False
        """
        return self.status_coll.find({"_id": self.get_request_id_key(request_id)}).count() > 0

    def park_request(self, request_id):
        """
        Park a request so that other crawlers won't crawl it before it is timeout
        :param request_id: request id
        :return: None
        :raises: ConcurrentParkingException if other crawlers tried to park the request at the same time
        """
        status_record = self.status_coll.find_one(self.get_request_id_key(request_id))
        if status_record is None:
            try:
                self.status_coll.insert_one(
                                             {"_id": self.get_request_id_key(request_id),
                                             "crawler": self.crawler_name,
                                              "request_id": serialize_request_id(request_id),
                                              "requested_time": now_timestamp(),
                                              "status": REQUESTED},
                                             )
            except pymongo.errors.DuplicateKeyError:
                raise ConcurrentParkingException
        elif status_record["status"] in {REQUESTED, REDO} and status_record["requested_time"] < self._timeout_threshold():
            result = self.status_coll.replace_one({"_id": self.get_request_id_key(request_id),
                                                   "requested_time": status_record["requested_time"]},
                                                  {"crawler": self.crawler_name,
                                                   "request_id": serialize_request_id(request_id),
                                                   "requested_time": now_timestamp(),
                                                   "status": REQUESTED},
                                                  False
                                                  )
            if result.matched_count == 0:
                raise ConcurrentParkingException
            elif result.matched_count != 1:
                raise RuntimeError("Something went terribly wrong!")
        else:
            raise ConcurrentParkingException

    def commit_request(self, request_id, meta=None, skip_park=False):
        """
        Commit a request. After a request is committed, it is considered finished.
        :param request_id: request id
        :param meta: metadata to attach the committed request
        :param skip_park: whether to allow the commit without parking the request first
        :return: None
        """
        status_record = self.status_coll.find_one(self.get_request_id_key(request_id))
        if (status_record is not None and status_record["status"] == REQUESTED) or skip_park:
            self.status_coll.replace_one({"_id": self.get_request_id_key(request_id)},
                                         {"crawler": self.crawler_name,
                                          "request_id": serialize_request_id(request_id),
                                          "committed_time": now_timestamp(),
                                          "metadata": meta,
                                          "status": CRAWLED},
                                         True
                                         )
        else:
            raise CommitBeforeParkingException()

    def get_request_meta(self, request_id) -> Optional[Any]:
        """
        :param request_id: request id
        :return: the metadata attached to the request when it was committed
        """
        record = self.status_coll.find_one(self.get_request_id_key(request_id))
        if record is not None:
            return record.get("metadata")

    def get_data_ids(self, request_id):
        """
        :param request_id: request id
        :return: Return the set of data ids that were uploaded for this request
        """
        prefix = self._get_data_key_prefix(request_id)
        data_ids = set()
        for key in mongo_list_by_prefix(self.s3_key_cache, prefix):
            request_id, data_id, version = self._parse_data_key(key)
            assert request_id == request_id
            data_ids.add(data_id)
        return data_ids

    def get_timeout_request_id(self) -> Optional[Any]:
        """
        :return: a request (if any) which was unable to be finished before timeout
        """
        record = self.status_coll.find_one({"status": REQUESTED,
                                            "crawler": self.crawler_name,
                                            "requested_time": {"$lt": self._timeout_threshold()}})
        if record is None:
            record = self.status_coll.find_one({"status": REDO,
                                                "crawler": self.crawler_name})
        if record is not None:
            return deserialize_request_id(record["request_id"])
        else:
            return None

    def save_data(self, request_id, data_id: Any, data, version: int = 0):
        """

        :param request_id: request id (number, datetime, or string)
        :param data_id: data id, anything that can be json seerialized
        :param data: bytes or file like object
        :param version: int (default 0)
        :return:
        """
        file_key = self._get_data_key(request_id, data_id, version)
        self.s3.put_object(Bucket=self.S3_BUCKET_NAME, Body=data, Key=file_key)
        self.s3_key_cache.replace_one({"_id": file_key}, {}, True)

    def get_data(self, request_id, data_id, version=0):
        if self.data_exists(request_id, data_id, version):
            file_key = self._get_data_key(request_id, data_id, version)
            return self.s3.get_object(Bucket=self.S3_BUCKET_NAME, Key=file_key)
        else:
            raise ValueError("request: %s, data id %s does not exist" % (request_id, data_id))

    def data_exists(self, request_id, data_id, version=0):
        return self.s3_key_cache.find_one(self._get_data_key(request_id, data_id, version)) is not None

    def list_data(self) -> Iterable[Tuple]:
        """
        :return: A iterable of tuples of (request_id, data_id, version)
        """
        return map(self._parse_data_key, mongo_list_by_prefix(self.s3_key_cache, self.crawler_name + "/"))

    def parallel_scan_items(self, map_fn, thread_count=None, n_chunks=None, reverse=True, backend=Process) -> Iterable[Any]:
        global _map_fn, _db
        _map_fn = map_fn
        _db = self

        if thread_count is None:
            thread_count = multiprocessing.cpu_count()
        if n_chunks is None:
            n_chunks = thread_count
        requests = list(mongo_list_by_prefix(self.s3_key_cache, self.crawler_name + "/"))
        if reverse:
            requests = list(reversed(requests))

        def chunks(l, n):
            """Yield successive n-sized chunks from l."""
            for i in range(0, len(l), n):
                yield l[i:i + n]
        queue = Queue()
        for chunk in chunks(requests, int((len(requests) / n_chunks) + 1)):
            queue.put(chunk)
        processes = []
        try:
            for i in range(thread_count):
                p = backend(target=worker, args=(i, queue))
                processes.append(p)
                p.start()
        finally:
            for p in processes:
                p.join()

    def delete_request(self, request_id):
        self.status_coll.remove({"_id": self.get_request_id_key(request_id)})

    def uncommit_request(self, request_id):
        self.status_coll.replace_one({"_id": self.get_request_id_key(request_id)},
                                     {"crawler": self.crawler_name,
                                      "request_id": serialize_request_id(request_id),
                                      "requested_time": -1,
                                      "status": REDO},
                                     True
                                     )

    def delete_data(self, request_id, data_id, version=0):
        file_key = self._get_data_key(request_id, data_id, version)
        self.uncommit_request(request_id)
        self.s3_key_cache.remove({"_id": self._get_data_key(request_id, data_id, version)})
        self.s3.delete_object(Bucket=self.S3_BUCKET_NAME, Key=file_key)

        logging.warning("delete data id %s and uncommit request id %s to force re-crawl", data_id, request_id)

class SequentialCrawlDB(CrawlDB):
    def __init__(self, crawler_name, request_timeout, initial_request_id, max_request_id, mongo_db=None):
        super(SequentialCrawlDB, self).__init__(crawler_name, request_timeout, mongo_db)

        assert isinstance(initial_request_id, int)
        assert isinstance(max_request_id, int)

        self.initial_request_id = initial_request_id
        self.max_request_id = max_request_id

    def _divide_and_conquer_search(self, start_request_id=None, end_request_id=None):
        """

        :param start_request_id: the current id
        :param end_request_id: the max id
        :return: a request id that is not been crawled between the start_request_id (inclusive) and end_request_id (inclusive)
        """
        if start_request_id is None:
            start_request_id = self.initial_request_id
        if end_request_id is None:
            end_request_id = self.max_request_id

        for i in range(start_request_id, min(end_request_id + 1, start_request_id + 10)):
            if not self.is_crawled_or_being_crawled(i):
                return i
            if i == end_request_id:
                return None

        n_requests = self.status_coll.find({"_id": {"$gte": self.get_request_id_key(start_request_id),
                                                    "$lte": self.get_request_id_key(end_request_id)}}).count()
        if n_requests == start_request_id - end_request_id + 1:
            # no more gap!
            return None

        # do divide and counter
        middle = int(start_request_id + ((end_request_id - start_request_id) / 2))
        ret = self._divide_and_conquer_search(middle, end_request_id)
        if ret is not None:
            return ret
        else:
            ret = self._divide_and_conquer_search(start_request_id, middle)
            return ret

    def next_request_id(self):
        if self.initial_request_id is not None:
            ret = self._divide_and_conquer_search()
            self.initial_request_id = ret
            return ret
        else:
            return None

class SequentialTimeCrawlDB(CrawlDB):
    def __init__(self, crawler_name, request_timeout, initial_date: datetime,
                 time_interval: timedelta = timedelta(days=1), mongo_db=None):
        self.epoch = datetime(1900, 1, 1)
        self.initial_date = int((initial_date - self.epoch) / time_interval) * time_interval + self.epoch
        self.time_interval = time_interval
        super(SequentialTimeCrawlDB, self).__init__(crawler_name, request_timeout, mongo_db)

    def next_request(self, cur_request_date=None) -> datetime:
        if cur_request_date is None:
            cur_request_date = self.initial_date
        while self.is_crawled_or_being_crawled(cur_request_date):
            cur_request_date += self.time_interval
            if cur_request_date >= datetime.now():
                break
        return cur_request_date

class StringIDCrawlDB(CrawlDB):
    def __init__(self, crawler_name, request_timeout, mongo_db=None):
        super(StringIDCrawlDB, self).__init__(crawler_name, request_timeout, mongo_db)


import pytest
import sqlite3
#from talon import talon, init_refs
#from .helper_fns import fetch_correct_ID
@pytest.mark.integration

class TestMonoexonicGenomicEdgeCase(object):
    """ The aim here is to make sure that TALON does not apply the same type of
        edge constraint to monoexonic transcripts of monoexonic genes that it
        does to spliced alignments. This can result in spurious genomic transcripts.
    """
    def test_within_bounds(self):
        """ Example where the monoexonic start and end positions are within the
            allowable cutoff distances. Expectation is for the transcript to
            be called as known. TALON run is done as part of the build_dataabases
            step, so we simply check the results."""

        with sqlite3.connect("scratch/monoexon.db") as conn:
            conn.row_factory = sqlite3.Row
            cursor = conn.cursor()

            cursor.execute("""SELECT * FROM observed
                              WHERE read_name = "read_1" """)
            result = cursor.fetchone()
            print([x for x in result])
            assert result['gene_ID'] == 1
            assert result['transcript_ID'] == 1
            assert result['start_delta'] == 10
            assert result['end_delta'] == -91

    def test_out_of_bounds(self):
        """ Example where the distance at the ends is > the cutoff of 99 bp.
        """

        with sqlite3.connect("scratch/monoexon.db") as conn:
            conn.row_factory = sqlite3.Row
            cursor = conn.cursor()

            cursor.execute("""SELECT * FROM observed
                              WHERE read_name = "read_2" """)
            result = cursor.fetchone()
            print([x for x in result])
            assert result['gene_ID'] == 1
            assert result['transcript_ID'] == 1
            assert result['start_delta'] == 100
            assert result['end_delta'] == -401


# -*- coding: utf-8 -*-
from IndexerBase import IndexerBase
import numpy as np
import talib

class RSI(IndexerBase):
    indexer_name = 'RSI'
    indexer_name_list = ['RSI']
    default_para_dic = {
        'N': 5,
    }

    def __init__(self, raw_data, plt):
        super(RSI, self).__init__(raw_data, plt)
        self.indexer_name_list = ['RSI']  # MA的指标名和参数名都跟参数有关，所以要随参数进行设置
        self.indexer_color_dic = {
            'RSI': 'blue'
        }

    def calculate_indexer_value(self):
        n = self.para_dic['N']
        close_array = np.array(self.raw_data['close'].values, dtype='float')
        self.indexer_value_dic['RSI'] = talib.RSI(close_array, n)

    def draw_indexer(self):
        i = 0
        for indexer_name, values in self.indexer_value_dic.items():
            c = self.indexer_color_dic[indexer_name][0]
            self.plt_dic[indexer_name] = self.plt.plot(name=indexer_name, pen=c)
            self.plt_dic[indexer_name].setData(values)
            i += 1

    def re_draw_indexer(self):
        for pname, values in self.indexer_value_dic.items():
            self.plt_dic[pname].setData(values)

    def get_polar_value(self,start_pos, end_pos):
        max_v = max(self.indexer_value_dic['RSI'][start_pos:end_pos])
        min_v = min(self.indexer_value_dic['RSI'][start_pos:end_pos])
        return max_v, min_v


from typing import Union
from .elementdb import ElementDB

class Element(object):
    """ Element from periodic table

    Parameters
    ----------
    identifier : :obj:`str`, :obj:`int`
        Symbol, Name, or Atomic Number of element
    mass : :obj:`float`, optional
        Mass [amu] of element. Default is most common isotope atomic
        weight

    Examples
    --------
    Constructing a Helium Atom.

    >>> Element('He')
    <Element symbol:He name:Helium mass:4.00>

    >>> Element('Helium')
    <Element symbol:He name:Helium mass:4.00>

    >>> Element(2)
    <Element symbol:He name:Helium mass:4.00>

    >>> Element('He', 4.3)
    <Element symbol:He name:Helium mass:4.30>
    """

    def __init__(self, identifier: Union[str, int], mass: float = 0.0):
        """Initializes element from identifier and mass
        Identifier: str | int (symbol, name or atomic_number)"""
        element_dict = ElementDB.lookup(identifier)

        self._symbol = element_dict['symbol']
        self._name = element_dict['name']
        self._atomic_number = element_dict['z']
        self._mass = mass or element_dict['mass']

    def __eq__(self, other_element: object) -> bool:
        if not isinstance(other_element, Element):
            return False
        elif (self.symbol == other_element.symbol and
              self.name == other_element.name and
              self.atomic_number == other_element.atomic_number and
              self.mass == other_element.mass):
            return True
        else:
            return False

    def __repr__(self) -> str:
        return "<Element symbol:{} name:{} mass:{:2.2f}>".format(
            self.symbol, self.name, self.mass)

    def __hash__(self) -> int:
        return sum(hash(item) for item in [
            self._mass, self._symbol, self._name, self.atomic_number
        ])

    @property
    def symbol(self) -> str:
        """Element's atomic symbol"""
        assert isinstance(self._symbol, str)
        return self._symbol

    @property
    def name(self) -> str:
        """Element's formal name"""
        assert isinstance(self._name, str)
        return self._name

    @property
    def atomic_number(self) -> int:
        """Element's atomic number"""
        assert isinstance(self._atomic_number, int)
        return self._atomic_number

    @property
    def mass(self) -> float:
        """Element's mass"""
        assert isinstance(self._mass, float)
        return self._mass


from flask import Flask, request, jsonify, json
import requests

from apps import devices as dv
from apps import db_mgmt as dbm
from apps import app

@app.route('/')
def index():
    return ('Obiwan : Hello there!')

@app.route('/listnwkdev')
def discovery_device():
    resp = dv.get_dev_list()
    return resp

@app.route('/devlist')
def get_device_ip_mac():
    resp = dv.device_IPMAC_match()
    return resp

'''
curl --header "Content-Type: application/json" \
  --request POST \
  --data '{"MAC":"C4:4F:33:0C:AE:19","IP":"192.168.43.72", "NAME":"R0304 - EMERG", "BROKERIP":"192.168.43.200"}' \
  http://192.168.43.200:7343/bbmap
'''
@app.route('/bbmap', methods=['POST'])
def device_register():
    content = request.get_json(silent=True)
    if content != None:
        mac = content["MAC"]
        name = content["NAME"]
        ip = content["IP"]
        brokerip = content["BROKERIP"] #rpi ip
        building = content["BUILDING"]
        bed = content["BED"]
        dev_des = content["DEVICE_DES"]

        #modify data master sensor database on rpi by inserting blackbox mac
        url = "http://" + brokerip + ":7343/regdev"
        obj = {"MAC":mac, "NAME":name}
        print("Sendin to {}".format(url))
        try:
            requests.post(url, json=content)
        except requests.exceptions.RequestException as e:
            print(e)
            return json.dumps({'status':'ERROR'})

        #Black box Broker IP modify via http
        url = "http://" + ip + "/api/v1/system/config"
        obj = {"mqtt_uri":"mqtt://" + brokerip}
        print("Connect to bb")
        try:
            requests.post(url, json=obj, timeout = 1) #black box do not return timeout
        except requests.exceptions.RequestException as e:
            pass

        return json.dumps({'status':'OK'})
    else:
        return json.dumps({'status':'ERROR, EMPTY CONTENT'})

@app.route('/regdev', methods=['POST'])
def wifi_module_register():
    content = request.get_json(silent=True)
    mac = content["MAC"]
    name = content["NAME"]
    building = content["BUILDING"]
    bed = content["BED"]
    dev_des = content["DEVICE_DES"]

    result = dbm.register_new_device(mac, name)

    # attach new async task for db sync
    dbm.db_asycn_service_attach(dbm.remotedb_async_add, content)

    if result == True:
        return json.dumps({'status':'OK'})
    else:
        return json.dumps({'status':'ERROR'})


from __future__ import absolute_import
import datetime
import glob
import heapq
import os
import subprocess
import time

import six

import dependency_manager  # pylint: disable=import-error

from telemetry.internal.util import local_first_binary_manager

def _ParseCrashpadDateTime(date_time_str):
  # Python strptime does not support time zone parsing, strip it.
  date_time_parts = date_time_str.split()
  if len(date_time_parts) >= 3:
    date_time_str = ' '.join(date_time_parts[:2])
  return datetime.datetime.strptime(date_time_str, '%Y-%m-%d %H:%M:%S')

class MinidumpFinder():
  """Handles finding Crashpad/Breakpad minidumps.

  In addition to whatever data is expected to be returned, most public methods
  also return a list of strings. These strings are what would normally be
  logged, but returned in the list instead of being logged directly to help
  cut down on log spam from uses such as
  BrowserBackend.GetRecentMinidumpPathWithTimeout().
  """
  def __init__(self, os_name, arch_name):
    super().__init__()
    self._os = os_name
    self._arch = arch_name
    self._minidump_path_crashpad_retrieval = {}
    self._explanation = []

  def MinidumpObtainedFromCrashpad(self, minidump):
    """Returns whether the given minidump was found via Crashpad or not.

    Args:
      minidump: the path to the minidump being checked.

    Returns:
      False if the dump was found via globbing, else True.
    """
    # Default to Crashpad where we hope to be eventually
    return self._minidump_path_crashpad_retrieval.get(minidump, True)

  def GetAllCrashpadMinidumps(self, minidump_dir):
    """Returns all minidumps in the given directory findable by Crashpad.

    Args:
      minidump_dir: The directory to look for minidumps in.

    Returns:
      A tuple. The first element is a list of paths to minidumps found by
      Crashpad, or None in the event of an error. The second element is a list
      of strings containing logging information generated while searching for
      the minidumps.
    """
    self._explanation = ['Attempting to find all Crashpad minidump files for a '
                         'suspected Chrome crash.']
    return self._GetAllCrashpadMinidumps(minidump_dir), self._explanation

  def GetAllMinidumpPaths(self, minidump_dir):
    """Finds all minidumps for either Crashpad or Breakpad.

    If any Crashpad minidumps are found, only they will be returned. Otherwise,
    Breakpad minidumps will be returned.

    Args:
      minidump_dir: The directory to look for minidumps in.

    Returns:
      A tuple. The first element is a list of paths to the found minidumps, or
      None in the event of an error. The second element is a list of strings
      containing logging information generated while searching for the
      minidumps.
    """
    self._explanation = ['Attempting to find all minidump files for a '
                         'suspected Chrome crash. Crashpad minidumps will be '
                         'searched for first, falling back to Breakpad '
                         'minidumps if none are found.']
    return self._GetAllMinidumpPaths(minidump_dir), self._explanation

  def GetMostRecentMinidump(self, minidump_dir):
    """Finds the most recently created Crashpad or Breakpad minidump.

    Args:
      minidump_dir: The directory to look for minidumps in.

    Returns:
      A tuple. The first element is a path to the found minidump, or None if
      no minidumps were found. The second element is a list of strings
      containing logging information generated while searching for the
      minidumps.
    """
    self._explanation = ['Attempting to find the most recent minidump file for '
                         'a suspected Chrome crash. Crashpad minidumps will be '
                         'searched for first, falling back to Breakpad '
                         'minidumps if none are found.']
    return self._GetMostRecentMinidump(minidump_dir), self._explanation

  def _GetAllCrashpadMinidumps(self, minidump_dir):
    if not minidump_dir:
      self._explanation.append('No minidump directory provided. Likely '
                               'attempted to retrieve the Crashpad minidumps '
                               'after the browser was already closed.')
      return None
    try:
      crashpad_database_util = \
          local_first_binary_manager.GetInstance().FetchPath(
              'crashpad_database_util')
      if not crashpad_database_util:
        self._explanation.append('Unable to find crashpad_database_util. This '
                                 'is likely due to running on a platform that '
                                 'does not support Crashpad.')
        return None
    except dependency_manager.NoPathFoundError:
      self._explanation.append(
          'Could not find a path to look for crashpad_database_util.')
      return None

    self._explanation.append(
        'Found crashpad_database_util at %s' % crashpad_database_util)

    report_output = subprocess.check_output([
        crashpad_database_util, '--database=' + minidump_dir,
        '--show-pending-reports', '--show-completed-reports',
        '--show-all-report-info'])
    # This can be removed once fully switched to Python 3 by setting text=True
    # when calling check_output above.
    if not isinstance(report_output, six.string_types):
      report_output = report_output.decode('utf-8')

    last_indentation = -1
    reports_list = []
    report_dict = {}
    for report_line in report_output.splitlines():
      # Report values are grouped together by the same indentation level.
      current_indentation = 0
      for report_char in report_line:
        if not report_char.isspace():
          break
        current_indentation += 1

      # Decrease in indentation level indicates a new report is being printed.
      if current_indentation >= last_indentation:
        report_key, report_value = report_line.split(':', 1)
        if report_value:
          report_dict[report_key.strip()] = report_value.strip()
      elif report_dict:
        try:
          report_time = _ParseCrashpadDateTime(report_dict['Creation time'])
          report_path = report_dict['Path'].strip()
          reports_list.append((report_time, report_path))
        except (ValueError, KeyError) as e:
          self._explanation.append('Expected to find keys "Path" and "Creation '
                                   'time" in Crashpad report, but one or both '
                                   'don\'t exist: %s' % e)
        finally:
          report_dict = {}

      last_indentation = current_indentation

    # Include the last report.
    if report_dict:
      try:
        report_time = _ParseCrashpadDateTime(report_dict['Creation time'])
        report_path = report_dict['Path'].strip()
        reports_list.append((report_time, report_path))
      except (ValueError, KeyError) as e:
        self._explanation.append('Expected to find keys "Path" and "Creation '
                                 'time" in Crashpad report, but one or both '
                                 'don\'t exist: %s' % e)

    return reports_list

  def _GetAllMinidumpPaths(self, minidump_dir):
    reports_list = self._GetAllCrashpadMinidumps(minidump_dir)
    if reports_list:
      for report in reports_list:
        self._minidump_path_crashpad_retrieval[report[1]] = True
      return [report[1] for report in reports_list]
    self._explanation.append('No minidumps found via crashpad_database_util, '
                             'falling back to globbing for Breakpad '
                             'minidumps.')
    dumps = self._GetBreakpadMinidumpPaths(minidump_dir)
    if dumps:
      self._explanation.append('Found Breakpad minidump via globbing.')
      for dump in dumps:
        self._minidump_path_crashpad_retrieval[dump] = False
      return dumps
    self._explanation.append(
        'Failed to find any Breakpad minidumps via globbing.')
    return []

  def _GetMostRecentCrashpadMinidump(self, minidump_dir):
    reports_list = self._GetAllCrashpadMinidumps(minidump_dir)
    if reports_list:
      _, most_recent_report_path = max(reports_list)
      return most_recent_report_path

    return None

  def _GetBreakpadMinidumpPaths(self, minidump_dir):
    if not minidump_dir:
      self._explanation.append('Attempted to fetch Breakpad minidump paths '
                               'without a minidump directory. The browser was '
                               'likely closed before attempting to fetch.')
      return None
    return glob.glob(os.path.join(minidump_dir, '*.dmp'))

  def _GetMostRecentMinidump(self, minidump_dir):
    # Crashpad dump layout will be the standard eventually, check it first.
    crashpad_dump = True
    most_recent_dump = self._GetMostRecentCrashpadMinidump(minidump_dir)

    # Typical breakpad format is simply dump files in a folder.
    if not most_recent_dump:
      crashpad_dump = False
      self._explanation.append('No minidump found via crashpad_database_util, '
                               'falling back to globbing for Breakpad '
                               'minidump.')
      dumps = self._GetBreakpadMinidumpPaths(minidump_dir)
      if dumps:
        most_recent_dump = heapq.nlargest(1, dumps, os.path.getmtime)[0]
        if most_recent_dump:
          self._explanation.append('Found Breakpad minidump via globbing.')

    # As a sanity check, make sure the crash dump is recent.
    if (most_recent_dump and
        os.path.getmtime(most_recent_dump) < (time.time() - (5 * 60))):
      self._explanation.append(
          'Crash dump is older than 5 minutes. May not be correct.')

    self._minidump_path_crashpad_retrieval[most_recent_dump] = crashpad_dump
    return most_recent_dump


from PIL import Image, ImageTk
import customtkinter as ctk
import tkinter as tk
import EasyPySpin
import logging
import time
import cv2
import os

class CameraApp:
    def __init__(self, window, window_title):
        self.window = window
        self.window.title(window_title)

        # Initialize camera
        self.cap = EasyPySpin.VideoCapture(0)
        # Set resolution for the camera
        self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
        self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)

        # Set appearance mode and color theme
        ctk.set_appearance_mode("Dark")
        ctk.set_default_color_theme("blue")

        # Button to start/stop live video feed
        self.toggle_video_button = ctk.CTkButton(window, text="Start Live Feed", command=self.toggle_video_feed)
        self.toggle_video_button.pack(pady=20, padx=20)

        # Label to hold video frames
        self.video_label = ctk.CTkLabel(window, text="")
        self.video_label.pack()

        # Slider for exposure control
        self.exposure_slider = ctk.CTkSlider(window, from_=1, to=100000, command=self.adjust_exposure)
        self.exposure_slider.set(50)
        self.exposure_slider.pack(pady=20)

        # Exposure control label
        exposure = self.cap.get(cv2.CAP_PROP_EXPOSURE)
        self.exposure_label = ctk.CTkLabel(window, text=f"Exposure (us): {exposure}")
        self.exposure_label.pack(pady=(10, 0))

        # Control variable for video feed status
        self.is_live = False

        # Button for image saving
        self.record_button = ctk.CTkButton(window, text="Start Saving Images", command=self.toggle_image_saving)
        self.record_button.pack(pady=20)

        # Variables for image saving
        self.is_saving_images = False
        self.image_folder = "images"

        # Create the images directory if it doesn't exist
        if not os.path.exists(self.image_folder):
            os.makedirs(self.image_folder)

        self.window.protocol("WM_DELETE_WINDOW", self.on_closing)  # Handle window close event

        self.window.mainloop() # Start the GUI

    def adjust_exposure(self, exposure_value):
        try:
            exposure_us = float(exposure_value)  # Convert the slider value to a float

            self.cap.set(cv2.CAP_PROP_EXPOSURE, exposure_us)  # Set the exposure value for the camera

            self.exposure_label.configure(text=f"Exposure (us): {exposure_us}") # Update the exposure label

        except AttributeError:
            logging.error("Exposure not set.")

    def toggle_video_feed(self):
        if self.is_live:
            self.is_live = False
            self.toggle_video_button.configure(text="Start Live Feed")  # Corrected line
        else:
            self.is_live = True
            self.toggle_video_button.configure(text="Stop Live Feed")  # Corrected line
            self.update_video_label()  # Start video feed

    def toggle_image_saving(self):
            if self.is_saving_images:
                self.is_saving_images = False
                self.record_button.configure(text="Start Saving Images")
            else:
                self.is_saving_images = True
                self.record_button.configure(text="Stop Saving Images")

    def save_image(self, frame):
        timestamp = time.strftime("%Y%m%d-%H%M%S")
        filename = os.path.join(self.image_folder, f"image_{timestamp}.png")
        cv2.imwrite(filename, frame)

    def update_video_label(self):
        if self.is_live:
            ret, frame = self.cap.read()
            if ret:
                if self.is_saving_images:
                    self.save_image(frame)

                cv_img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                pil_img = Image.fromarray(cv_img)
                imgtk = ImageTk.PhotoImage(image=pil_img)
                self.video_label.imgtk = imgtk
                self.video_label.configure(image=imgtk)

            self.window.after(30, self.update_video_label)

    def on_closing(self):
        self.is_live = False
        self.cap.release()
        self.window.destroy()

# Create a window and pass it to the CameraApp class
root = tk.Tk()
app = CameraApp(root, "DART")


import sklearn
from sklearn.linear_model import SGDRegressor, SGDClassifier
from sklearn.neural_network import MLPRegressor, MLPClassifier
from sklearn.naive_bayes import MultinomialNB, GaussianNB
from sklearn.neighbors import KNeighborsRegressor, KNeighborsClassifier
from sklearn.ensemble import GradientBoostingRegressor, GradientBoostingClassifier
from sklearn.cluster import KMeans
from .mixed_nb import MixedNB
import logging
import numpy as np
import titus.prettypfa

def sklearn_to_pfa(estimator, types, featurizer=None):
    """
    Convert scikit-learn estimator to PFA format.
    :param estimator: Scikit-learn estimator, must be supported
    :param types: List of tuples (name, type)
    """
    types = _fix_types_compatibility(types)

    featurizer = featurizer or _construct_featurizer(types)

    if isinstance(estimator, SGDRegressor):
        return _pfa_sgdregressor(estimator, types, featurizer)
    elif isinstance(estimator, SGDClassifier):
        return _pfa_sgdclassifier(estimator, types, featurizer)
    elif isinstance(estimator, MLPRegressor):
        return _pfa_mlpregressor(estimator, types, featurizer)
    elif isinstance(estimator, MLPClassifier):
        return _pfa_mlpclassifier(estimator, types, featurizer)
    elif isinstance(estimator, MultinomialNB):
        return _pfa_multinomialnb(estimator, types, featurizer)
    elif isinstance(estimator, GaussianNB):
        return _pfa_gaussiannb(estimator, types, featurizer)
    elif isinstance(estimator, MixedNB):
        return _pfa_mixednb(estimator, types, featurizer)
    elif isinstance(estimator, KMeans):
        return _pfa_kmeans(estimator, types, featurizer)
    elif isinstance(estimator, KNeighborsRegressor):
        return _pfa_kneighborsregressor(estimator, types, featurizer)
    elif isinstance(estimator, KNeighborsClassifier):
        return _pfa_kneighborsclassifier(estimator, types, featurizer)
    elif isinstance(estimator, GradientBoostingRegressor):
        return _pfa_gradientboostingregressor(estimator, types, featurizer)
    elif isinstance(estimator, GradientBoostingClassifier):
        return _pfa_gradientboostingclassifier(estimator, types, featurizer)
    else:
        raise NotImplementedError('Estimator {} is not yet supported'.format(estimator.__class__.__name__))

def _pfa_sgdregressor(estimator, types, featurizer):
    """
    Example output:
        input: record(Data,
            feature0: double,
            feature1: double,
        )
        output: double
        action:
            -2.525797382870301 + 31.7004451488 * input.feature0 + 42.5005713274 * input.feature1
    """
    input_record = _input_record(types)

    # construct template
    pretty_pfa = """
types:
    Query = record(Query,
                   sql: string,
                   variable: string,
                   covariables: array(string));
    Regression = record(Regression, const: double, coeff: array(double));
    Input = {input_record}
input: Input
output: double
cells:
    // query(Query) = {{}};
    model(Regression) = {{const: 0.0, coeff: []}};
fcns:
    {functions}
action:
    var x = {featurizer};
    model.reg.linear(x, model)
    """.format(
        input_record=input_record, featurizer=featurizer, functions=_functions()
    ).strip()

    # compile
    logging.info(pretty_pfa)
    pfa = titus.prettypfa.jsonNode(pretty_pfa)

    # add model from scikit-learn
    pfa['cells']['model']['init'] = {'const': estimator.intercept_[0], 'coeff': list(estimator.coef_)}

    return pfa

def _pfa_sgdclassifier(estimator, types, featurizer):
    """
    Example output:
        types:
            Query = record(Query,
                           sql: string,
                           variable: string,
                           covariables: array(string));
            Input = record(Input,
                a: double,
                b: double
            );
            Regression = record(Regression, const: double, coeff: array(double));
        input: Input
        output: string
        cells:
            query(Query) = {sql: SELECT, variable: dep, covariables: [a, b, c]};
            model(array(Regression)) = [
                {const: 3, coeff: [1, 2]},
                {const: 3, coeff: [1, 4]}
            ];
            classes(array(string)) = [a, b, c];
        fcns:
            featurize = fcn(x: Input -> array(double))
                new(array(double),
                    x.a,
                    x.b
                )
        action:
            // TODO: define this as a function when functions start working with titus python 3
            var x = new(array(double),
                input.a,
                input.b
            );
            var scores = a.map(model, fcn(r: Regression -> double) model.reg.linear(x, r));
            classes[a.argmax(scores)]
    """
    input_record = _input_record(types)

    # construct template
    pretty_pfa = """
types:
    Query = record(Query,
                   sql: string,
                   variable: string,
                   covariables: array(string));
    Regression = record(Regression, const: double, coeff: array(double));
    Input = {input_record}
input: Input
output: string
cells:
    // query(Query) = {{}};
    model(array(Regression)) = [];
    classes(array(string)) = [];
fcns:
    {functions}
action:
    var x = {featurizer};
    var scores = a.map(model, fcn(r: Regression -> double) model.reg.linear(x, r));
    classes[a.argmax(scores)]
    """.format(
        input_record=input_record, featurizer=featurizer, functions=_functions()
    ).strip()

    # compile
    logging.info(pretty_pfa)
    pfa = titus.prettypfa.jsonNode(pretty_pfa)

    # add model from scikit-learn
    pfa['cells']['classes']['init'] = list(estimator.classes_)
    pfa['cells']['model']['init'] = [
        {
            'const': const,
            'coeff': list(coeff)
        } for const, coeff in zip(estimator.intercept_, estimator.coef_)
    ]

    return pfa

def _pfa_mlpregressor(estimator, types, featurizer):
    """
    See https://github.com/opendatagroup/hadrian/wiki/Basic-neural-network
    """
    input_record = _input_record(types)

    # construct template
    pretty_pfa = """
types:
    Query = record(Query,
                   sql: string,
                   variable: string,
                   covariables: array(string));
    Layer = record(Layer,
                   weights: array(array(double)),
                   bias: array(double));
    Input = {input_record}
input: Input
output: double
cells:
    neuralnet(array(Layer)) = [];
fcns:
{functions}
action:
    var x = {featurizer};
    var activation = model.neural.simpleLayers(x, neuralnet, fcn(x: double -> double) m.link.relu(x));
    activation[0]
    """.format(
        input_record=input_record, featurizer=featurizer, functions=_functions()
    ).strip()

    # compile
    logging.info(pretty_pfa)
    pfa = titus.prettypfa.jsonNode(pretty_pfa)

    # add model from scikit-learn
    # NOTE: `model.neural.simpleLayers` accepts transposed matrices
    pfa['cells']['neuralnet']['init'] = [
        {
            'bias': bias.tolist(),
            'weights': weights.T.tolist()
        } for bias, weights in zip(estimator.intercepts_, estimator.coefs_)
    ]

    return pfa

def _pfa_mlpclassifier(estimator, types, featurizer):
    """
    See https://github.com/opendatagroup/hadrian/wiki/Basic-neural-network
    """
    input_record = _input_record(types)

    # construct template
    pretty_pfa = """
types:
    Query = record(Query,
                   sql: string,
                   variable: string,
                   covariables: array(string));
    Layer = record(Layer,
                   weights: array(array(double)),
                   bias: array(double));
    Input = {input_record}
input: Input
output: string
cells:
    neuralnet(array(Layer)) = [];
    classes(array(string)) = [];
fcns:
{functions}
action:
    var x = {featurizer};
    var activations = model.neural.simpleLayers(x, neuralnet, fcn(x: double -> double) m.link.relu(x));
    classes[a.argmax(activations)]
    """.format(
        input_record=input_record, featurizer=featurizer, functions=_functions()
    ).strip()

    # compile
    logging.info(pretty_pfa)
    pfa = titus.prettypfa.jsonNode(pretty_pfa)

    # add model from scikit-learn
    pfa['cells']['classes']['init'] = list(estimator.classes_)
    # NOTE: `model.neural.simpleLayers` accepts transposed matrices
    pfa['cells']['neuralnet']['init'] = [
        {
            'bias': bias.tolist(),
            'weights': weights.T.tolist()
        } for bias, weights in zip(estimator.intercepts_, estimator.coefs_)
    ]

    return pfa

def _pfa_multinomialnb(estimator, types, featurizer):
    """
    See https://github.com/opendatagroup/hadrian/wiki/Basic-naive-bayes
    NOTE: in our use case we use mostly one-hot encoded variables, so using BernoulliNB might make
        more sense
    """
    input_record = _input_record(types)

    # construct template
    pretty_pfa = """
types:
    Query = record(Query,
                   sql: string,
                   variable: string,
                   covariables: array(string));
    Distribution = record(Distribution,
                        logLikelihoods: array(double),
                        logPrior: double);
    Input = {input_record}
input: Input
output: string
cells:
    model(array(Distribution)) = [];
    classes(array(string)) = [];
fcns:
{functions}
action:
    var x = {featurizer};

    var classLL = a.map(model, fcn(dist: Distribution -> double) {{
        model.naive.multinomial(x, dist.logLikelihoods) + dist.logPrior
    }});
    var norm = a.logsumexp(classLL);
    var probs = a.map(classLL, fcn(x: double -> double) m.exp(x - norm));
    classes[a.argmax(probs)]
    """.format(
        input_record=input_record, featurizer=featurizer, functions=_functions()
    ).strip()

    # compile
    logging.info(pretty_pfa)
    pfa = titus.prettypfa.jsonNode(pretty_pfa)

    # add model from scikit-learn
    logprior = np.maximum(estimator.class_log_prior_, -1e10)
    pfa['cells']['classes']['init'] = list(estimator.classes_)
    pfa['cells']['model']['init'] = [
        {
            'logLikelihoods': ll.tolist(),
            'logPrior': log_prior.tolist()
        } for log_prior, ll in zip(logprior, np.exp(estimator.feature_log_prob_))
    ]

    return pfa

def _pfa_gaussiannb(estimator, types, featurizer):
    """
    See https://github.com/opendatagroup/hadrian/wiki/Basic-naive-bayes
    """
    input_record = _input_record(types)

    # construct template
    pretty_pfa = """
types:
    Query = record(Query,
                   sql: string,
                   variable: string,
                   covariables: array(string));
    Distribution = record(Distribution,
                          stats: array(record(M, mean: double, variance: double)),
                          logPrior: double);
    Input = {input_record}
input: Input
output: string
cells:
    model(array(Distribution)) = [];
    classes(array(string)) = [];
fcns:
{functions}
action:
    var x = {featurizer};

    var classLL = a.map(model, fcn(dist: Distribution -> double) {{
      model.naive.gaussian(x, dist.stats) + dist.logPrior
    }});

    var norm = a.logsumexp(classLL);
    var probs = a.map(classLL, fcn(x: double -> double) m.exp(x - norm));
    classes[a.argmax(probs)]
    """.format(
        input_record=input_record, featurizer=featurizer, functions=_functions()
    ).strip()

    # compile
    logging.info(pretty_pfa)
    pfa = titus.prettypfa.jsonNode(pretty_pfa)

    # add model from scikit-learn
    prior = np.maximum(estimator.class_prior_, 1e-10)
    pfa['cells']['classes']['init'] = list(estimator.classes_)
    pfa['cells']['model']['init'] = [
        {
            'stats': [{
                'mean': m,
                'variance': s
            } for m, s in zip(means, sigmas)],
            'logPrior': np.log(prior).tolist()
        } for prior, means, sigmas in zip(prior, estimator.theta_, estimator.sigma_)
    ]

    return pfa

def _pfa_mixednb(estimator, types, featurizer):
    input_record = _input_record(types)

    # construct template
    pretty_pfa = """
types:
    Query = record(Query,
                   sql: string,
                   variable: string,
                   covariables: array(string));
    GaussianDistribution = record(GaussianDistribution,
                                  stats: array(record(M, mean: double, variance: double)));
    MultinomialDistribution = record(MultinomialDistribution,
                                     logLikelihoods: array(double));
    Input = {input_record}
input: Input
output: string
cells:
    gaussModel(array(GaussianDistribution)) = [];
    multinomialModel(array(MultinomialDistribution)) = [];
    classes(array(string)) = [];
    logPrior(array(double)) = [];
fcns:
{functions}
action:
    var x = {featurizer};

    var gaussFeatures = if( a.len(gaussModel) > 0 ) a.len(gaussModel[0,"stats"]) else 0;

    var gaussianLL = a.map(gaussModel, fcn(dist: GaussianDistribution -> double) {{
        model.naive.gaussian(a.subseq(x, 0, gaussFeatures), dist.stats)
    }});

    var multinomialLL = a.map(multinomialModel, fcn(dist: MultinomialDistribution -> double) {{
        model.naive.multinomial(a.subseq(x, gaussFeatures, a.len(x)), dist.logLikelihoods)
    }});

    var classLL = logPrior;
    if (a.len(gaussianLL) > 0)
        classLL = la.add(classLL, gaussianLL);
    if (a.len(multinomialLL) > 0)
        classLL = la.add(classLL, multinomialLL);

    var norm = a.logsumexp(classLL);
    var probs = a.map(classLL, fcn(x: double -> double) m.exp(x - norm));
    classes[a.argmax(probs)]
    """.format(
        input_record=input_record, featurizer=featurizer, functions=_functions()
    ).strip()

    # compile
    logging.info(pretty_pfa)
    pfa = titus.prettypfa.jsonNode(pretty_pfa)

    # add model from scikit-learn
    pfa['cells']['classes']['init'] = list(estimator.classes_)
    # avoid null values from log prior
    logprior = np.maximum(estimator.class_log_prior_, -1e10)
    pfa['cells']['logPrior']['init'] = logprior.tolist()

    # assumes that continuous features go before nominal ones
    if hasattr(estimator.gauss_nb, 'theta_'):
        pfa['cells']['gaussModel']['init'] = [
            {
                'stats': [{
                    'mean': m,
                    'variance': s
                } for m, s in zip(means, sigmas)]
            } for means, sigmas in zip(estimator.gauss_nb.theta_, estimator.gauss_nb.sigma_)
        ]
    if hasattr(estimator.multi_nb, 'feature_log_prob_'):
        pfa['cells']['multinomialModel']['init'] = [
            {
                'logLikelihoods': ll.tolist()
            } for ll in np.exp(estimator.multi_nb.feature_log_prob_)
        ]

    return pfa

def _pfa_kmeans(estimator, types, featurizer):
    input_record = _input_record(types)

    # construct template
    pretty_pfa = """
types:
    Query = record(Query,
                   sql: string,
                   variable: string,
                   covariables: array(string));
    Cluster = record(Cluster, center: array(double), id: int);
    Input = {input_record}
input: Input
output: int
cells:
    clusters(array(Cluster)) = [];
fcns:
{functions}
action:
    var x = {featurizer};
    var cluster = model.cluster.closest(x, clusters, metric.simpleEuclidean);
    cluster.id
    """.format(
        input_record=input_record, featurizer=featurizer, functions=_functions()
    ).strip()

    # compile
    logging.info(pretty_pfa)
    pfa = titus.prettypfa.jsonNode(pretty_pfa)

    # add model from scikit-learn
    pfa['cells']['clusters']['init'] = [
        {
            'center': c.tolist(),
            'id': i
        } for i, c in enumerate(estimator.cluster_centers_)
    ]
    return pfa

def _pfa_kneighborsregressor(estimator, types, featurizer):
    """See https://github.com/opendatagroup/hadrian/wiki/Basic-nearest-neighbors"""
    input_record = _input_record(types)

    # construct template
    pretty_pfa = """
types:
    Query = record(Query,
                   sql: string,
                   variable: string,
                   covariables: array(string));
    Point = record(Point,
                   x: array(double),
                   y: double);
    Codebook = array(Point);
    Input = {input_record}
input: Input
output: double
cells:
    codebook(Codebook) = [];
    nNeighbors(int) = 5;
fcns:
{functions}
action:
    var x = {featurizer};
    var neighbors = model.neighbor.nearestK(nNeighbors, x, codebook, fcn(x: array(double), p: Point -> double) {{
        metric.simpleEuclidean(x, p.x)
    }});
    a.mean(a.map(neighbors, fcn(p: Point -> double) p.y))
    """.format(
        input_record=input_record, featurizer=featurizer, functions=_functions()
    ).strip()

    # compile
    logging.info(pretty_pfa)
    pfa = titus.prettypfa.jsonNode(pretty_pfa)

    # add model from scikit-learn
    pfa['cells']['codebook']['init'] = [
        {
            'x': x.tolist(),
            'y': y
        } for x, y in zip(estimator._fit_X, estimator._y)
    ]
    pfa['cells']['nNeighbors']['init'] = estimator.n_neighbors

    return pfa

def _pfa_kneighborsclassifier(estimator, types, featurizer):
    """See https://github.com/opendatagroup/hadrian/wiki/Basic-nearest-neighbors"""
    input_record = _input_record(types)

    # construct template
    pretty_pfa = """
types:
    Query = record(Query,
                   sql: string,
                   variable: string,
                   covariables: array(string));
    Point = record(Point,
                   x: array(double),
                   y: string);
    Codebook = array(Point);
    Input = {input_record}
input: Input
output: string
cells:
    codebook(Codebook) = [];
    nNeighbors(int) = 5;
fcns:
{functions}
action:
    var x = {featurizer};
    var neighbors = model.neighbor.nearestK(nNeighbors, x, codebook, fcn(x: array(double), p: Point -> double) {{
        metric.simpleEuclidean(x, p.x)
    }});
    a.mode(a.map(neighbors, fcn(p: Point -> string) p.y))
    """.format(
        input_record=input_record, featurizer=featurizer, functions=_functions()
    ).strip()

    # compile
    logging.info(pretty_pfa)
    pfa = titus.prettypfa.jsonNode(pretty_pfa)

    # add model from scikit-learn
    y_labels = [estimator.classes_[e] for e in estimator._y]
    pfa['cells']['codebook']['init'] = [
        {
            'x': x.tolist(),
            'y': y
        } for x, y in zip(estimator._fit_X, y_labels)
    ]
    pfa['cells']['nNeighbors']['init'] = estimator.n_neighbors

    return pfa

def make_tree(tree, node_id=0):
    if tree.children_left[node_id] == sklearn.tree._tree.TREE_LEAF:
        leaf_value = float(tree.value[node_id][0, 0])

        # special case for empty tree with just root
        if node_id == 0:
            return {'TreeNode': {
                'feature': 0,
                'operator': '<=',
                'value': 0.,
                'pass': {'double': leaf_value},
                'fail': {'double': leaf_value}
            }}
        else:
            return {'double': leaf_value}

    return {'TreeNode': {
        'feature': int(tree.feature[node_id]),
        'operator': '<=',
        'value': float(tree.threshold[node_id]),
        'pass': make_tree(tree, tree.children_left[node_id]),
        'fail': make_tree(tree, tree.children_right[node_id])
    }}

def _pfa_gradientboostingregressor(estimator, types, featurizer):
    """See https://github.com/opendatagroup/hadrian/wiki/Basic-decision-tree"""
    input_record = _input_record(types)

    # construct template
    pretty_pfa = """
types:
    Query = record(Query,
                   sql: string,
                   variable: string,
                   covariables: array(string));
    TreeNode = record(TreeNode,
                    feature: int,
                    operator: string,
                    value: double,
                    pass: union(double, TreeNode),
                    fail: union(double, TreeNode));
    Row = record(Row, values: array(double));
    Input = {input_record}
input: Input
output: double
cells:
    // empty tree to satisfy type constraint; will be filled in later
    trees(array(TreeNode)) = [];
    // model intercept to which tree predictions are added
    intercept(double) = 0.0;
    learningRate(double) = 0.0;
fcns:
{functions}
action:
    var x = {featurizer};
    var row = new(Row, values: x);

    var scores = a.map(trees, fcn(tree: TreeNode -> double) {{
      model.tree.simpleWalk(row, tree, fcn(d: Row, t: TreeNode -> boolean) {{
        d.values[t.feature] <= t.value
      }})
    }});

    intercept + learningRate * a.sum(scores)
    """.format(
        input_record=input_record, featurizer=featurizer, functions=_functions()
    ).strip()

    # compile
    pfa = titus.prettypfa.jsonNode(pretty_pfa)

    # add model from scikit-learn
    tree_dicts = []
    for tree in estimator.estimators_[:, 0]:
        tree_dicts.append(make_tree(tree.tree_)['TreeNode'])
    pfa["cells"]["trees"]["init"] = tree_dicts
    pfa['cells']['intercept']['init'] = estimator.init_.mean
    pfa['cells']['learningRate']['init'] = estimator.learning_rate

    return pfa

def _pfa_gradientboostingclassifier(estimator, types, featurizer):
    """See https://github.com/opendatagroup/hadrian/wiki/Basic-decision-tree"""
    input_record = _input_record(types)

    # construct template
    pretty_pfa = """
types:
    Query = record(Query,
                   sql: string,
                   variable: string,
                   covariables: array(string));
    TreeNode = record(TreeNode,
                    feature: int,
                    operator: string,
                    value: double,
                    pass: union(double, TreeNode),
                    fail: union(double, TreeNode));
    Row = record(Row, values: array(double));
    Input = {input_record}
input: Input
output: string
cells:
    classes(array(string)) = [];
    // set of trees for each class
    classesTrees(array(array(TreeNode))) = [];
    // model priors to which tree predictions are added
    priors(array(double)) = [];
    learningRate(double) = 0.0;
fcns:
{functions}
action:
    var x = {featurizer};
    var row = new(Row, values: x);

    // trees activations
    var activations = a.map(classesTrees, fcn(trees: array(TreeNode) -> double) {{
        var scores = a.map(trees, fcn(tree: TreeNode -> double) {{
          model.tree.simpleWalk(row, tree, fcn(d: Row, t: TreeNode -> boolean) {{
            d.values[t.feature] <= t.value
          }})
        }});
        learningRate * a.sum(scores)
    }});

    // add priors
    activations = la.add(priors, activations);

    // probabilities
    var norm = a.logsumexp(activations);
    var probs = a.map(activations, fcn(x: double -> double) m.exp(x - norm));

    classes[a.argmax(probs)]
    """.format(
        input_record=input_record, featurizer=featurizer, functions=_functions()
    ).strip()

    # compile
    pfa = titus.prettypfa.jsonNode(pretty_pfa)

    # add model from scikit-learn
    tree_dicts = [[make_tree(tree.tree_)['TreeNode'] for tree in trees] for trees in estimator.estimators_.T]
    pfa["cells"]["classesTrees"]["init"] = tree_dicts
    pfa['cells']['classes']['init'] = list(estimator.classes_)
    pfa['cells']['priors']['init'] = list(estimator.init_.priors)
    pfa['cells']['learningRate']['init'] = estimator.learning_rate

    return pfa

def _construct_featurizer(types):
    inputs = []
    for name, typ in types:
        if typ == 'int':
            inputs.append('cast.double(input.{})'.format(name))
        else:
            inputs.append('input.{}'.format(name))

    return """
new(array(double),
    {inputs}
    )
    """.format(inputs=',\n'.join(inputs)).strip()

def _input_record(types):
    s = 'record(Input'
    for name, typ in types:
        s += ',\n    {}: {}'.format(name, typ)
    return s + '\n);'

def _fix_types_compatibility(types):
    new_types = []
    for name, typ in types:
        if typ == 'real':
            typ = 'double'
        elif typ == 'integer':
            typ = 'int'
        elif typ in ('polynominal', 'binominal'):
            typ = 'string'
        new_types.append((name, typ))
    return new_types

def _functions():
    return """
    arr = fcn(x: double -> array(double))
        new(array(double), x);

    C = fcn(x: string, categories: array(string) -> array(double))
        a.map(categories, fcn(cat: string -> double) if(cat == x) 1 else 0);

    standardize = fcn(x: double, mu: double, sigma: double -> double)
        (x - mu) / sigma;
""".strip()


from utils import directory_path, define_h2o_dataframe, load_pickle, dump_pickle
import shap
import pandas as pd
from matplotlib import pyplot as plt

def shap_visualisation_xgboost(version:str, point:int, do_train=False):
    """This function do the shap visualisation

    :param version: version
    :type version: str
    :param point: point to visualise
    :type point: int
    :param do_train:  do shap value and tree explainer
    :type do_train: bool
    :param point: point to visualize
    :type point: int
    """

    test = pd.read_csv(f"{directory_path}data/processed/application_test.csv")
    model = load_pickle(f"{directory_path}/models/{version}/xgboostclassifier/xgboostclassifier")

    # understanding model predictions
    X = test
    if do_train:
        explainer = shap.TreeExplainer(model)
        shap_values = explainer.shap_values(X)
        dump_pickle(f"{directory_path}visualisation/xgboost/explainer", explainer)
        dump_pickle(f"{directory_path}visualisation/xgboost/shap_values", shap_values)
    else:
        shap_values = load_pickle(f"{directory_path}visualisation/xgboost/shap_values")
        explainer = load_pickle(f"{directory_path}visualisation/xgboost/explainer")

    image_name = f"{directory_path}visualisation/xgboost/plots/plot"
    plots_shap(image_name, shap_values, explainer, X, point)

def plots_shap(image_name: str, shap_values, explainer, X: pd.DataFrame, point:int):
    """This function do several visualisations

    :param image_name: path to the image
    :type image_name: str
    :param shap_values: shap value
    :type shap_values:
    :param explainer: explainer
    :type explainer:
    :param X: test set
    :type X: pd.DataFrame
    :param point: point to visualize
    :type point: int
    """
    # Visualize explanations for a specific point of your data set,
    save_plot_specific_point(f'{image_name}_specific_point_{point}.png', X,
                             explainer.expected_value, shap_values, point)

    save_plot_local_waterfall(f'{image_name}_waterfall_point_{point}.png', X,
                              explainer.expected_value, shap_values, point)

    # Visualize explanations for all points of your data set at once,
    save_plot_all_points(f'{image_name}_allpoint.png', X,
                         explainer.expected_value, shap_values)

    # Visualize a summary plot for each class on the whole dataset.
    save_plot_all_featureimportance(f'{image_name}_summary.png', X, shap_values)
    save_plot_summary(f'{image_name}_summary_bar.png', X, shap_values)

def save_plot_specific_point(image_name: str, set: pd.DataFrame, expected_value, shap_values, row: int):
    """This function do visualisation on a single points and
    store the result.

    :param image_name: path to the image
    :type image_name: str
    :param set: test set
    :type set: pd.DataFrame
    :param expected_value: expected_value
    :type expected_value:
    :param shap_values: shap value
    :type shap_values:
    :param row: point to visualize
    :type row: int
    """
    point = shap.force_plot(expected_value, shap_values[row], features=set.iloc[-1:],
                                feature_names=list(set.columns),
                                matplotlib=True, show=False
                        )
    plt.savefig(image_name,  bbox_inches='tight')
    plt.close()

def save_plot_local_waterfall(image_name: str, set: pd.DataFrame, expected_value, shap_values, row: int):
    """This function do visualisation on a single points using waterfall and
    store the result.

    :param image_name: path to the image
    :type image_name: str
    :param set: path to the image
    :type set: pd.DataFrame
    :param expected_value: expected_value
    :type expected_value:
    :param shap_values: shap value
    :type shap_values:
    :param row: point to visualize
    :type row: int
    """
    fig=plt.gcf()
    shap.plots._waterfall.waterfall_legacy(expected_value, shap_values[row],
                                                   feature_names=set.columns,
                                                   )
    fig.savefig(image_name,  bbox_inches='tight')

def save_plot_all_points(image_name:str, set: pd.DataFrame, expected_value, shap_values):
    """This function do visualisation on all points and
    store the result.

    :param image_name: path to the image
    :type image_name: str
    :param set: test set
    :type set: pd.DataFrame
    :param expected_value: expected_value
    :type expected_value:
    :param shap_values: shap value
    :type shap_values:
    """
    shap_plot = shap.force_plot(expected_value,
                                shap_values[-1:], features=set.iloc[-1:],
                                feature_names=list(set.columns),
                                matplotlib=True, show=False)

    plt.savefig(image_name,  bbox_inches='tight')
    plt.close()

def save_plot_all_featureimportance(image_name, set, shap_values):
    """This function do feature importance plot and
    store the result.

    :param image_name: path to the image
    :type image_name: str
    :param set: test set
    :type set: pd.DataFrame
    :param shap_values: shap value
    :type shap_values:
    """
    fig = plt.gcf()
    shap.summary_plot(shap_values, set.values, class_names=[0,1], feature_names=set.columns
                             )
    fig.savefig(image_name,  bbox_inches='tight')

def save_plot_summary(image_name:str, set: pd.DataFrame, shap_values):
    """This function do summary bar plot and
    store the result.

    :param image_name: path to the image
    :type image_name: str
    :param set: test set
    :type set: pd.DataFrame
    :param shap_values: shap value
    :type shap_values:
    """
    fig = plt.gcf()
    shap.summary_plot(shap_values, set.values, plot_type="bar",
                             class_names=[0,1], feature_names=set.columns)
    fig.savefig(image_name,  bbox_inches='tight')


import os
import argparse
import numpy as np
import trimesh

from scripts.helper import sample_mesh, visualize_pc

class PointCloud:
    def __init__(self, model_name):
        self.model_name = model_name.split('.')[0] + '.ply'

    def sample(self, num_points=40960, centralize=False):
        # sample points on the mesh
        mesh_path = os.path.join(mesh_regions_dir, self.model_name)
        mesh = trimesh.load(mesh_path, process=False)
        try:
            mesh.area_faces
        except AttributeError:
            return None

        # bring the pc to the center if necessary
        if centralize:
            mesh.vertices -= np.mean(mesh.vertices, axis=0)

        # find the number of points to be sampled.
        pc, _ = sample_mesh(mesh, num_points=num_points)

        return pc

def derive_pc(args, region_names):
    for region_name in region_names:
        visited = set(os.listdir(results_dir))
        region_name = region_name.split('.')[0] + '.npy'
        if region_name not in visited:
            pc_object = PointCloud(region_name)
            # sample point clouds
            pc = pc_object.sample(num_points=args.num_points, centralize=False)
            if pc is not None:
                # visualize_pc(pc)

                # save point clouds and labels
                output_path = os.path.join(results_dir, region_name)
                np.save(output_path, pc)

                visited.add(region_name)

def build_colored_pc(pc):
    radii = np.linalg.norm(pc, axis=1)
    colors = trimesh.visual.interpolate(radii, color_map='viridis')
    pc_vis = trimesh.points.PointCloud(pc, colors=colors)

    return pc_vis

def build_cube(center, pc_extents, coverage_percent):
    # find the scale of the cube
    scale = pc_extents * coverage_percent
    # build the cube
    cube = {'extents': scale, 'center': center}

    return cube

def sample_cube_centers(pc, pc_extents):
    denominator = 4
    is_center = []
    while np.sum(is_center) == 0:
        indices = np.arange(len(pc))
        is_center = np.abs(pc) <= (1/denominator * pc_extents / 2)
        is_center = np.sum(is_center, axis=1) == 3
        denominator -= 1

    return indices[is_center]

def find_sampled_crop(pc, cube, num_points):
    # take points inside the cube.
    is_inside = np.abs(pc - cube['center']) <= (cube['extents'] / 2.0)
    is_inside = np.sum(is_inside, axis=1) == 3
    crop = pc[is_inside, :]

    sampled_crop = None
    if len(crop) >= num_points:
        sampled_indices = np.random.choice(range(len(crop)), num_points, replace=False)
        sampled_crop = crop[sampled_indices, :]

    return sampled_crop

def sample_crop_pc(pc_file_names, crop_bounds, num_points, num_tries=200):
    # for each object sample points and extract a crop.
    coverage_percent = np.random.uniform(crop_bounds[0], crop_bounds[1], 3)
    for pc_file_name in pc_file_names:
        # load the pc and sample the center of the cube near the center of the object.
        input_path = os.path.join(pc_dir, pc_file_name)
        pc = np.load(input_path)
        pc_extents = trimesh.points.PointCloud(pc).extents
        sampled_indices = sample_cube_centers(pc, pc_extents)

        # try at least x times to get enough points.
        num_curr_tries = np.minimum(num_tries, len(sampled_indices))
        sampled_indices = np.random.choice(sampled_indices, num_curr_tries, replace=False)

        # iterate through the points and exit once you find enough points in the crop.
        sampled_crop = None
        for sampled_index in sampled_indices:
            center = pc[sampled_index, :]
            cube = build_cube(center, pc_extents, coverage_percent)

            # find the points inside the cube crop
            sampled_crop = find_sampled_crop(pc, cube, num_points)

            # if you find enough points, you found the crop so exit.
            if sampled_crop is not None:
                break

        # return the entire crop if sampling was not successful.
        if sampled_crop is None:
            print('Cropping not successful. Taking the entire pc')
            sampled_indices = np.random.choice(range(len(pc)), num_points, replace=False)
            sampled_crop = pc[sampled_indices, :]

        # build_colored_pc(pc).show()
        # build_colored_pc(sampled_crop).show()

        # save the crop
        output_path = os.path.join(results_dir, pc_file_name)
        np.save(output_path, sampled_crop)

def get_args():
    parser = argparse.ArgumentParser('Extract Regions', add_help=False)
    parser.add_argument('--action', default='extract', help='extract | sample_crop')
    parser.add_argument('--dataset', default='matterport3d')
    parser.add_argument('--mode', default='train', help='train | val | test')
    parser.add_argument('--mesh_regions_dir', default='../data/{}/mesh_regions')
    parser.add_argument('--pc_dir', default='../data/{}/pc_regions')
    parser.add_argument('--results_dir', default='../data/{}/pc_regions')
    parser.add_argument('--num_points', default=4096, type=int, help='number of points randomly sampled form the pc.')
    parser.add_argument('--crop_bounds', type=float, nargs='+', default=(0.9, 0.9))
    parser.add_argument('--seed', default=0, type=int, help='use different seed for parallel runs')
    parser.add_argument('--num_chunks', default=1, type=int, help='number of chunks for parallel run')
    parser.add_argument('--chunk_idx', default=0, type=int, help='chunk id for parallel run')

    return parser

def adjust_paths(args, exceptions):
    base_dir = os.path.dirname(os.path.abspath(__file__))
    for k, v in vars(args).items():
        if (type(v) is str) and ('/' in v) and k not in exceptions:
            v = v.format(args.dataset)
            vars(args)[k] = os.path.join(base_dir, v)

def main():
    region_names = os.listdir(mesh_regions_dir)
    np.random.seed(args.seed)
    if args.action == 'extract':
        np.random.shuffle(region_names)
        chunk_size = int(np.ceil(len(region_names) / args.num_chunks))
        derive_pc(args, region_names=region_names[args.chunk_idx * chunk_size: (args.chunk_idx + 1) * chunk_size])
    elif args.action == 'sample_crop':
        pc_file_names = os.listdir(pc_dir)
        sample_crop_pc(pc_file_names, crop_bounds=args.crop_bounds, num_points=args.num_points, num_tries=200)

if __name__ == '__main__':
    # read the args
    parser_ = argparse.ArgumentParser('Extract Regions', parents=[get_args()])
    args = parser_.parse_args()
    adjust_paths(args, exceptions=[])

    # load metadata
    mesh_regions_dir = os.path.join(args.mesh_regions_dir, args.mode)
    pc_dir = os.path.join(args.pc_dir, args.mode)
    results_dir = os.path.join(args.results_dir, args.mode)
    if not os.path.exists(results_dir):
        try:
            os.makedirs(results_dir)
        except FileExistsError:
            pass

    main()
    # To run in parallel you can use the command:
    # parallel -j5 "python3 -u extract_point_clouds.py --mode {1} --seed {2} --num_chunks {3} --chunk_idx {4}" ::: test ::: 0 ::: 5 ::: 0 1 2 3 4
    # parallel -j5 "python3 -u extract_point_clouds.py --mode {1} --results_folder_name {2} --seed {3} --num_chunks {4} --chunk_idx {5}" ::: test ::: crops_1 ::: 0 ::: 5 ::: 0 1 2 3 4


from copy import deepcopy
import numpy as np
import os
from datetime import datetime

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.tensorboard import SummaryWriter

from dataloaders import get_cifar
from utils import get_activation_function, plot_grid, make_folder
from models.classifiers import evaluate_classifier, train_classifier
from models.autoencoders import Autoencoder
from metrics import compute_fid, compute_confidence_diversity, compute_inception_score
from arguments import get_args_parser

# replace all occurences of comma followed by non-white space character with comma+space: ,(?=[^\s])

def train_one_epoch(args, model, epoch=None, beta=None, train_dataloader=None, optim=None, lr_scheduler=None, device=None, writer=None):
    model.train()
    total_train_rec_loss = 0
    total_train_kl_loss = 0
    count = 0

    for image, label in train_dataloader:
        image = image.to(device)
        reconstructed = model(image)

        train_rec_loss = F.mse_loss(reconstructed, image)
        total_train_rec_loss += train_rec_loss
        train_loss = train_rec_loss

        if args.variational:
            mu = model.stats[:, :, 0]
            log_var = model.stats[:, :, 1]
            sigma = torch.exp(0.5 * log_var)
            if args.mse:
                mu_loss = F.mse_loss(mu, torch.zeros_like(mu), reduction='none')
                sigma_loss = F.mse_loss(sigma, torch.ones_like(sigma), reduction='none')
                train_kl_loss = mu_loss.sum(1).mean(0) + sigma_loss.sum(1).mean(0)
            else:
                # explained in https://arxiv.org/abs/1906.02691
                train_kl_loss = torch.mean(-0.5 * torch.sum(1 + log_var - mu ** 2 - log_var.exp(), dim = 1), dim = 0)

            total_train_kl_loss += beta * train_kl_loss
            train_loss += beta * train_kl_loss

        optim.zero_grad()
        train_loss.backward()
        optim.step()
        lr_scheduler.step()

        if writer is not None:
            i = epoch * len(train_dataloader) + count
            writer.add_scalar('Train Loss/KL', train_kl_loss, i)
            writer.add_scalar('Train Loss/Reconstruction', train_rec_loss, i)
            writer.add_scalar('Train Loss/Total', train_loss, i)
            writer.add_scalar('Train Loss/beta', beta, i)

        count += 1

    # to combat vanishing KL loss:
    if args.variational:
        beta *= args.beta_mult

    avg_train_rec_loss = total_train_rec_loss / len(train_dataloader)
    avg_train_kl_loss = total_train_kl_loss / len(train_dataloader)
    return beta, avg_train_rec_loss, avg_train_kl_loss

def compute_test_loss(args, model, beta=None, test_dataloader=None, device=None):
    # compute test loss (on test dataset)
    model.eval()
    with torch.no_grad():
        total_test_rec_loss = 0
        total_test_kl_loss = 0
        latent_mean = 0
        latent_std = 0

        for test_image, label in test_dataloader:
            test_image = test_image.to(device)
            test_reconstructed = model(test_image)

            test_rec_loss = F.mse_loss(test_reconstructed, test_image)
            total_test_rec_loss += test_rec_loss

            if args.variational:
                mu = model.stats[:, :, 0]
                log_var = model.stats[:, :, 1]
                sigma = torch.exp(0.5 * log_var)
                if args.mse:
                    mu_loss = F.mse_loss(mu, torch.zeros_like(mu), reduction='none')
                    sigma_loss = F.mse_loss(sigma, torch.ones_like(sigma), reduction='none')
                    test_kl_loss = mu_loss.sum(1).mean(0) + sigma_loss.sum(1).mean(0)
                else:
                    test_kl_loss = torch.mean(-0.5 * torch.sum(1 + log_var - mu ** 2 - log_var.exp(), dim = 1), dim = 0)

                total_test_kl_loss += test_kl_loss
                latent_mean += mu.abs().mean()
                latent_std += sigma.abs().mean()

        num_test_batches = len(test_dataloader)
        avg_test_kl_loss = beta * total_test_kl_loss / num_test_batches
        avg_test_rec_loss = total_test_rec_loss / num_test_batches
        avg_latent_mean = latent_mean / num_test_batches
        avg_latent_std = latent_std / num_test_batches

        return avg_latent_mean, avg_latent_std, avg_test_rec_loss, avg_test_kl_loss

def compute_metrics(autoencoder=None, classifier=None, real_images=None, samples=None, debug=False):
    #fid_ae_dynamic = compute_fid2(model=model_we_are_training, images1=real_images, images2=samples)
    fid_ae = compute_fid(model=autoencoder, images1=real_images, images2=samples)  # use pretrained cifar model to compute features
    fid_classifier = compute_fid(model=classifier, images1=real_images, images2=samples)
    confidence, diversity = compute_confidence_diversity(classifier, samples, debug=debug)
    inception_score = compute_inception_score(classifier=classifier, images=samples)
    return fid_classifier, fid_ae, confidence, diversity, inception_score

def main(args):
    # create necessary directories as needed
    for dir_name in ['checkpoints', 'data', 'plots', 'logs']:
        make_folder(dir_name)

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

    train_dataloader, test_dataloader = get_cifar(
        data_dir=args.data_dir,
        train_batch_size=args.train_batch_size,
        test_batch_size=args.test_batch_size
        )

    # load images from disk for plotting (otherwise train images are randomly picked every time)
    if os.path.isfile('data/train_examples.pth') and os.path.isfile('data/test_examples.pth'):
        train_input_images = torch.load('data/train_examples.pth')
        test_input_images = torch.load('data/test_examples.pth')
    else:
        train_input_images = next(iter(train_dataloader))[0]
        test_input_images = next(iter(test_dataloader))[0]
        torch.save(train_input_images, 'data/train_examples.pth')
        torch.save(test_input_images, 'data/test_examples.pth')

    if test_input_images.shape[0] < args.latent_size:  # number of images should be >= number of features (to compute FID)
        print(f'\n\nNumber of test images {test_input_images.shape[0]} should be >= number of '
            f'features {args.latent_size} in order to compute FID properly (only applies to VAE)\n\n')
        if args.variational:
            raise(SystemExit)

    args.act_fn = get_activation_function(act_str=args.act)

    if args.checkpoint is not None:
        print(f'\n\nLoading model checkpoint from {args.checkpoint}')
        checkpoint = torch.load(args.checkpoint)
        saved_args = checkpoint['args']

        print(f'\n\nCurrent arguments:\n')
        for current_arg in vars(args):
            print(current_arg, getattr(args, current_arg))
        print(f'\n\nCheckpoint arguments:\n')
        for saved_arg in vars(saved_args):
            print(saved_arg, getattr(saved_args, saved_arg))

    model = Autoencoder(
        variational=args.variational,
        latent_size=args.latent_size,
        num_channels=args.num_channels,
        kernel_size=args.kernel_size,
        act=args.act_fn,
        sigmoid=args.sigmoid,
        interpolation=args.interpolation,
        no_pool=args.no_pool,
        no_upsample=args.no_upsample,
        device=device,
        debug=args.debug,
        ).to(device)

    # train see https://github.com/orybkin/sigma-vae-pytorch

    beta = args.beta
    init_epoch = 0

    if args.variational:
        model_type = f'vae_{args.beta}x{args.beta_mult}'
        if os.path.isfile(args.ae_checkpoint):
            print(f'\n\nLoading pretrained VAE for FID computation from {args.ae_checkpoint}')
            plain_ae = torch.load(args.ae_checkpoint)
        else:
            print(f'\n\nTo compute FID score using a plain autoencoder, provide valid path to checkpoint, or train one:')
            print(f'\n\npython main.py --latent_size 256 --sigmoid --wd 0.01 --epochs 100 --train --no_upsample --no_pool\n\n')
            raise(SystemExit)

        if args.classifier_checkpoint is None:
            classifier_checkpoint = f'checkpoints/cifar_classifier_{args.latent_size}.pth'
        print(f'\n\nInstantiating Classifier for FID computation')
        if os.path.isfile(classifier_checkpoint):
            print(f'\n\nFound checkpoint at {classifier_checkpoint}, loading...')
            classifier = torch.load(classifier_checkpoint)
            test_acc = evaluate_classifier(classifier, test_dataloader=test_dataloader, device=device)
            print(f'\n\nCIFAR-10 Test Accuracy {test_acc:.2f}')
        else:
            print(f'\n\nClassifier checkpoint is not found at {classifier_checkpoint}')
            classifier = train_classifier(args, classifier=None, train_dataloader=train_dataloader, test_dataloader=test_dataloader, device=device)

        if isinstance(plain_ae, dict) or isinstance(classifier, dict):
            raise NotImplementedError('\n\nmodel checkpoint must be a full saved model, not a state_dict\n\n')
    else:
        model_type = 'plain-ae'

    optim = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.wd)
    lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optim, T_max=args.epochs * len(train_dataloader))

    if args.checkpoint is not None:
        model.load_state_dict(checkpoint['state_dict'])
        optim.load_state_dict(checkpoint['optimizer'])
        lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
        init_epoch = checkpoint['epoch']
        beta = checkpoint['current_beta']

    pool_str = 'pool-stride' if args.no_pool else 'pool-max'
    upsample_str = 'upsample-deconv' if args.no_upsample else f'upsample-{args.interpolation}'

    experiment_str = args.tag + f'{model_type}_latent{args.latent_size}_chan{args.num_channels}_{pool_str}_{upsample_str}_bs{args.train_batch_size}_lr{args.lr}_wd{args.wd}_e{args.epochs}'
    print(f'\n\n{experiment_str}')

    if args.log_tb:
        # launch with 'tensorboard --host 0.0.0.0 --logdir args.tb_dir' on the gateway and go to gateway_ip:6006
        writer = SummaryWriter(args.tb_dir + '/' + experiment_str)
        # TODO should we save this writer to checkpoints?
    else:
        writer = None

    if args.evaluate:
        print(f'\n\nEvaluating model')
        plot_grid(model, train_input_images[:4], name=experiment_str+'_train')
        plot_grid(model, test_input_images[:4], name=experiment_str+'_test')
        print(f'\n\nPlots saved to plots/{experiment_str}_train-test.png\n\n')

    if args.sample:
        print(f'\n\nSampling 8 images from a random latent vector')
        model.sample(name=experiment_str)
        print(f'\n\nPlot saved to plots/samples_{experiment_str}.png\n\n')

    if args.train:
        if args.variational:
            fid_classifier, fid_ae, confidence, diversity, inception_score = compute_metrics(
                autoencoder=plain_ae,
                classifier=classifier,
                real_images=train_input_images,
                samples=test_input_images,
                debug=True,
            )
            ref_metrics_str = f'Reference (real images) metrics:  FID (classifier/AE): ' + \
                            f'{fid_classifier:.1f}/{fid_ae:.1f}  confidence {confidence:.1f}  diversity {diversity:.1f}  IS {inception_score:.3f}\n\n'
        else:
            ref_metrics_str = ''

        print(f'\n\nTraining {model_type} model on CIFAR-10 images\n\n{ref_metrics_str}')

        for epoch in range(init_epoch, args.epochs, 1):
            beta, avg_train_rec_loss, avg_train_kl_loss = train_one_epoch(
                args,
                model,
                epoch=epoch,
                beta=beta,
                train_dataloader=train_dataloader,
                optim=optim,
                lr_scheduler=lr_scheduler,
                device=device,
                writer=writer,
            )

            avg_latent_mean, avg_latent_std, avg_test_rec_loss, avg_test_kl_loss = compute_test_loss(
                args,
                model,
                beta=beta,
                test_dataloader=test_dataloader,
                device=device,
            )
            avg_test_loss = avg_test_rec_loss + avg_test_kl_loss if args.variational else avg_test_rec_loss

            if args.variational:
                samples = model.sample(num_images=args.test_batch_size, return_samples=True)
                fid_classifier, fid_ae, confidence, diversity, inception_score = compute_metrics(
                    autoencoder=plain_ae,
                    classifier=classifier,
                    real_images=test_input_images,
                    samples=samples,
                    debug=True,
                )
                metrics_str = f'  FID (classifier/AE): {fid_classifier:.1f}/{fid_ae:.1f}  confidence {confidence:.1f}  diversity {diversity:.1f}  IS {inception_score:.2f}'
            else:
                metrics_str = ''

            if writer is not None:
                i = epoch * len(train_dataloader)
                writer.add_scalar('Test Loss/KL', avg_test_kl_loss, i)
                writer.add_scalar('Test Loss/Reconstruction', avg_test_rec_loss, i)
                writer.add_scalar('Test Loss/Total', avg_test_loss, i)
                writer.add_scalar('Test Loss/beta', beta, i)
                writer.add_scalar('Latent Vector/mean', avg_latent_mean, i)
                writer.add_scalar('Latent Vector/std', avg_latent_std, i)
                writer.add_scalar('FID/classifier', fid_classifier, i)
                writer.add_scalar('FID/plain_ae', fid_ae, i)
                writer.add_scalar('Metrics/confidence', confidence, i)
                writer.add_scalar('Metrics/diversity', diversity, i)
                writer.add_scalar('Metrics/inception score', inception_score, i)

            latent_stats_str = f'  mean {avg_latent_mean:.4f} std {avg_latent_std:.4f}' if args.variational else ''
            kl_loss_str = f'  kl train {1000*avg_train_kl_loss:.2f} test {1000*avg_test_kl_loss:.2f}' if args.variational else ''
            loss_str = f'losses:  train {1000*avg_train_rec_loss:.2f} test {1000*avg_test_rec_loss:.2f}{kl_loss_str}'
            changes_str = f'  LR {lr_scheduler.get_last_lr()[0]:.5f}' + (f' beta {beta:.6f}' if args.variational else '')
            time_str = f'{str(datetime.now())[:-7]}'
            print(f'{time_str}  Epoch {epoch:>3d}   {loss_str}{latent_stats_str}{metrics_str}{changes_str}')

            plot_grid(model, train_input_images[:4], name=experiment_str+'_train')
            plot_grid(model, test_input_images[:4], name=experiment_str+'_test')

            # checkpointing:
            checkpoint = {}
            checkpoint['state_dict'] = model.state_dict()
            checkpoint['optimizer'] = optim.state_dict()
            checkpoint['lr_scheduler'] = lr_scheduler.state_dict()
            checkpoint['epoch'] = epoch
            checkpoint['args'] = args
            checkpoint['current_beta'] = beta
            checkpoint['losses'] = loss_str

            path = 'checkpoints/' + experiment_str + ".pth"
            torch.save(checkpoint, path)

            model.sample(name=experiment_str)

        if not os.path.isfile(args.ae_checkpoint):  # save full model checkpoint
            torch.save(model, 'checkpoints/' + experiment_str + "_full_model.pth")

if __name__ == "__main__":
    # use below for inline vscode cell execution
    # parser = get_args_parser()
    # args, unknown = parser.parse_known_args()
    args = get_args_parser().parse_args()
    main(args)

    # TODO:
    # 2. Inception Score - DONE
    # 3. Tensorboard support - DONE
    # 5. save logs of train output
    # 6. add CelebA dataset

    # add vae-vq
    # implement MUSIQ
    # add larger models (resnets)