# Copyright 1999-2020 Alibaba Group Holding Ltd.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import functools
import operator
import logging
import sys
import time
from collections import defaultdict
from enum import Enum

from .. import promise
from ..config import options
from ..errors import DependencyMissing, ExecutionInterrupted, WorkerDead
from ..serialize import dataserializer
from ..scheduler.chunkmeta import WorkerMeta
from ..utils import log_unhandled, build_exc_info
from .events import EventContext, EventCategory, EventLevel, ProcedureEventType
from .storage import DataStorageDevice
from .utils import WorkerActor

logger = logging.getLogger(__name__)

class ReceiveStatus(Enum):
    NOT_STARTED = 0
    PENDING = 1
    RECEIVING = 2
    RECEIVED = 3
    ERROR = 4

class EndpointTransferState(object):
    """
    Structure providing transfer status in an endpoint
    """
    __slots__ = 'parts', 'total_size', 'keys', 'end_marks', 'send_future'

    def __init__(self):
        self.reset()
        self.send_future = None

    def reset(self):
        self.parts = []
        self.total_size = 0
        self.keys = []
        self.end_marks = []

class SenderActor(WorkerActor):
    """
    Actor handling sending data to ReceiverActors in other workers
    """
    def __init__(self):
        super().__init__()
        self._dispatch_ref = None
        self._events_ref = None

    def post_create(self):
        from .dispatcher import DispatchActor
        from .events import EventsActor

        super().post_create()

        self._events_ref = self.ctx.actor_ref(EventsActor.default_uid())
        if not self.ctx.has_actor(self._events_ref):
            self._events_ref = None

        self._dispatch_ref = self.promise_ref(DispatchActor.default_uid())
        self._dispatch_ref.register_free_slot(self.uid, 'sender')

    @promise.reject_on_exception
    @log_unhandled
    def send_data(self, session_id, chunk_keys, target_endpoints, ensure_cached=True,
                  compression=None, block_size=None, pin_token=None, timeout=None, callback=None):
        """
        Send data to other workers
        :param session_id: session id
        :param chunk_keys: chunks to be sent
        :param target_endpoints: endpoints to receive this chunk
        :param ensure_cached: if True, make sure the data is in the shared storage of the target worker
        :param compression: compression type when transfer in network
        :param block_size: size of data block
        :param pin_token: token to pin the data
        :param timeout: timeout of data sending
        :param callback: promise callback
        """
        chunk_keys = list(chunk_keys)
        target_endpoints = list(target_endpoints)
        block_size = block_size or options.worker.transfer_block_size

        data_sizes = self.storage_client.get_data_sizes(session_id, chunk_keys)
        if any(s is None for s in data_sizes):
            raise DependencyMissing('Dependencies %r not met when sending.'
                                    % [k for k, s in zip(chunk_keys, data_sizes) if s is None])
        if options.vineyard.enabled:
            compression = None
        else:
            compression = compression or dataserializer.CompressType(options.worker.transfer_compression)

        wait_refs = []
        addrs_to_chunks = dict()
        keys_to_readers = dict()
        receiver_refs = []
        receiver_manager_ref_dict = dict(
            (ep, self.promise_ref(ReceiverManagerActor.default_uid(), address=ep))
            for ep in target_endpoints)

        @log_unhandled
        def _create_local_readers():
            promises = []
            for k in chunk_keys:
                promises.append(self.storage_client.create_reader(
                    session_id, k, source_devices, packed=True, packed_compression=compression)
                    .then(functools.partial(operator.setitem, keys_to_readers, k)))
            return promise.all_(promises)

        @log_unhandled
        def _create_remote_writers():
            nonlocal data_sizes
            create_write_promises = []
            data_sizes = self.storage_client.get_data_sizes(session_id, chunk_keys)
            for ref in receiver_manager_ref_dict.values():
                # register transfer actions
                create_write_promises.append(
                    ref.create_data_writers(
                        session_id, chunk_keys, data_sizes, self.ref(), ensure_cached=ensure_cached,
                        timeout=timeout, pin_token=pin_token, _timeout=timeout, _promise=True
                    ).then(_handle_created)
                )
            return promise.all_(create_write_promises).then(lambda *_: keys_to_readers)

        @log_unhandled
        def _handle_created(ref, statuses):
            manager_ref = receiver_manager_ref_dict[ref.address]
            # filter out endpoints already transferred or already started transfer
            keys_receiving, keys_to_receive = [], []
            for k, status in zip(chunk_keys, statuses):
                if status == ReceiveStatus.RECEIVING:
                    keys_receiving.append(k)
                elif status is None:
                    keys_to_receive.append(k)

            if keys_to_receive:
                addrs_to_chunks[ref.address] = keys_to_receive
                receiver_refs.append(self.promise_ref(ref))
            if keys_receiving:
                wait_refs.append(manager_ref.add_keys_callback(
                    session_id, keys_receiving, _timeout=timeout, _promise=True))

        @log_unhandled
        def _finalize(*_):
            self._dispatch_ref.register_free_slot(self.uid, 'sender', _tell=True)
            self.tell_promise(callback, data_sizes)

        @log_unhandled
        def _handle_rejection(*exc):
            logger.exception('Transfer chunks %r to %r failed', chunk_keys, target_endpoints, exc_info=exc)
            self._dispatch_ref.register_free_slot(self.uid, 'sender', _tell=True)

            for reader in keys_to_readers.values():
                reader.close()
            keys_to_readers.clear()

            for ref in receiver_refs:
                ref.cancel_receive(session_id, addrs_to_chunks[ref.address], _tell=True, _wait=False)
            self.tell_promise(callback, *exc, _accept=False)

        try:
            if options.vineyard.enabled:
                in_memory_device = DataStorageDevice.VINEYARD
            else:
                in_memory_device = DataStorageDevice.SHARED_MEMORY
            source_devices = [in_memory_device, DataStorageDevice.DISK]  # pragma: no cover
            _create_local_readers().then(_create_remote_writers) \
                .then(lambda *_: self._compress_and_send(
                    session_id, addrs_to_chunks, receiver_refs, keys_to_readers,
                    block_size=block_size, timeout=timeout,
                )) \
                .then(lambda *_: promise.all_(wait_refs)) \
                .then(_finalize, _handle_rejection)
        except:  # noqa: E722
            _handle_rejection(*sys.exc_info())
            return

    @log_unhandled
    def _compress_and_send(self, session_id, addrs_to_chunks, receiver_refs, keys_to_readers,
                           block_size, timeout=None):
        """
        Compress and send data to receivers in chunked manner
        :param session_id: session id
        :param addrs_to_chunks: dict mapping endpoints to chunks to send
        :param receiver_refs: refs to send data to
        """
        # collect data targets
        chunks_to_addrs = defaultdict(set)
        for addr, chunks in addrs_to_chunks.items():
            for key in chunks:
                chunks_to_addrs[key].add(addr)
        all_chunk_keys = sorted(chunks_to_addrs.keys(), key=lambda k: len(chunks_to_addrs[k]))
        addr_statuses = dict((k, EndpointTransferState()) for k in addrs_to_chunks.keys())
        addr_to_refs = dict((ref.address, ref) for ref in receiver_refs)

        # start compress and send data into targets
        logger.debug('Data writer for chunks %r allocated at targets, start transmission', all_chunk_keys)

        # filter out endpoints we need to send to
        try:
            if not receiver_refs:
                self._dispatch_ref.register_free_slot(self.uid, 'sender', _tell=True, _wait=False)
                return
            cur_key_id = 0
            cur_key = all_chunk_keys[cur_key_id]
            cur_reader = keys_to_readers[cur_key]
            with EventContext(self._events_ref, EventCategory.PROCEDURE, EventLevel.NORMAL,
                              ProcedureEventType.NETWORK, self.uid):
                while cur_key_id < len(all_chunk_keys):
                    # read a data part from reader we defined above
                    pool = cur_reader.get_io_pool()
                    next_part = pool.submit(cur_reader.read, block_size).result()
                    file_eof = len(next_part) < block_size

                    for addr in chunks_to_addrs[cur_key]:
                        addr_status = addr_statuses[addr]
                        addr_status.parts.append(next_part)
                        addr_status.keys.append(cur_key)
                        addr_status.total_size += len(next_part)
                        addr_status.end_marks.append(file_eof)

                        if addr_status.total_size >= block_size:
                            if addr_status.send_future:
                                addr_status.send_future.result(timeout=timeout)
                            addr_status.send_future = addr_to_refs[addr].receive_data_part(
                                session_id, addr_status.keys, addr_status.end_marks,
                                *addr_status.parts, _wait=False)
                            addr_status.reset()

                    # when some part goes to end, move to the next chunk
                    if file_eof:
                        cur_reader.close()
                        cur_key_id += 1
                        if cur_key_id < len(all_chunk_keys):
                            # still some chunks left unhandled
                            cur_key = all_chunk_keys[cur_key_id]
                            cur_reader = keys_to_readers[cur_key]
                        else:
                            # all chunks handled
                            for addr, addr_status in addr_statuses.items():
                                if addr_status.send_future:
                                    addr_status.send_future.result(timeout=timeout)
                                if addr_status.parts:
                                    # send remaining chunks
                                    addr_status.end_marks[-1] = True
                                    addr_status.send_future = addr_to_refs[addr].receive_data_part(
                                        session_id, addr_status.keys, addr_status.end_marks,
                                        *addr_status.parts, _wait=False)
                                addr_status.reset()
                            for addr_status in addr_statuses.values():
                                if addr_status.send_future:
                                    addr_status.send_future.result(timeout=timeout)
                                addr_status.reset()

        except:  # noqa: E722
            for ref in receiver_refs:
                ref.cancel_receive(session_id, addrs_to_chunks[ref.address], _tell=True, _wait=False)
            raise
        finally:
            for reader in keys_to_readers.values():
                reader.close()

class ReceiverDataMeta(object):
    __slots__ = 'start_time', 'chunk_size', 'source_address',\
                'status', 'transfer_event_id', 'receiver_worker_uid', \
                'callback_ids', 'callback_args', 'callback_kwargs'

    def __init__(self, start_time=None, chunk_size=None, source_address=None,
                 transfer_event_id=None, receiver_worker_uid=None, status=None,
                 callback_ids=None, callback_args=None, callback_kwargs=None):
        self.start_time = start_time
        self.chunk_size = chunk_size
        self.source_address = source_address
        self.status = status
        self.transfer_event_id = transfer_event_id
        self.receiver_worker_uid = receiver_worker_uid
        self.callback_ids = callback_ids or []
        self.callback_args = callback_args or ()
        self.callback_kwargs = callback_kwargs or {}

    def update(self, **kwargs):
        kwargs['callback_ids'] = list(set(self.callback_ids) | set(kwargs.get('callback_ids') or ()))
        for k, v in kwargs.items():
            setattr(self, k, v)

class ReceiverManagerActor(WorkerActor):
    def __init__(self):
        super().__init__()
        self._data_metas = dict()
        self._max_callback_id = 0
        self._callback_id_to_callbacks = dict()
        self._callback_id_to_keys = dict()

        self._dispatch_ref = None

    def post_create(self):
        super().post_create()

        from .dispatcher import DispatchActor
        self._dispatch_ref = self.promise_ref(DispatchActor.default_uid())

    def _update_data_meta(self, session_id, data_key, **kwargs):
        try:
            self._data_metas[(session_id, data_key)].update(**kwargs)
        except KeyError:
            self._data_metas[(session_id, data_key)] = ReceiverDataMeta(**kwargs)

    @promise.reject_on_exception
    @log_unhandled
    def create_data_writers(self, session_id, data_keys, data_sizes, sender_ref,
                            ensure_cached=True, pin_token=None, timeout=0,
                            use_promise=True, callback=None):
        sender_address = None if sender_ref is None else sender_ref.address
        logger.debug('Begin creating transmission data writer for chunks %r from %s',
                     data_keys, sender_address)
        data_locations = dict(zip(
            data_keys, self.storage_client.get_data_locations(session_id, data_keys)))
        keys_to_fetch = []
        sizes_to_fetch = []
        statuses = []

        slot_ref = self.promise_ref(self._dispatch_ref.get_hash_slot('receiver', repr(data_keys)))
        for chunk_key, data_size in zip(data_keys, data_sizes):
            session_chunk_key = (session_id, chunk_key)

            try:
                data_meta = self._data_metas[session_chunk_key]
            except KeyError:
                data_meta = self._data_metas[session_chunk_key] = \
                    ReceiverDataMeta(chunk_size=data_size, source_address=sender_address,
                                     status=ReceiveStatus.NOT_STARTED)

            if data_locations.get(chunk_key):
                data_meta.status = ReceiveStatus.RECEIVED
                statuses.append(ReceiveStatus.RECEIVED)
                self._update_data_meta(session_id, chunk_key, status=ReceiveStatus.RECEIVED)
                continue
            elif data_meta.status == ReceiveStatus.RECEIVING:
                # data transfer already started
                logger.debug('Chunk %s already started transmission', chunk_key)
                statuses.append(ReceiveStatus.RECEIVING)
                continue
            elif data_meta.status == ReceiveStatus.RECEIVED:
                data_meta = self._data_metas[session_chunk_key] = \
                    ReceiverDataMeta(chunk_size=data_size, source_address=sender_address,
                                     status=ReceiveStatus.NOT_STARTED)

            data_meta.start_time = time.time()
            data_meta.receiver_worker_uid = slot_ref.uid
            data_meta.source_address = sender_address
            data_meta.status = ReceiveStatus.RECEIVING

            self._update_data_meta(session_id, chunk_key, chunk_size=data_size,
                                   source_address=sender_address, status=ReceiveStatus.RECEIVING)
            keys_to_fetch.append(chunk_key)
            sizes_to_fetch.append(data_size)
            statuses.append(None)  # this notifies the sender to transmit data

        if use_promise:
            if keys_to_fetch:
                slot_ref.create_data_writers(
                    session_id, keys_to_fetch, sizes_to_fetch, sender_ref, ensure_cached=ensure_cached,
                    timeout=timeout, pin_token=pin_token, use_promise=use_promise, _promise=True) \
                    .then(lambda *_: self.tell_promise(callback, slot_ref, statuses))
            else:
                self.tell_promise(callback, slot_ref, statuses)
        else:
            slot_ref.create_data_writers(
                session_id, keys_to_fetch, sizes_to_fetch, sender_ref, ensure_cached=ensure_cached,
                timeout=timeout, pin_token=pin_token, use_promise=use_promise)
            return slot_ref, statuses

    def register_pending_keys(self, session_id, data_keys):
        for key in data_keys:
            session_data_key = (session_id, key)
            if session_data_key not in self._data_metas \
                    or self._data_metas[session_data_key].status == ReceiveStatus.ERROR:
                self._update_data_meta(session_id, key, status=ReceiveStatus.PENDING,
                                       callback_args=(), callback_kwargs={})

    def filter_receiving_keys(self, session_id, data_keys):
        keys = []
        receiving_status = (ReceiveStatus.PENDING, ReceiveStatus.RECEIVING)
        for k in data_keys:
            try:
                if self._data_metas[(session_id, k)].status in receiving_status:
                    keys.append(k)
            except KeyError:
                pass
        return keys

    def add_keys_callback(self, session_id, data_keys, callback):
        cb_id = self._max_callback_id
        self._max_callback_id += 1

        receiving_status = (ReceiveStatus.PENDING, ReceiveStatus.RECEIVING)
        registered_session_keys = []
        args, kwargs = (), {}
        for k in data_keys:
            session_data_key = (session_id, k)
            data_meta = self._data_metas[session_data_key]  # type: ReceiverDataMeta
            if data_meta.status in receiving_status:
                registered_session_keys.append(session_data_key)
                data_meta.callback_ids.append(cb_id)
            else:
                args, kwargs = data_meta.callback_args, data_meta.callback_kwargs

        if registered_session_keys:
            self._callback_id_to_callbacks[cb_id] = callback
            self._callback_id_to_keys[cb_id] = set(registered_session_keys)
            logger.debug('Callback for transferring %r registered', registered_session_keys)
        else:
            self._max_callback_id = cb_id
            self.tell_promise(callback, *args, **kwargs)

    def notify_keys_finish(self, session_id, data_keys, *args, **kwargs):
        keys_to_clear = []
        for data_key in data_keys:
            session_data_key = (session_id, data_key)
            try:
                data_meta = self._data_metas[session_data_key]  # type: ReceiverDataMeta
            except KeyError:
                logger.debug('Record of %s not found.', data_key)
                continue

            try:
                data_meta.callback_args = args
                data_meta.callback_kwargs = kwargs
                if kwargs.get('_accept', True):
                    data_meta.status = ReceiveStatus.RECEIVED
                else:
                    data_meta.status = ReceiveStatus.ERROR

                cb_ids = data_meta.callback_ids
                data_meta.callback_ids = []
                if not cb_ids:
                    continue

                kwargs['_wait'] = False
                notified = 0
                for cb_id in cb_ids:
                    cb_keys = self._callback_id_to_keys[cb_id]
                    cb_keys.remove(session_data_key)
                    if not cb_keys:
                        del self._callback_id_to_keys[cb_id]
                        cb = self._callback_id_to_callbacks.pop(cb_id)
                        notified += 1
                        self.tell_promise(cb, *args, **kwargs)
                logger.debug('%d transfer listeners of %s notified.', notified, data_key)
            finally:
                if data_meta.status == ReceiveStatus.RECEIVED:
                    keys_to_clear.append(data_key)
        self.ref().clear_keys(session_id, keys_to_clear, _tell=True, _delay=5)

    def clear_keys(self, session_id, keys):
        for k in keys:
            self._data_metas.pop((session_id, k), None)

class ReceiverWorkerActor(WorkerActor):
    """
    Actor handling receiving data from a SenderActor
    """
    def __init__(self):
        super().__init__()
        self._chunk_holder_ref = None
        self._dispatch_ref = None
        self._receiver_manager_ref = None
        self._events_ref = None
        self._status_ref = None

        self._data_writers = dict()
        self._writing_futures = dict()
        self._data_metas = dict()

    def post_create(self):
        from .events import EventsActor
        from .status import StatusActor
        from .dispatcher import DispatchActor

        super().post_create()

        self._events_ref = self.ctx.actor_ref(EventsActor.default_uid())
        if not self.ctx.has_actor(self._events_ref):
            self._events_ref = None

        self._status_ref = self.ctx.actor_ref(StatusActor.default_uid())
        if not self.ctx.has_actor(self._status_ref):
            self._status_ref = None

        self._receiver_manager_ref = self.ctx.actor_ref(ReceiverManagerActor.default_uid())
        if not self.ctx.has_actor(self._receiver_manager_ref):
            self._receiver_manager_ref = None

        self._dispatch_ref = self.promise_ref(DispatchActor.default_uid())
        self._dispatch_ref.register_free_slot(self.uid, 'receiver')

    @log_unhandled
    def check_status(self, session_id, chunk_key):
        """
        Check if data exists or is being transferred in the target worker
        :param session_id: session id
        :param chunk_key: chunk key
        """
        session_chunk_key = (session_id, chunk_key)

        if self.storage_client.get_data_locations(session_id, [chunk_key])[0]:
            return ReceiveStatus.RECEIVED
        if session_chunk_key in self._data_writers:
            # data still being transferred
            return ReceiveStatus.RECEIVING
        return ReceiveStatus.NOT_STARTED

    @promise.reject_on_exception
    @log_unhandled
    def create_data_writers(self, session_id, chunk_keys, data_sizes, sender_ref,
                            ensure_cached=True, timeout=0, pin_token=None,
                            use_promise=True, callback=None):
        """
        Create a data writer for subsequent data transfer. The writer can either work on
        shared storage or spill.

        :param session_id: session id
        :param chunk_keys: chunk keys
        :param data_sizes: uncompressed data sizes
        :param sender_ref: ActorRef of SenderActor
        :param ensure_cached: if True, the data should be stored in shared memory,
                              otherwise spill is acceptable
        :param timeout: timeout if the chunk receiver does not close
        :param pin_token: token to pin the data
        :param use_promise: if True, we use promise callback to notify accomplishment
                            of writer creation, otherwise the function returns directly
                            and when sill is needed, a StorageFull will be raised instead.
        :param callback: promise callback
        """
        promises = []
        failed = False
        if options.vineyard.enabled:
            device_order = [DataStorageDevice.VINEYARD]  # pragma: no cover
        else:
            device_order = [DataStorageDevice.SHARED_MEMORY]
        source_address = sender_ref.address if sender_ref is not None else None
        if not ensure_cached:
            device_order += [DataStorageDevice.DISK]

        def _handle_accept_key(key, writer):
            if failed:
                writer.close(finished=False)
            else:
                self._data_writers[(session_id, key)] = writer

        @log_unhandled
        def _handle_reject_key(key, *exc):
            nonlocal failed
            if self.check_status(session_id, key) == ReceiveStatus.RECEIVED:
                logger.debug('Chunk %s already received', key)
            else:
                logger.debug('Rejecting %s from putting into plasma.', key)
                failed = True
                self._stop_transfer_with_exc(session_id, chunk_keys, exc)
                if callback is not None:
                    self.tell_promise(callback, *exc, _accept=False)

        # configure timeout callback
        if timeout:
            self.ref().handle_receive_timeout(session_id, chunk_keys, _delay=timeout, _tell=True)

        for chunk_key, data_size in zip(chunk_keys, data_sizes):
            self._data_metas[(session_id, chunk_key)] = ReceiverDataMeta(
                start_time=time.time(), chunk_size=data_size, source_address=source_address)

            use_device_order = device_order \
                if not isinstance(chunk_key, tuple) or not options.worker.write_shuffle_to_disk \
                else [DataStorageDevice.DISK]
            if use_promise:
                promises.append(self.storage_client.create_writer(
                    session_id, chunk_key, data_size, use_device_order, packed=True,
                    pin_token=pin_token, _promise=True)
                    .then(functools.partial(_handle_accept_key, chunk_key),
                          functools.partial(_handle_reject_key, chunk_key)))
            else:
                try:
                    _writer = self.storage_client.create_writer(
                        session_id, chunk_key, data_size, use_device_order, packed=True,
                        pin_token=pin_token, _promise=False)
                    _handle_accept_key(chunk_key, _writer)
                    return self.address, None
                except:  # noqa: E722
                    _handle_reject_key(chunk_key, *sys.exc_info())
                    raise

        promise.all_(promises).then(lambda *_: self.tell_promise(callback))

    def _wait_unfinished_writing(self, session_id, data_key):
        try:
            self._writing_futures[(session_id, data_key)].result()
            del self._writing_futures[(session_id, data_key)]
        except KeyError:
            pass

    @log_unhandled
    def receive_data_part(self, session_id, chunk_keys, end_marks, *data_parts):
        """
        Receive data part from sender
        :param session_id: session id
        :param chunk_keys: chunk keys
        :param end_marks: array with same number of boolean elements as chunk keys.
                          if one element is True, the corresponding data in data_parts
                          is the last part of the chunk.
        :param data_parts: data parts to be written
        """
        try:
            finished_keys, finished_meta_keys, finished_metas = [], [], []
            for chunk_key, data_part, end_mark in zip(chunk_keys, data_parts, end_marks):
                self._wait_unfinished_writing(session_id, chunk_key)
                session_chunk_key = (session_id, chunk_key)
                try:
                    data_meta = self._data_metas[session_chunk_key]  # type: ReceiverDataMeta

                    # if error occurred, interrupts
                    if data_meta.status == ReceiveStatus.ERROR:
                        raise data_meta.callback_args[1].with_traceback(data_meta.callback_args[2])
                    writer = self._data_writers[session_chunk_key]
                    pool = writer.get_io_pool()
                    self._writing_futures[session_chunk_key] = pool.submit(
                        writer.write, data_part)

                    if end_mark:
                        finished_keys.append(chunk_key)
                        if not isinstance(chunk_key, tuple):
                            finished_meta_keys.append(chunk_key)
                            finished_metas.append(WorkerMeta(chunk_size=data_meta.chunk_size,
                                                             workers=(self.address,)))
                except:  # noqa: E722
                    self._stop_transfer_with_exc(session_id, chunk_keys, sys.exc_info())
                    raise

            if finished_keys:
                for chunk_key in finished_keys:
                    session_chunk_key = (session_id, chunk_key)
                    data_meta = self._data_metas[session_chunk_key]  # type: ReceiverDataMeta

                    self._wait_unfinished_writing(session_id, chunk_key)
                    # update transfer speed stats
                    if self._status_ref:
                        time_delta = time.time() - data_meta.start_time
                        self._status_ref.update_mean_stats(
                            'net_transfer_speed', data_meta.chunk_size * 1.0 / time_delta,
                            _tell=True, _wait=False)
                    self._data_writers[session_chunk_key].close()
                    data_meta.status = ReceiveStatus.RECEIVED
                    logger.debug('Transfer for data %s finished.', chunk_key)
                    del self._data_writers[session_chunk_key]

                self._invoke_finish_callbacks(session_id, finished_keys)
            if finished_meta_keys:
                self.get_meta_client().batch_set_chunk_meta(session_id, finished_meta_keys, finished_metas)
        finally:
            del data_parts

    def _is_receive_running(self, session_id, chunk_key):
        receive_done_statuses = (ReceiveStatus.ERROR, ReceiveStatus.RECEIVED)
        try:
            return self._data_metas[(session_id, chunk_key)].status not in receive_done_statuses
        except KeyError:
            return False

    @log_unhandled
    def cancel_receive(self, session_id, chunk_keys, exc_info=None):
        """
        Cancel data receive by returning an ExecutionInterrupted
        :param session_id: session id
        :param chunk_keys: chunk keys
        :param exc_info: exception to raise
        """
        receiving_keys = []
        for k in chunk_keys:
            if self._is_receive_running(session_id, k):
                receiving_keys.append(k)
            self._wait_unfinished_writing(session_id, k)

        logger.debug('Transfer for %r cancelled.', chunk_keys)

        if exc_info is None:
            exc_info = build_exc_info(ExecutionInterrupted)

        self._stop_transfer_with_exc(session_id, receiving_keys, exc_info)

    @log_unhandled
    def notify_dead_senders(self, dead_workers):
        """
        When some peer workers are dead, corresponding receivers will be cancelled
        :param dead_workers: endpoints of dead workers
        """
        dead_workers = set(dead_workers)
        exc_info = build_exc_info(WorkerDead)
        session_to_keys = defaultdict(set)
        for session_chunk_key in self._data_writers.keys():
            if self._data_metas[session_chunk_key].source_address in dead_workers:
                session_to_keys[session_chunk_key[0]].add(session_chunk_key[1])
        for session_id, data_keys in session_to_keys.items():
            self.ref().cancel_receive(session_id, list(data_keys), exc_info=exc_info, _tell=True)

    @log_unhandled
    def handle_receive_timeout(self, session_id, chunk_keys):
        if not any(self._is_receive_running(session_id, k) for k in chunk_keys):
            # if transfer already finishes, no needs to report timeout
            return
        logger.debug('Transfer for %r timed out, cancelling.', chunk_keys)
        self._stop_transfer_with_exc(session_id, chunk_keys, build_exc_info(TimeoutError))

    def _stop_transfer_with_exc(self, session_id, chunk_keys, exc):
        for chunk_key in chunk_keys:
            self._wait_unfinished_writing(session_id, chunk_key)

        if not isinstance(exc[1], ExecutionInterrupted):
            logger.exception('Error occurred in receiving %r. Cancelling transfer.',
                             chunk_keys, exc_info=exc)

        for chunk_key in chunk_keys:
            session_chunk_key = (session_id, chunk_key)

            # stop and close data writer
            try:
                # transfer is not finished yet, we need to clean up unfinished stuffs
                self._data_writers[session_chunk_key].close(finished=False)
                del self._data_writers[session_chunk_key]
            except KeyError:
                # transfer finished and writer cleaned, no need to clean up
                pass

            try:
                data_meta = self._data_metas[session_chunk_key]  # type: ReceiverDataMeta
                data_meta.status = ReceiveStatus.ERROR
            except KeyError:
                pass

        self._invoke_finish_callbacks(session_id, chunk_keys, *exc, **dict(_accept=False))

    def _invoke_finish_callbacks(self, session_id, chunk_keys, *args, **kwargs):
        # invoke registered callbacks for chunk
        for k in chunk_keys:
            try:
                data_meta = self._data_metas.pop((session_id, k))  # type: ReceiverDataMeta
            except KeyError:
                continue

            if data_meta.transfer_event_id is not None and self._events_ref is not None:
                self._events_ref.close_event(data_meta.transfer_event_id, _tell=True, _wait=False)
            if not kwargs.get('_accept', True):
                if not data_meta.callback_args or data_meta.callback_args[0] is ExecutionInterrupted:
                    data_meta.callback_args = args
                    data_meta.callback_kwargs = kwargs
                else:
                    args = data_meta.callback_args
                    kwargs = data_meta.callback_kwargs
            else:
                data_meta.callback_args = args
                data_meta.callback_kwargs = kwargs
        kwargs['_tell'] = True
        if self._receiver_manager_ref:
            self._receiver_manager_ref.notify_keys_finish(session_id, chunk_keys, *args, **kwargs)

class ResultCopyActor(WorkerActor):
    def start_copy(self, session_id, chunk_key, targets):
        locations = [v[1] for v in self.storage_client.get_data_locations(session_id, [chunk_key])[0]]
        if set(locations).intersection(targets):
            return
        ev = self.ctx.event()
        self.storage_client.copy_to(session_id, [chunk_key], targets) \
            .then(lambda *_: ev.set())
        return ev

class ResultSenderActor(WorkerActor):
    """
    Actor handling sending result to user client
    """
    def __init__(self):
        super().__init__()
        self._result_copy_ref = None
        self._serialize_pool = None

    def post_create(self):
        super().post_create()
        self._serialize_pool = self.ctx.threadpool(1)
        self._result_copy_ref = self.ctx.create_actor(ResultCopyActor, uid=ResultCopyActor.default_uid())

    def pre_destroy(self):
        self._result_copy_ref.destroy()
        super().pre_destroy()

    def fetch_batch_data(self, session_id, chunk_keys, index_objs=None, compression_type=None):
        results = []
        if index_objs is not None:
            for chunk_key, index_obj in zip(chunk_keys, index_objs):
                results.append(self.fetch_data(session_id, chunk_key, index_obj, compression_type=compression_type))
        else:
            for chunk_key in chunk_keys:
                results.append(self.fetch_data(session_id, chunk_key, compression_type=compression_type))
        return results

    def fetch_data(self, session_id, chunk_key, index_obj=None, compression_type=None):
        logger.debug('Sending data %s from %s, using compression: %r', chunk_key, self.uid, compression_type)
        if options.vineyard.enabled:
            compression_type = None
        else:
            compression_type = compression_type or dataserializer.CompressType(options.worker.transfer_compression)

        if options.vineyard.enabled:
            in_memory_device = DataStorageDevice.VINEYARD
        else:
            in_memory_device = DataStorageDevice.SHARED_MEMORY

        if index_obj is None:
            target_devs = [in_memory_device, DataStorageDevice.DISK]  # pragma: no cover
            ev = self._result_copy_ref.start_copy(session_id, chunk_key, target_devs)
            if ev:
                ev.wait(options.worker.prepare_data_timeout)

            reader = self.storage_client.create_reader(
                session_id, chunk_key, target_devs, packed=True,
                packed_compression=compression_type, _promise=False)

            with reader:
                pool = reader.get_io_pool()
                return pool.submit(reader.read).result()
        else:
            try:
                value = self.storage_client.get_object(
                    session_id, chunk_key, [in_memory_device], _promise=False)
            except IOError:
                reader = self.storage_client.create_reader(
                    session_id, chunk_key, [DataStorageDevice.DISK], packed=False, _promise=False)
                with reader:
                    pool = reader.get_io_pool()
                    value = dataserializer.deserialize(pool.submit(reader.read).result())

            try:
                sliced_value = value.iloc[tuple(index_obj)]
            except AttributeError:
                sliced_value = value[tuple(index_obj)]

            return self._serialize_pool.submit(
                dataserializer.dumps, sliced_value, compress=compression_type).result()

def put_remote_chunk(session_id, chunk_key, data, receiver_manager_ref):
    """
    Put a chunk to target machine using given receiver_ref
    """
    from .dataio import ArrowBufferIO
    buf = dataserializer.serialize(data).to_buffer()
    receiver_ref, _ = receiver_manager_ref.create_data_writers(
        session_id, [chunk_key], [buf.size], None, ensure_cached=False, use_promise=False)
    receiver_ref = receiver_manager_ref.ctx.actor_ref(receiver_ref)
    block_size = options.worker.transfer_block_size

    reader = None
    try:
        reader = ArrowBufferIO(buf, 'r', block_size=block_size)
        futures = []
        while True:
            next_part = reader.read(block_size)
            is_last = not next_part or len(next_part) < block_size
            [f.result() for f in futures]
            futures.append(receiver_ref.receive_data_part(
                session_id, [chunk_key], [is_last], next_part, _wait=False))
            if is_last:
                [f.result() for f in futures]
                break
    except:  # noqa: E722
        receiver_ref.cancel_receive(session_id, [chunk_key])
        raise
    finally:
        if reader:
            reader.close()
        del reader


from fenics import *

mesh = Mesh('mesh/res32.xdmf')	# mm

plot(mesh,interactive=True)

# Since the mesh is in mm pressure units in pascal must be scaled by alpha = (1e6)**(-1)
alpha = (1e6)**(-1)

# Mark boundaries
class Neumann_boundary(SubDomain):
	def inside(self, x, on_boundry):
		return on_boundry

mf = FacetFunction("size_t", mesh)
mf.set_all(0)

Neumann_boundary().mark(mf, 1)
ds = ds[mf]

# Continuum mechanics
E = 16*1e3 *alpha
nu = 0.25
mu, lambda_ = Constant(E/(2*(1 + nu))), Constant(E*nu/((1 + nu)*(1 - 2*nu)))
epsilon = lambda u: sym(grad(u))

p_outside = 133 *alpha
n = FacetNormal(mesh)
f = Constant((0, 0, 0))

V = VectorFunctionSpace(mesh, "Lagrange", 1)

# --------------- Handle Neumann-problem --------------- #
R = FunctionSpace(mesh, 'R', 0)        		 # space for one Lagrange multiplier
M = MixedFunctionSpace([R]*6)          		 # space for all multipliers
W = MixedFunctionSpace([V, M])
u, mus = TrialFunctions(W)
v, nus = TestFunctions(W)

# Establish a basis for the nullspace of RM
e0 = Constant((1, 0, 0))				# translations
e1 = Constant((0, 1, 0))
e2 = Constant((0, 0, 1))

e3 = Expression(('-x[1]', 'x[0]', '0')) # rotations
e4 = Expression(('-x[2]', '0', 'x[0]'))
e5 = Expression(('0', '-x[2]', 'x[1]'))
basis_vectors = [e0, e1, e2, e3, e4, e5]

a = 2*mu*inner(epsilon(u),epsilon(v))*dx + lambda_*inner(div(u),div(v))*dx
L = inner(f, v)*dx + p_outside*inner(n,v)*ds(1)

# Lagrange multipliers contrib to a
for i, e in enumerate(basis_vectors):
	mu = mus[i]
	nu = nus[i]
	a += mu*inner(v, e)*dx + nu*inner(u, e)*dx

# -------------------------------------------------------- #

# Assemble the system
A = PETScMatrix()
b = PETScVector()
assemble_system(a, L, A_tensor=A, b_tensor=b)

# Solve
uh = Function(W)
solver = PETScLUSolver('mumps') # NOTE: we use direct solver for simplicity
solver.set_operator(A)
solver.solve(uh.vector(), b)

# Split displacement and multipliers. Plot
u, ls = uh.split(deepcopy=True)
plot(u, mode='displacement', title='Neumann_displacement',interactive=True)

file = File('deformed_brain.pvd')
file << u


from mpl_toolkits.basemap import Basemap
import matplotlib.pyplot as plt
import xlrd
from matplotlib.patches import Polygon
import numpy as np
import pandas
import xlrd
import xlwt
import numpy
import pandas as pd
import math
import traceback
from mpl_toolkits.basemap import Basemap
denglu_index = {}
total_Coast_point = 22#从0开始计算
up_index = 19
passindex = {}
def GET_DATA():
    keyword = ['name_datetime', 'I', 'LAT', 'LON', 'PRES', 'WND', 'OWD', 'END_TAG']
    ALL_DATA = {'name_datetime': [], 'I': [], 'LAT': [], 'LON': [], 'PRES': [], 'WND': [], 'OWD': [], 'END_TAG': []}
    data = xlrd.open_workbook('1981-2018总表.xls').sheet_by_index(0)
    for i in range(len(keyword)):
        if keyword[i] == 'LAT' or keyword[i] == 'LON':
            ALL_DATA[keyword[i]] += [float(item) / 10 for item in data.col_values(i)[1:]]
        else:
            ALL_DATA[keyword[i]] += data.col_values(i)[1:]
    return ALL_DATA
def if_influence_this_province_one_point(shapefile,lat,lon):
    left_up = False
    left_down = False
    right_up = False
    right_down = False
    for Boundary in shapefile:
        if Boundary[0] > lon and Boundary[1] > lat:
            right_up = True
        if Boundary[0] > lon and Boundary[1] < lat:
            right_down = True
        if Boundary[0] < lon and Boundary[1] > lat:
            left_up = True
        if Boundary[0] < lon and Boundary[1] < lat:
            left_down = True
        if  left_up and left_down and right_up and right_down:
            return True
    return False
def if_influence_this_province_path(shapefile,path,pathid,province):
    for i in range(len(path)):
        point = path[i]
        if if_influence_this_province_one_point(shapefile,point[0],point[1]):
            denglu_index[pathid] = [i,point,province]
            return True
    denglu_index[pathid] = 0
    return  False
def if_influence(SHAPE_FILE,path,pathid):
        for province in SHAPE_FILE:
            if if_influence_this_province_path(SHAPE_FILE[province],path,pathid,province):
                return True
        return False
def path_data():
    ALL_path = {}
    data = GET_DATA()
    path_id = []
    for i in range(len(data['name_datetime'])):
        ID =  data['name_datetime'][i][:9]
        if ID not in ALL_path:
            ALL_path[ID] = []
            path_id .append(ID)
        TIM = data['name_datetime'][i][5:]
        lat = data['LAT'][i]
        lon = data['LON'][i]
        PRE = data['PRES'][i]
        WND = data['WND'][i]
        END = data['END_TAG'][i]
        LEL = data['I'][i]
        ALL_path[ID].append((lat,lon,PRE,WND,END,TIM,LEL))
    return (ALL_path,path_id)
def load_4shapfile(filepath):
    map = Basemap(projection='cyl', lat_0=90, lon_0=160, \
                  llcrnrlat=0, urcrnrlat=45, \
                  llcrnrlon=90, urcrnrlon=180, \
                  rsphere=6371200., resolution='h', area_thresh=1000)

    map.readshapefile(filepath, 'states', drawbounds=False, color='black', linewidth=0.4)
    SHAPE_FILE = {}
    for info, shp in zip(map.states_info, map.states):
        province = info['NAME_1']  # name of state
        if province not in ['Zhejiang','Shanghai','Guangdong','Fujian']:
            continue
        if province not in SHAPE_FILE:
            SHAPE_FILE[province] = [shp]
        else:
            SHAPE_FILE[province].append(shp)
    for key in SHAPE_FILE:
        b = list(sorted(SHAPE_FILE[key], key=lambda e: len(e), reverse=True))
        SHAPE_FILE[key] = b[0]
    print('读取4省shapefile成功')
    return SHAPE_FILE
def wrtie_influence():
    a  = pandas.read_csv('influence.csv')
    return (list(a['ID'].values),list(a['influence'].values))
def piecelinear_of_coastline():
    data = {}
    coastline = coast_line()
    for i in range(total_Coast_point):
        x1 = coastline[i]
        x2 = coastline[i+1]
        data [i] = {
            'coef':(x1[1] - x2[1]) / (x1[0] - x2[0]),
            'start_y':x1[1],
            'start_x':x1[0],
            'end_x'  :x2[0],
            'end_y'  :x2[1]
        }
    return data
def if_point_in_coastline(coast_line_piecelinear,point):
    '判断一个点是否在海岸线内'
    '注意高精度需求'
    for i in range(up_index + 1):
        point_x = point[1]
        point_y = point[0]
        start_x = coast_line_piecelinear[i]['start_x']
        start_y = coast_line_piecelinear[i]['start_y']
        end_x   = coast_line_piecelinear[i]['end_x']
        end_y   = coast_line_piecelinear[i]['end_y']
        if point_x >= start_x and point_x < end_x and point_y <= 32:
            if point_y*(end_x - start_x) >= (end_y - start_y)*(point_x - start_x) + start_y * (end_x - start_x):
                if point_x > 121.972356:
                    if point_x <= 122.037456:
                        if  point_y*(122.037456 - 121.972356) <= (30.633228  - 31.591419)*(point_x - 121.972356) + 31.591419*(122.037456 - 121.972356):
                            return i
                    elif point_x <= 122.098542:
                        if  point_y*(122.098542 - 122.037456) <= (30.099967  - 30.633228)*(point_x - 122.037456) + 30.633228*(122.098542 - 122.037456):
                            return i
                else:
                    return  i
def get_denglu_index():
    '返回每个台风路径第几个点已经登录 与登录前的一个点,与在海岸线的哪个片段'
    'point[0] = lat'
    'point[1] = lon'
    pathdata = path_data()[0]
    coast_line_piecelinear = piecelinear_of_coastline()
    data = {}
    ALL = {'ID':[],'denglu':[]}
    count = 0
    for pathid in pathdata:
        path = pathdata[pathid]
        denglu_index = -1
        data[pathid] = denglu_index
        ALL['ID'].append(pathid)
        find = False
        for i in range(len(path)):
            point = path[i]
            RETURN = if_point_in_coastline(coast_line_piecelinear,point)
            if RETURN != None :
                count += 1
                data[pathid] = [point,path[i-1],RETURN]
                find= True
                break
        if find:
            ALL['denglu'].append(1)
        else:
            ALL['denglu'].append(0)
    print('登录台风个数为:',count)
    pandas.DataFrame(ALL).to_csv('登录台风统计.csv')
    return data
def path_vector_knn():
    data = {
        'pathid':[],
        'start_x':[],
        'start_y':[],
        'destination_x':[],
        'destination_y': [],
        'angle':[]
    }
    pathdata = path_data()[0]
    for pathid in pathdata:
        s_x = pathdata[pathid][0][1]
        s_y = pathdata[pathid][0][0]
        e_x = pathdata[pathid][-1][1]
        e_y = pathdata[pathid][-1][0]
        data['pathid'] .append(pathid)
        data['start_x'].append(s_x)
        data['start_y'].append(s_y)
        data['destination_x'].append(e_x)
        data['destination_y'].append(e_y)
        data['angle'].append((s_x*e_x + s_y*e_y)/(math.sqrt(s_x*s_x+s_y*s_y) * math.sqrt(e_x*e_x+e_y*e_y)))
    pandas.DataFrame(data).to_csv('KNN_path_vector.csv')
def str_to_latlon(x):
    x = x.replace("\\",'').replace('\n','').replace('"','')
    du_index = x.index('°')
    fen_index = x.index("'")
    du =  int(x[:du_index])
    fen = int(x[du_index + 1: fen_index])
    sec = float(x[fen_index +1 :])
    return round(du + fen/60 + sec/3600,6)
def coast_line():
    global total_Coast_point
    '获得海岸线分割点数据'
    with open('coastline.txt') as f:
         raw = f.readlines()
    coastline = {}
    total_Coast_point = -1
    for this_item in raw:
        total_Coast_point += 1
        item = this_item.split(' ')
        coastline[total_Coast_point] = (str_to_latlon(item[2]),str_to_latlon(item[1]))
    return coastline
def line_of_two_point(x1, x2):
    coef = (x1[1] - x2[1]) / (x1[0] - x2[0])
    x = numpy.linspace(x1[0], x2[0], 1000)
    return x,numpy.array(list([coef*(x0 - x1[0]) + x1[1] for x0 in x]))
def cross_point(line2, line1):  # 计算交点函数
    x1 = line1[0]  # 取四点坐标
    y1 = line1[1]
    x2 = line1[2]
    y2 = line1[3]
    x3 = line2[0]
    y3 = line2[1]
    x4 = line2[2]
    y4 = line2[3]
    k1 = (y2 - y1) * 1.0 / (x2 - x1)  # 计算k1,由于点均为整数，需要进行浮点数转化
    b1 = y1 * 1.0 - x1 * k1 * 1.0  # 整型转浮点型是关键
    if (x4 - x3) == 0:  # L2直线斜率不存在操作
        k2 = None
        b2 = 0
    else:
        k2 = (y4 - y3) * 1.0 / (x4 - x3)  # 斜率存在操作
        b2 = y3 * 1.0 - x3 * k2 * 1.0
    if k2 == None:
        x = x3
    else:
        x = (b2 - b1) * 1.0 / (k1 - k2)
    y = k1 * x * 1.0 + b1 * 1.0
    return [x, y]
def get_denglu_point():
    '得到所有台风在海岸线上的登陆点'
    denglu = get_denglu_index()
    coastline = piecelinear_of_coastline()
    ALL = []
    for pathid in denglu:
        if denglu[pathid] == -1:
            passindex[pathid] = -1
            continue
        denglu_point = denglu[pathid][0]
        last_point = denglu[pathid][1]
        find = False
        for piece_index in range(0, total_Coast_point):
            piece = coastline[piece_index]
            crosspoint = cross_point([denglu_point[1], denglu_point[0], last_point[1], last_point[0]],
                                     [piece['start_x'], piece['start_y'], piece['end_x'], piece['end_y']])
            if piece['start_x'] < piece['end_x']:
                if crosspoint[0] >= piece['start_x'] and crosspoint[0] <= piece['end_x'] and crosspoint[0] >= min(
                        denglu_point[1], last_point[1]) and crosspoint[0] <= max(denglu_point[1], last_point[1]):
                    d_l = (crosspoint[0] - last_point[1]) ** 2 + (crosspoint[1] - last_point[0]) ** 2
                    d_d = (crosspoint[0] - denglu_point[1]) ** 2 + (crosspoint[1] - denglu_point[0]) ** 2
                    l_l = int(last_point[-1])
                    l_d = int(denglu_point[-1])
                    l_c = (d_l * l_l + d_d * l_d) / (d_l + d_d)
                    # 用加权平均计算登陆点台风的等级
                    crosspoint.append(l_c)
                    ALL.append(crosspoint)
                    passindex[pathid] = piece_index
                    find = True
            else:
                if crosspoint[0] <= piece['start_x'] and crosspoint[0] >= piece['end_x'] and crosspoint[0] >= min(
                        denglu_point[1], last_point[1]) and crosspoint[0] <= max(denglu_point[1], last_point[1]):
                    d_l = (crosspoint[0] - last_point[1]) ** 2 + (crosspoint[1] - last_point[0]) ** 2
                    d_d = (crosspoint[0] - denglu_point[1]) ** 2 + (crosspoint[1] - denglu_point[0]) ** 2
                    l_l = int(last_point[-1])
                    l_d = int(denglu_point[-1])
                    l_c = (d_l * l_l + d_d * l_d) / (d_l + d_d)
                    # 用加权平均计算登陆点台风的等级
                    crosspoint.append(l_c)
                    ALL.append(crosspoint)
                    passindex[pathid] = piece_index
                    find = True
    return ALL
def get_denglu_point2():
    '得到所有台风在海岸线上的登陆点'
    denglu = get_denglu_index()
    coastline = piecelinear_of_coastline()
    ALL = {}
    for pathid in denglu:
        if denglu[pathid] == -1:
            passindex[pathid] = -1
            continue
        denglu_point = denglu[pathid][0]
        last_point = denglu[pathid][1]
        find = False
        for piece_index in range(0, total_Coast_point):
            piece = coastline[piece_index]
            crosspoint = cross_point([denglu_point[1], denglu_point[0], last_point[1], last_point[0]],
                                     [piece['start_x'], piece['start_y'], piece['end_x'], piece['end_y']])
            if piece['start_x'] < piece['end_x']:
                if crosspoint[0] >= piece['start_x'] and crosspoint[0] <= piece['end_x'] and crosspoint[0] >= min(
                        denglu_point[1], last_point[1]) and crosspoint[0] <= max(denglu_point[1], last_point[1]):
                    d_l = (crosspoint[0] - last_point[1]) ** 2 + (crosspoint[1] - last_point[0]) ** 2
                    d_d = (crosspoint[0] - denglu_point[1]) ** 2 + (crosspoint[1] - denglu_point[0]) ** 2
                    l_l = int(last_point[-1])
                    l_d = int(denglu_point[-1])
                    l_c = (d_l * l_l + d_d * l_d) / (d_l + d_d)
                    # 用加权平均计算登陆点台风的等级
                    crosspoint.append(l_c)
                    passindex[pathid] = piece_index
                    find = True
                    ALL[pathid] = crosspoint
            else:
                if crosspoint[0] <= piece['start_x'] and crosspoint[0] >= piece['end_x'] and crosspoint[0] >= min(
                        denglu_point[1], last_point[1]) and crosspoint[0] <= max(denglu_point[1], last_point[1]):
                    d_l = (crosspoint[0] - last_point[1]) ** 2 + (crosspoint[1] - last_point[0]) ** 2
                    d_d = (crosspoint[0] - denglu_point[1]) ** 2 + (crosspoint[1] - denglu_point[0]) ** 2
                    l_l = int(last_point[-1])
                    l_d = int(denglu_point[-1])
                    l_c = (d_l * l_l + d_d * l_d) / (d_l + d_d)
                    # 用加权平均计算登陆点台风的等级
                    crosspoint.append(l_c)
                    ALL[pathid] = crosspoint
                    passindex[pathid] = piece_index
                    find = True
        if not find:
            ALL[pathid] = -2
            print(pathid)
    return ALL
def get_denglu_pathid():
    denglus = get_denglu_index()
    all = []
    for id in denglus:
        if denglus[id] != -1:
            all.append(id)
    return all
def get_type_of_denlgu():
    return  pandas.read_csv('sum_denglu_cluster.csv',encoding='gbk')['类别'].values
def get_denglu_count():
    denglu_id =  get_denglu_pathid()
    pathes = path_data()[0]
    ALL = {}
    count = 0
    for id in pathes:
        if id not in denglu_id:
            continue
        count += 1
        point = pathes[id][0]
        times = point[-2][:6]
        if times not in ALL:
            ALL[times] = 1
        else:
            ALL[times] += 1
    DATA  = {'时间':[],'登录数':[]}
    for year in range(1981,2019):
        for month in range(1,13):
            if month < 10:
                this_time = '{}0{}'.format(year,month)
            else:
                this_time = '{}{}'.format(year, month)
            DATA['时间'].append(this_time)
            if this_time not in ALL:
                DATA['登录数'].append(0)
            else:
                DATA['登录数'].append(ALL[this_time])
    pandas.DataFrame(DATA).to_csv('登录次数年月统计.csv')
def get_denglu_province():
    get_denglu_point()
    ALL = {'ID':[],'province':[]}
    data = {}
    for id in passindex:
        ALL['ID'].append(id)
        if passindex[id] == -1:
            ALL['province'].append('未登录')
            data[id]='未登录'
        elif passindex[id] <= 11:
            ALL['province'].append('广东')
            data[id] = '广东'
        elif passindex[id] <= 14:
            ALL['province'].append('福建')
            data[id] = '福建'
        elif passindex[id] <= 20:
            ALL['province'].append('浙江')
            data[id] = '浙江'
        else:
            ALL['province'].append('上海')
            data[id] = '上海'
    print(data)
    pandas.DataFrame(ALL).to_csv('台风登录省份统计.csv')
    return data
def get_province_count():
    data = pandas.read_csv('denglu.csv',encoding='gbk')
    y = []
    raw_y = data['cishu'].values
    for item in raw_y:
        y.append(item)
    return y
def dataset(cut):
    seq = get_province_count()
    x = []
    y = []
    for i in range(cut,len(seq)):
        y.append(seq[i])
        x.append([item[0]for item in seq[i-cut:i]])
    return (x,y)
def detali_denglu_point():
    denglupoints = get_denglu_point2()
    data = get_denglu_index()
    count = 0
    book = xlwt.Workbook()
    sheet = book.add_sheet('0')
    for id in data:
        detail = data[id]
        if detail == -1:
            continue
        denglupoint = denglupoints[id]
        sheet.write(count, 0,id)
        #format of point:(lat,lon,PRE,WND,END,TIM,LEL)
        if denglupoint != -2:
            sheet.write(count, 1, denglupoint[0])
            sheet.write(count, 2, denglupoint[1])
        else:
            sheet.write(count, 1, '无')
            sheet.write(count, 2, '无')
        afterpoint = detail[0]
        beforepoint = detail[1]
        sheet.write(count,3,beforepoint[1])
        sheet.write(count,4,beforepoint[0])
        sheet.write(count,5,int(beforepoint[-1]))
        sheet.write(count,6,int(beforepoint[3]))
        sheet.write(count,7,int(beforepoint[2]))
        sheet.write(count,8,afterpoint[1])
        sheet.write(count,9,afterpoint[0])
        sheet.write(count,10,int(afterpoint[-1]))
        sheet.write(count,11,int(afterpoint[3]))
        sheet.write(count,12,int(afterpoint[2]))
        count += 1
    book.save('outcome.xls')
colormap_name = 'Oranges'
year = 2010
#调色板颜色
def draw_lon_lat():
  parallels = np.arange(0,45,20)
  map.drawparallels(parallels,labels=[1,0,0,0],fontsize=10) # 绘制纬线
  meridians = np.arange(80,181,20)
  map.drawmeridians(meridians,labels=[0,0,0,1],fontsize=10) # 绘制经线
def load_shapfile(filepath):
    map.readshapefile(filepath, 'states', drawbounds=False, color='black', linewidth=0.4)
    SHAPE_FILE = {}
    for info, shp in zip(map.states_info, map.states):
        province = info['NAME_1']  # name of state
        if (province) not in SHAPE_FILE:
            SHAPE_FILE[province] = [shp]
        else:
            SHAPE_FILE[province].append(shp)
    for key in SHAPE_FILE:
        b = list(sorted(SHAPE_FILE[key], key=lambda e: len(e), reverse=True))
        SHAPE_FILE[key] = b[0]
    print('读取shapefile成功')
    return SHAPE_FILE
def draw_province_name():
    label_latlon = {
        'Guangdong ' : (113.07-0.5,23.10 - 0.3),
        'Fujian' : (118.22-2,27.04-1.5),
        'Zhejiang': (120.51 - 1.5,29.09-0.8),
        'Shanghai':(121.23+0.5,31.04)
    }
    for province in label_latlon:
        print(province)
        x,y = map(label_latlon[province][0],label_latlon[province][1])
        plt.text(x, y,province, fontsize=11, fontweight='bold', ha='left', va='bottom', color='k')
def draw_my_coastline(zorder):
    coastline = coast_line()
    for i in range(total_Coast_point):
        print('描绘海岸线至第{}个点 /{}'.format(i,total_Coast_point))
        line_x,line_y = line_of_two_point(coastline[i],coastline[i+1])
        map.plot(line_x,line_y,linewidth = 1,color = 'green',latlon = True,zorder = zorder)
def fill_province():
    SHAPE_FILE = load_shapfile('shapefile\\gadm36_CHN_1')
    province_color = {
    'Guangdong' : 'yellow',
    'Fujian' :'orange',
    'Zhejiang' :'green',
    'Shanghai':'red'
    }
    for province in SHAPE_FILE:
        if province in ['Guangdong','Fujian','Zhejiang','Shanghai']:
            poly = Polygon(SHAPE_FILE[province],fill = True,facecolor=province_color[province],edgecolor='black', lw=0.6,alpha = 0.6)
            ax1.add_patch(poly)
def taifeng_average_level():
     data = xlrd.open_workbook('登陆台风整理.xlsx').sheet_by_index(0)
     levels = data.col_values(7)[1:]
     id     = data.col_values(2)[1:]
     all = {}
     for i in range(len(id)):
         if id[i] not in all:
             if i != 0:
                 print(all[id[i - 1]])
                 print(need)
                 all[id[i - 1]] = math.ceil(all[id[i - 1]]/need)
             all[id[i]]  = levels[i]
             need  = 1
         else:
             all[id[i]] += levels[i]
             need += 1
     all[id[- 1]] = math.ceil(all[id[- 1]]/need)
     return all
fig = plt.figure(figsize=(20 ,18))
ax1 = fig.add_axes([0.1,0.1,0.8,0.8])
map = Basemap(projection='cyl',lat_0=90,lon_0=120,\
            llcrnrlat=15 ,urcrnrlat=25,\
            llcrnrlon=105,urcrnrlon=115,\
            rsphere=6371200.,resolution='h',area_thresh=1000)
#map.fillcontinents(color = '#ddaa66', lake_color = 'aqua')
#map.drawmapboundary(fill_color = 'aqua')
pathes = path_data()[0]
dengluid = get_denglu_pathid()
dengluty = get_type_of_denlgu()
type_color = {
    0: 'purple',
    '1-0' : 'purple',
    '1月1日' :  'purple',
    '1月2日':  'purple',
     2: 'purple'
}
count = 0
province = get_denglu_province()
for i in range(len(dengluid)):
    print('描绘第{}个台风中/{}'.format(i+1,250))
    ID = dengluid[i]
    print(ID)
    path = pathes[ID]
    path_type = dengluty[i]
    pathcolor = 'black'
    if ID not in province or province[ID] != '广东':
        continue
    for j in range(len(path) - 1 ):
        this_point = path[j]
        next_point = path[j + 1]
        #map.scatter(x=this_point[1], y=this_point[0], marker='o', color='purple', s=5, ax=ax1, edgecolor='purple',zorder = 999999)
        #map.scatter(x=next_point[1], y=next_point[0], marker='o', color='purple', s=5, ax=ax1, edgecolor='purple',zorder = 999999)
        x1,y1 =map(this_point[1],this_point[0])
        x2,y2 =map(next_point[1],next_point[0])
        if x1 == x2 and y1 == y2:
            continue
        #print(this_point[1], this_point[0])
        #print(next_point[1], next_point[0])
        #print(x1,y1)
        ax1.arrow(x1, y1,x2 - x1,y2 - y1,
             length_includes_head=True,  # 增加的长度包含箭头部分
             head_width=0.01, head_length=0.01, fc= pathcolor, ec=pathcolor,linewidth = 0.1)
        #print()
#item = (lat,lon,PRE,WND,END,TIM)
#map.drawcountries()
#map.etopo(scale=1, alpha=0.8)
#map.bluemarble(scale = 1,alpha = 0.4)# 绘制蓝色卫星图
map.shadedrelief()
#map.etopo(scale=1, alpha=0.5)
fill_province()
draw_province_name()
#draw_my_coastline(999999 + 1)
plt.savefig('pic.png')
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False
plt.title('台风路径图')
plt.show()

130_不同省份登录台风路线.py

130_不同省份登录台风路线.py


import numpy as np
import cv2 as cv

shrub_img = cv.imread('Screenshot.jpg', cv.IMREAD_UNCHANGED)
needle_img = cv.imread('scrub.jpg', cv.IMREAD_UNCHANGED)

needle_w = needle_img.shape[1]
needle_h = needle_img.shape[0]
result = cv.matchTemplate(shrub_img, needle_img, cv.TM_CCOEFF_NORMED)

# get the best match position
min_val, max_val, min_loc, max_loc = cv.minMaxLoc(result)

result = cv.matchTemplate(shrub_img, needle_img, cv.TM_SQDIFF_NORMED)

print(result)

threshold = 0.09
locations = np.where(result <= threshold)
locations = list(zip(*locations[::-1]))
rectangles = []
for loc in locations:
    rect = [int(loc[0]), int(loc[1]), needle_w, needle_h]
    rectangles.append(rect)
print(rectangles)
if len(rectangles):
    print("Found needle")

    line_color = (0, 255, 0)
    line_type = cv.LINE_4
    thickness = 4

    # need to loop over all the locations and draw the rectangles around them
    for loc in rectangles:
        # determine the box position
        top_left = loc
        bottom_right = (top_left[0]+needle_w, top_left[1]+needle_h)

        # draw the box
        cv.rectangle(shrub_img, tuple(top_left), tuple(bottom_right),
                     line_color, thickness)
cv.imshow('Matches', shrub_img)
cv.waitKey()

'''
threshold = 0.8

if max_val >= threshold:
    print('Found needle')
# get dimensions of the needle image
    needle_w = needle_img.shape[1]
    needle_h = needle_img.shape[0]

    top_left = max_loc
    bottom_right = (top_left[0]+needle_w, top_left[1]+needle_h)

    cv.rectangle(shrub_img, top_left, bottom_right, color=(
        0, 255, 0), thickness=2, lineType=cv.LINE_4)

    cv.imshow('Result', shrub_img)
    cv.waitKey()
else:
    print("Needle not found")
'''


try:
    import Image
except ImportError:       # pragma: no cover
    from PIL import Image # pragma: no cover

class Resizer(object):

    def __init__(self, expand=False, filter=Image.ANTIALIAS):
        self.expand = expand
        self.filter = filter

    def transformations(self, img, target_size):
        # should return a list of JSON-serializable commands
        # that are processed with transform method
        raise NotImplementedError

    def transform(self, img, transformation, params):
        # transformations MUST be idempotent
        if transformation == 'crop':
            return img.crop(params)
        elif transformation == 'resize':
            return img.resize(params, self.filter)
        else:
            raise NotImplementedError(transformation)

    def __call__(self, img, target_size):
        transforms = self.transformations(img, target_size)
        for transformation, params in transforms:
            img = self.transform(img, transformation, params)
        return img

class ResizeFit(Resizer):

    def transformations(self, img, target_size):
        sw, sh = img.size
        tw, th = target_size
        if not self.expand and sw<=tw and sh<=th:
            return []

        if sw*th>sh*tw:
            h = sh*tw//sw
            return [('resize', (tw, h))]
        else:
            w = sw*th//sh
            return [('resize', (w, th))]

class ResizeCrop(Resizer):

    def transformations(self, img, target_size):
        sw, sh = img.size
        tw, th = target_size
        if not self.expand and sw<=tw and sh<=th:
            return []

        transforms = []
        if sw*th>sh*tw:
            if sh!=th and (sh>th or self.expand):
                w = sw*th//sh
                transforms.append(('resize', (w, th)))
                sw, sh = w, th
            if sw>tw:
                wd = (sw-tw)//2
                transforms.append(('crop', (wd, 0, tw+wd, sh)))
        else:
            if sw!=tw and (sw>tw or self.expand):
                h = sh*tw//sw
                transforms.append(('resize', (tw, h)))
                sw, sh = tw, h
            if sh>th:
                hd = (sh-th)//2
                transforms.append(('crop', (0, hd, sw, th+hd)))
        return transforms

class ResizeMixed(Resizer):

    def __init__(self, hor_resize, vert_resize):
        self.hor_resize = hor_resize
        self.vert_resize = vert_resize

    def get_resizer(self, img, target_size):
        sw, sh = img.size
        if sw >= sh:
            return self.hor_resize
        else:
            return self.vert_resize

    def transformations(self, img, size):
        return self.get_resizer(img, size).transformations(img, size)

    #def transform(self, img, *args):
    # XXX is this method needed?
    #    return self.get_resizer(img).transform(img, *args)

    def __call__(self, img, size):
        return self.get_resizer(img, size)(img, size)

class ResizeFixedWidth(Resizer):
    def transformations(self, img, target_size):
        sw, sh = img.size
        tw, th = target_size
        if not self.expand and sw<=tw:
            return []
        h = sh*tw//sw
        return [('resize', (tw, h))]

class ResizeFixedHeight(Resizer):

    def transformations(self, img, target_size):
        sw, sh = img.size
        tw, th = target_size
        if not self.expand and  sh<=th:
            return []
        w = sw*th//sh
        return [('resize', (w, th))]


import ast
import inspect
import cast
from pycparser import c_ast

def constant(name):
    """
    Check if *name* is a known constant and return an equivalent
    c_ast.Constant(). If it is unknown, return None.
    """
    known = ('pi', 'e', 'ln2', 'ln10')

    if name in known:
        return c_ast.Constant('float', name)

    return None

class PythonToC(ast.NodeVisitor):
    def __init__(self):
        self.result = None

    def generic_visit(self, node):
        ops = {ast.Add: '+', ast.Sub: '-', ast.Mult: '*', ast.Div: '/',
               ast.Gt: '>', ast.GtE: '>=', ast.Lt: '<', ast.LtE: '<=',
               ast.Eq: '==', ast.NotEq: '!=',
               ast.And: '&&', ast.Or: '||',
               ast.Invert: '~', ast.Not: '!', ast.UAdd: '+', ast.USub: '-'}

        if type(node) in ops:
            self.result = ops[type(node)]
        else:
            for _, c in node.children():
                self.visit(c)

    def visit_Return(self, node):
        self.result = c_ast.Return(python_to_c_ast(node.value))

    def visit_UnaryOp(self, node):
        self.result = c_ast.UnaryOp(python_to_c_ast(node.op),
                                    python_to_c_ast(node.operand))

    def visit_BinOp(self, node):
        if isinstance(node.op, ast.Pow):
            self.result = c_ast.FuncCall(c_ast.ID('pow'),
                                         c_ast.ExprList([python_to_c_ast(node.left),
                                                         python_to_c_ast(node.right)]))
        else:
            self.result = c_ast.BinaryOp(python_to_c_ast(node.op),
                                         python_to_c_ast(node.left),
                                         python_to_c_ast(node.right))

    def visit_BoolOp(self, node):
        conds = [python_to_c_ast(v) for v in node.values]
        self.result = cast.chain(python_to_c_ast(node.op), conds)

    def visit_If(self, node):
        else_branch = None if not node.orelse else python_to_c_ast(node.orelse)
        self.result = c_ast.If(python_to_c_ast(node.test),
                               python_to_c_ast(node.body),
                               else_branch)

    def visit_For(self, node):
        if isinstance(node.iter, ast.Name):
            # We are iterating over a named iterable which we will convert to
            # an array iteration
            body = python_to_c_ast(node.body)
            self.result = c_ast.For(None, None, None, body)

            # Mark this for-loop to be processed once input data is known
            setattr(self.result, '_extra', (node.target.id, node.iter.id))
        elif isinstance(node.iter, ast.Call) and node.iter.func.id == 'range':
            var = node.target.id
            args = [python_to_c_ast(arg) for arg in node.iter.args]
            frm = c_ast.Constant('int', '0') if len(args) == 1 else args[0]
            to = args[0] if len(args) == 1 else args[1]
            step = c_ast.Constant('int', '1') if len(args) <= 2 else args[2]

            init = cast.TypeDecl(var, 'int', frm)
            cond = c_ast.BinaryOp('<', c_ast.ID(var), to)
            update = c_ast.ExprList([c_ast.BinaryOp('+=', c_ast.ID(var), step)])

            self.result = c_ast.For(init, cond, update, python_to_c_ast(node.body))

    def visit_Compare(self, node):
        self.result = c_ast.BinaryOp(python_to_c_ast(node.ops[0]),
                                     python_to_c_ast(node.left),
                                     python_to_c_ast(node.comparators[0]))

    def visit_Name(self, node):
        self.result = constant(node.id)

        if not self.result:
            self.result = c_ast.ID(node.id)

    def visit_Num(self, node):
        self.result = c_ast.Constant('int', str(node.n))

    def visit_Index(self, node):
        self.result = python_to_c_ast(node.value)

    def visit_Subscript(self, node):
        self.result = c_ast.ArrayRef(python_to_c_ast(node.value),
                                     python_to_c_ast(node.slice))

    def visit_Assign(self, node):
        # We resolve tuple assignments by creating an assignment for each tuple
        # pair. All assignments are then converted to a compound statement,
        # that means all variables _must_ be declared beforehands.
        if isinstance(node.targets[0], ast.Tuple):
            assignments = []

            lvalues = (python_to_c_ast(e) for e in node.targets[0].elts)
            rvalues = (python_to_c_ast(e) for e in node.value.elts)

            for lvalue, rvalue in zip(lvalues, rvalues):
                assignments.append(c_ast.Assignment('=', lvalue, rvalue))

            self.result = c_ast.Compound(assignments)
        else:
            lvalue = python_to_c_ast(node.targets[0])
            rvalue =  python_to_c_ast(node.value)
            self.result = c_ast.Assignment('=', lvalue, rvalue)

    def visit_AugAssign(self, node):
        op = python_to_c_ast(node.op)
        lvalue = python_to_c_ast(node.target)
        rvalue =  python_to_c_ast(node.value)
        self.result = c_ast.Assignment(op + '=', lvalue, rvalue)

    def visit_Tuple(self, node):
        elements = [python_to_c_ast(e) for e in node.elts]
        self.result = c_ast.ExprList(elements)

    def visit_Attribute(self, node):
        # strip off attribute accesses
        self.result = constant(node.attr)

        if not self.result:
            self.result = c_ast.ID(node.attr)

    def visit_Call(self, node):
        args = [python_to_c_ast(arg) for arg in node.args]
        exprs = c_ast.ExprList(args)
        ctypes = ['int', 'float']

        if isinstance(node.func, ast.Name) and node.func.id in ctypes:
            self.result = cast.CastDecl(node.func.id, exprs)
        else:
            self.result = c_ast.FuncCall(python_to_c_ast(node.func), exprs)

    def visit_FunctionDef(self, node):
        params = [python_to_c_ast(arg) for arg in node.args.args]
        rtype = c_ast.IdentifierType(['void'])
        tdecl = c_ast.TypeDecl(node.name, ['__kernel'], rtype)
        ftype = c_ast.FuncDecl(c_ast.ParamList(params), tdecl)
        decl = c_ast.Decl(node.name, None, None, None, ftype, None, None)
        body = python_to_c_ast(node.body)
        self.result = c_ast.FuncDef(decl, None, body)

    def visit_Expr(self, node):
        # Ignore docstrings which are standalone expressions with a Str object.
        pass

def python_to_c_ast(py_node):
    if isinstance(py_node, list):
        body = (python_to_c_ast(stmt) for stmt in py_node)
        return c_ast.Compound([stmt for stmt in body if stmt is not None])
    else:
        v = PythonToC()
        v.visit(py_node)
        return v.result

def parse(func):
    """
    Turn *func*'s Python AST into a pycparser AST for subsequent
    optimization and OpenCL code generation.
    """
    source = inspect.getsource(func)
    tree = ast.parse(source)
    return python_to_c_ast(tree.body[0])


# -*- encoding: utf-8 -*-
import sys
r_input = sys.stdin.readline

# 석환이가 부르는 노래의 N번째 단어는 무엇일까?
N = int(r_input())

repeat = N // 14
number = N % 14

if number == 1 or number == 13:
    print('baby')

elif number == 2 or number == 0:
    print('sukhwan')

elif number == 5:
    print('very')

elif number == 6:
    print('cute')

elif number == 9:
    print('in')

elif number == 10:
    print('bed')

else:
    ru_count = repeat + 1

    if number in [3, 7, 11]:
        ru_count += 1

    print('tu' + 'ru' * ru_count if ru_count < 5 else 'tu+ru*' + str(ru_count))

Algorithm/Baekjoon/15947 아기 석환 뚜루루 뚜루/15947.py


from django import template
from ..models import PartQuestion, Question

register = template.Library()

@register.inclusion_tag('secondapp/savollar.html')
def savollar(id):
    data = Question.objects.filter(partquestion_id__id=id)
    return {'data': data}


# uncompyle6 version 3.7.4
# Python bytecode 3.4 (3310)
# Decompiled from: Python 3.6.9 (default, Apr 18 2020, 01:56:04)
# [GCC 8.4.0]
# Embedded file name: /Library/Frameworks/Python.framework/Versions/3.4/lib/python3.4/site-packages/org/py4grid/scripts/py4gridserver_entrypoint.py
# Compiled at: 2014-09-03 22:05:30
# Size of source mod 2**32: 2581 bytes
"""
PY4GRID : a little framework to simule multiprocessing over a lot of computers
Copyright (C) 2014  João Jorge Pereira Farias Junior
This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.
This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU General Public License for more details.
You should have received a copy of the GNU General Public License
along with this program. If not, see <http://www.gnu.org/licenses/>.
"""
import argparse, socketserver, sys, os, signal
MODEL_BASE = None
if hasattr(os, 'fork'):
    MODEL_BASE = socketserver.ForkingTCPServer
else:
    MODEL_BASE = socketserver.ThreadingTCPServer
    MODEL_BASE.daemon_threads = True

def main(init_args=None):
    import org.py4grid.GPR as gpr, org.py4grid.multicast.daemons.DiscoverServers as discover
    discover_server = None
    server = None
    server_thread = None
    parse = argparse.ArgumentParser(prog='PY4GRIDSERVER')
    parse.add_argument('-port', help='a port to server usage to listen, default is 4680', default=4680, type=int)
    init_args = parse.parse_args(init_args)
    try:
        try:
            dic = {'port': init_args.port,  'ip': ''}
            host = (
             dic['ip'], init_args.port)
            print('HOST', host)
            server = MODEL_BASE(host, gpr.ProcessRemoteClient)

            def handle_controlz(signum, frame):
                print(signum, frame, 'CONTROl + Z')
                raise Exception('Press CTRL + Z')

            signal.signal(signal.SIGTSTP, handle_controlz)
            discover_server = discover.ServerDaemon(port=init_args.port)
            discover_server.start()
            print('SERVER_ADDRESS', server.server_address)
            server.serve_forever()
        except KeyboardInterrupt as kex:
            print('KEYBOARD interrupt...')
        except Exception as ex:
            print('Exception :', ex, '...')

    finally:
        print('Stoping DiscoverDaemon')
        discover_server.StopDaemon()
        print('DiscoverDaemon Stoped')
        print('Stoping PY4GRIDSERVER')
        server.server_close()
        print('PY4GRIDSERVER Stoped')
        sys.exit(0)

if __name__ == '__main__':
    print(sys.argv)
    main(init_args=sys.argv[1:])
    sys.exit(0)


import vpython as vp
from robot_imu import RobotImu, ImuFusion
from delta_timer import DeltaTimer
import imu_settings

imu = RobotImu(gyro_offsets=imu_settings.gyro_offsets)

fusion = ImuFusion(imu)

vp.graph(xmin=0, xmax=60, scroll=True)
graph_pitch = vp.gcurve(color=vp.color.red)
graph_roll = vp.gcurve(color=vp.color.green)

timer = DeltaTimer()
while True:
    vp.rate(100)
    dt, elapsed = timer.update()
    fusion.update(dt)
    graph_pitch.plot(elapsed, fusion.pitch)
    graph_roll.plot(elapsed, fusion.roll)


#! /usr/bin/env python
#

"""
Inspired by https://github.com/ufeindt/marshaltools
"""
import json
import requests
import pandas
import os
import numpy as np
from . import io

import matplotlib.pyplot as mpl
# list of effects:
#    list_program_sources.cgi | auth, data={'programidx' : str(programidx)}  # list of sources associated to the program
#    list_programs.cgi | auth # list of program you belong to
#    view_source.cgi | auth, data={'name' : name}
#    source_summary.cgi | auth, data={'sourceid' : str(source['id'])}
#    add_spec.cgi       | auth, data=payload,files=files
#                =>  payload = {'sourceid' : str(source['id']),'spectype':args.spectype,'programid':get_prog_id(args.prog_name),'instrumentid':get_inst_id(telname),
#                            'format':fformat,'obsdate':obsdate,'exptime':exptime,'observer':user,'reducedby':reducer,'class':"",'redshift':"",
#                            'phase':"",'comment':"",'commit':'yes','submit':'upload the file'}
#
#   growth_treasures_transient.cgi?cutprogramidx=%d
#
#marshal_root = 'http://skipper.caltech.edu:8080/cgi-bin/growth/'
#summary_url 	= marshal_root  + 'source_summary.cgi?sourceid=%s'
#listprog_url = marshal_root + 'list_programs.cgi' # list of program you belong to
#scanning_url = marshal_root + 'growth_treasures_transient.cgi'
#saving_url = marshal_root   + 'save_cand_growth.cgi?candid=%s&program=%s'
#savedsources_url = marshal_root + 'list_program_sources.cgi'
#rawsaved_url = marshal_root + 'list_sources_bare.cgi'
#annotate_url = marshal_root + 'edit_comment.cgi'
#ingest_url = marshal_root + 'ingest_avro_id.cgi'
#
#

MARSHAL_BASEURL = "http://skipper.caltech.edu:8080/cgi-bin/growth/"
from .io import LOCALSOURCE

def _account_id_declined_(username, password):
    """ This returns True if the login information has been rejected"""
    r = requests.post( MARSHAL_BASEURL+'list_programs.cgi',  auth=(username, password) )
    return "This server could not verify that you" in r.text

#############################
#                           #
# Stand Alone Functions     #
#                           #
#############################
def get_target_data(name):
    """ provide a name (or list of names) and get its/there marshal information
    IMPORTANT: This function is slow, but it takes the same amount of time if you provide 1 or any number of targets.
               So better provide a long list of target name at once.
    Returns
    -------
    pandas.DataFrame
    """
    m = MarshalAccess()
    m.load_target_sources()
    return m.get_target_data(name)

def download_spectra(name, dirout="default", auth=None, verbose=False, **kwargs):
    """Download all spectra for a source in the marshal as a tar.gz file

    Parameters:
    -----------
    name: [str]
        Name of a target on the marshal.

    dirout: [str] -optional-
        Directory where the data should be stored.
        Additional options:
        - `dirout=None`: The spectra are not saved be returned
        - `dirout='default'`: The spectra will be saved in native target location
                              (`$ZTFDATA`/marshal/spectra/`name`)
                              Spectra saved here can be recovered using `get_local_spectra`
                              * This is favored *

    auth: [str,str] -optional-
        Marshal [username, password]

    verbose: [bool] -optional-
        Prints to know what is going on.

    **kwargs goes to ztfquery.io.download_single_url()

    Returns
    -------
    None (or list of data if `dirout=None`)
    """
    # fileout is saved later to manage decompression
    import tarfile
    from io import BytesIO
    response = io.download_single_url(MARSHAL_BASEURL+'batch_spec.cgi',
                                   fileout=None,
                                   data={"name":name},
                                   auth=io._load_id_("marshal") if auth is None else auth,
                                   cookies="no_cookies", show_progress=False,
                                   **kwargs)

    tar = tarfile.open(fileobj=BytesIO( response.content ), mode='r')

    # No directory out? Then reformated data returned
    if dirout is None or dirout in ["None"]:
        if verbose: print("Data returned (dirout=None)")
        out = {member.name:tar.extractfile(member).read().decode("utf-8").splitlines() for member in tar.getmembers()}
        return out
    # Directory given, then dump data there:
    if dirout in ["default"]:
        dirout = target_spectra_directory(name)

    if verbose: print("Data will be stored here: %s"%dirout)
    if not os.path.exists(dirout):
        os.makedirs(dirout)

    tar.extractall(dirout)

# -------------- #
#  PLOT LC       #
# -------------- #

GENERIC = dict(alpha=1, mew=0.4, mec="0.7", ecolor="0.7", ls="None")
PROP    = { # ZTF
            "ZTF:r":dict(marker="o",ms=7,  mfc="C3"),
            "ZTF:g":dict(marker="o",ms=7,  mfc="C2"),
            "ZTF:i":dict(marker="o",ms=7, mfc="C1"),
            # Swift
            "UVOT:B":   dict(marker="s",  ms=5, mfc="C0"),
            "UVOT:u":   dict(marker="s",  ms=5, mfc=mpl.cm.Blues(0.7)),
            "UVOT:UVM2":dict(marker="s", ms=5, mfc=mpl.cm.Purples(0.6)),
            "UVOT:UVW2":dict(marker="s", ms=5, mfc=mpl.cm.Purples(0.8)),
            "UVOT:UVW1":dict(marker="s", ms=5, mfc=mpl.cm.Purples(0.4)),
            "UVOT:V":   dict(marker="s", ms=5, mfc=mpl.cm.Greens(0.9)),
            #
            "IOO:u":   dict(marker="d", ms=6,mfc=mpl.cm.Blues(0.6)),
            "IOO:g":   dict(marker="d", ms=6,mfc=mpl.cm.Greens(0.6)),
            "IOO:r":   dict(marker="d", ms=6,mfc=mpl.cm.Reds(0.7)),
            "IOO:i":   dict(marker="d",ms=6, mfc=mpl.cm.Oranges(0.6)),
            "IOO:z":   dict(marker="d", ms=6,mfc=mpl.cm.binary(0.8))
            }
for v in PROP.values():
    for k,v_ in GENERIC.items():
        v[k]=v_

def plot_lightcurve(lc_dataframe, ax=None, title=None, show_legend=True):
    """ """
    import matplotlib.pyplot as mpl
    from astropy.time import Time
    if ax is None:
        fig = mpl.figure(figsize=[7,4])
        ax  = fig.add_axes([0.1,0.12,0.67,0.8])
    else:
        fig = ax.figure

    lc_dataframe["inst_filter"] = [d.split("+")[-1].replace('"',"")
                            for d in lc_dataframe["instrument"]+":"+lc_dataframe["filter"]]

    # DataPoints
    for filter_ in np.unique(lc_dataframe["inst_filter"]):
        if filter_ not in  PROP:
            print("WARNING: Unknown instrument: %s | magnitude not shown"%filter_)
            continue

        jd, mag, magerr = lc_dataframe[lc_dataframe["inst_filter"].isin([filter_]) &
                                       ~lc_dataframe["magpsf"].isin([99.00])][
                            ["jdobs","magpsf","sigmamagpsf"]
                        ].values.T

        ax.errorbar([Time(jd_, format="jd").datetime for jd_ in jd],
                     mag, yerr= magerr,
                     label="%s"%filter_, **PROP[filter_.replace('"',"")])
    # Upper Limits
    ax.invert_yaxis()
    for filter_ in np.unique(lc_dataframe["inst_filter"]):
        if filter_ not in  PROP:
            print("WARNING: Unknown instrument: %s | upper limits not shown"%filter_)
            continue

        jdup, upmag = lc_dataframe[lc_dataframe["inst_filter"].isin([filter_]) &
                                 lc_dataframe["magpsf"].isin([99.00])][
                            ["jdobs","limmag"]
                        ].values.T
        ax.errorbar([Time(jd_, format="jd").datetime for jd_ in jdup],
                                upmag, yerr=0.15, lolims=True,alpha=0.3,
                                    color=PROP[filter_.replace('"',"")]["mfc"],
                            ls="None",
                                    label="_no_legend_")

    ax.set_ylabel("magnitude", fontsize="large")
    ax.set_xlabel("Time", fontsize="large")
    if title is not None:
        ax.set_title(title)
    if show_legend:
        ax.legend(loc=[1.02,0.], fontsize="medium" )
    return {"ax":ax, "fig":fig}

# -------------- #
#  Data I/O      #
# -------------- #
# - What at ?
def target_spectra_directory(name):
    """ where Marshal spectra are stored """
    return LOCALSOURCE+"marshal/spectra/%s/"%name

def target_lightcurves_directory(name):
    """ where Marshal lightcurves are stored """
    return LOCALSOURCE+"marshal/lightcurves/%s/"%name

# - Get the FullPathes
def get_local_spectra(name):
    """ returns list of fullpath of spectra on your computer for the given target name.
    Remark: These spectra have to be stored in the native `$ZTFDATA`/marshal/spectra/`name`
    """
    dir_ = target_spectra_directory(name)
    return {d:open(dir_+"%s"%d).read().splitlines() for d in os.listdir( dir_ )}

def get_local_lightcurves(name):
    """ returns list of fullpath of lightcurves on your computer for the given target name.
    Remark: These lightcurves have to be stored in the native `$ZTFDATA`/marshal/lightcurves/`name`
    """
    dir_ = target_lightcurves_directory(name)

    return {d: pandas.read_csv(dir_+d) for d in os.listdir( dir_ )}

# -------------- #
#  Downloading   #
# -------------- #
def download_lightcurve(name, dirout="default",
                            auth=None, verbose=False, **kwargs):
    """Download all spectra for a source in the marshal as a tar.gz file

    Parameters:
    -----------
    name: [str]
        Name of a target on the marshal.

    dirout: [str] -optional-
        Directory where the data should be stored.
        Additional options:
        - `dirout=None`: The spectra are not saved be returned
        - `dirout='default'`: The lightcurve will be saved in native target location
                              (`$ZTFDATA`/marshal/lightcurves/`name`)
                              lightcurve saved here can be recovered using `get_local_lightcurve`
                              * This is favored *

    auth: [str,str] -optional-
        Marshal [username, password]

    verbose: [bool] -optional-
        Prints to know what is going on.

    **kwargs goes to ztfquery.io.download_single_url()

    Returns
    -------
    None (or pandas.DataFrame)
    """
    # fileout is saved later to manage decompression

    response = io.download_single_url(MARSHAL_BASEURL+'print_lc.cgi',
                                   fileout=None,
                                   data={"name":name},
                                   auth=io._load_id_("marshal") if auth is None else auth,
                                   cookies="no_cookies", show_progress=False,
                                   **kwargs)

    # Convert the response into DataFrame
    data = response.text.split("
")[-1].replace(' ', '').replace('\n', '').split("
")
    dataframe = pandas.DataFrame(data=[d.split(",")[:8] for d in data[1:] if len(d)>0], columns=data[0].split(",")[:8])
    # returns it
    if dirout is None or dirout in ["None"]:
        return dataframe

    # Directory given, then dump data there:
    if dirout in ["default"]:
        dirout = target_lightcurves_directory(name)

    if verbose: print("Data will be stored here: %s"%dirout)
    if not os.path.exists(dirout):
        os.makedirs(dirout)

    dataframe.to_csv(dirout+"marshal_lightcurve_%s.csv"%name)

#############################
#                           #
#   Marshall Class          #
#                           #
#############################

class MarshalAccess( object ):
    """ Access the Marshal """
    def __init__(self, load_programs=False, **kwargs):
        """

        """
        if load_programs:
            self.load_user_programs( **kwargs )

    # -------------- #
    #  Main Methods  #
    # -------------- #
    #
    # DOWNLOADER
    #
    def download_spectra(self, name, dirout="default", auth=None,  **kwargs):
        """

        Method calling ztfquery.marshal.download_spectra()

        Parameters:
        -----------
        name: [str or list of]
            Name of a target on the marshal.

        dirout: [str] -optional-
            Directory where the data should be stored.
            Additional options:
            - `dirout=None`: The spectra are not saved be returned
            - `dirout='default'`: The spectra will be saved in native target location
                                  (`$ZTFDATA`/marshal/spectra/`name`)
                                 Spectra saved here can be recovered using `get_local_spectra`
                                 * This is favored *

        auth: [str,str] -optional-
            Marshal [username, password]

        **kwargs goes to ztfquery.io.download_single_url()

        Returns
        -------
        dict
        // {name: `return_of ztfquery.marshal.download_spectra()`}
        """
        return {name_: download_spectra(name, dirout=dirout, auth=auth, **kwargs) for name_ in np.atleast_1d(name)}

    def download_lightcurve(self, name, dirout="default", auth=None,  **kwargs):
        """

        Method calling ztfquery.marshal.download_lightcurve()

        Parameters:
        -----------
        name: [str or list of]
            Name of a target on the marshal.

        dirout: [str] -optional-
            Directory where the data should be stored.
            Additional options:
            - `dirout=None`: The spectra are not saved be returned
            - `dirout='default'`: The lightcurve will be saved in native target location
                              (`$ZTFDATA`/marshal/lightcurves/`name`)
                              lightcurve saved here can be recovered using `get_local_lightcurve`
                              * This is favored *

        auth: [str,str] -optional-
            Marshal [username, password]

        **kwargs goes to ztfquery.io.download_single_url()

        Returns
        -------
        dict
        // {name: `return_of ztfquery.marshal.download_lightcurve()`}
        """
        return {name_: download_lightcurve(name, dirout=dirout, auth=auth, **kwargs) for name_ in np.atleast_1d(name)}

    #
    # LOADER
    #
    def load_user_programs(self, auth=None):
        """ """
        if auth is None:
           auth = io._load_id_("marshal", askit=True)

        r = requests.post(MARSHAL_BASEURL+'list_programs.cgi',  auth=auth)
        r.raise_for_status() # raise a status if issue, like wrong auth

        self.program_data = pandas.DataFrame.from_dict(json.loads(r.text))

    def load_target_sources(self, program="*",
                            getredshift=True, getclassification=True,
                            setit=True, auth=None):
        """ download target source information and store them as a
            pandas.DataFrame as self.target_sources
            (or returns it, see setit parameter)

        Parameters
        ----------
        program: [str or list of] --optional--
            You want targets only associated to this program?
            e.g. program="Redshift Completeness Factor"
                 program=["AMPEL Test","Redshift Completeness Factor"]

            -> use program = None or program="*" for no program selection

        getredshift, getclassification: [bool, bool] -optional-
            If redshift and/or classification have been made in the marshal,
            do you want them ?

        setit: [bool] -optional-
            Do you want to set downloaded data to self.target_sources (`setit=True`, default)
            or would you prefer no directly get the pandas.DataFrame without
            touching to self.target_sources ? (`setit=False`)

        auth: [str,str] -optional-
            Marshal's [username, password]
            CAUTION: if you are requesting program(s), make sure the `auth`
                      matches that of your loaded program if already loaded.
                      Remark: If you did not load the user_program yet
                      (`self.load_user_programs()`),  they are authomatically matched.
        Returns
        -------
        None (or pandas.DataFrame if setit=False, see above)
        """

        r = requests.post(MARSHAL_BASEURL+'list_program_sources.cgi',
                       auth=io._load_id_("marshal", askit=True) if auth is None else auth,
                       data={'programidx': self._program_to_programidx_(program, auth=auth),
                             'getredshift': int(getredshift), 'getclassification': int(getclassification)})

        if setit:
            self.set_target_sources( pandas.DataFrame.from_dict(json.loads(r.text)) )
        else:
            return pandas.DataFrame.from_dict(json.loads(r.text))

    #
    # SETTER
    #
    def set_target_sources(self, target_source_dataframe):
        """ Provide a Pandas.DataFrame containing the target source information
            as obtained by `load_target_sources`
        """

        self.target_sources =  target_source_dataframe

    #
    # GETTER
    #
    def get_target_data(self, name):
        """ target_sources entry corresponding to the given name(s)

        Parameters
        ----------
        name: [str or list of]
            one or several target name(s) from `target_sources`

        Returns
        -------
        DataFrame Row(s)
        """
        return self.target_sources[self.target_sources["name"].isin(np.atleast_1d(name))]

    def get_target_coordinates(self, name):
        """ Target(s) coordinates ["ra", "dec"] in degree

        [simply  `self.get_target_data(name)[["ra","dec"]]` ]

        Parameters
        ----------
        name: [str or list of]
            one or several target name(s) from `target_sources`

        Returns
        -------
        Ra, Dec
        """
        return self.get_target_data(name)[["ra","dec"]]

    def get_target_redshift(self, name):
        """ Target(s) redshift as saved in the Marshal.
        # Caution, Redshift in the marshal might not be accurate. #

        [simply  `self.get_target_data(name)[["redshift"]]` ]

        Parameters
        ----------
        name: [str or list of]
            one or several target name(s) from `target_sources`

        Returns
        -------
        Redshift
        """
        return self.get_target_data(name)["redshift"]

    def get_target_classification(self, name):
        """

        [simply  `self.get_target_data(name)[["classification"]]` ]

        Parameters
        ----------
        name: [str or list of]
            one or several target name(s) from `target_sources`

        Returns
        -------
        classification
        // SN Ia, SN Ia 91bg-like, SN Ib, AGN, None, etc.
        """
        return self.get_target_data(name)["classification"]

    # -------------- #
    #  Internal      #
    # -------------- #
    def _program_to_programidx_(self, program, **kwargs):
        """ Returns the programidx corresponding to your program """

        if program is None or program in ['*','all']:
            programidx = "*"
        else:
            if not hasattr(self, "program_data"):
                self.load_user_programs(**kwargs)

            program  = np.atleast_1d(program)
            programidx = self.program_data[self.program_data["name"].isin(program)]["programidx"].values
            if len(programidx) == 0:
                raise ValueError("None of you program corresponds to "+", ".join(list(program)) +"\n"+
                                     "Your programs: "+", ".join(list(self.programs)))

        return programidx

    # =================== #
    #   Properties        #
    # =================== #
    # Programs
    @property
    def programs(self):
        """ """
        if not hasattr(self, "program_data"):
            raise AttributeError("You did not load the program_data. See load_user_programs()")
        return self.program_data["name"].values

    # Sources
    @property
    def nsources(self):
        """ """
        if not hasattr(self, "target_sources"):
            raise AttributeError("You did not load/set the target_sources. See {load,set}_target_sources()")
        return len(self.target_sources)