import os
import re
import time
import openai
import tiktoken
import google.generativeai as palm
import google.api_core.exceptions
from typing import Set, Tuple
from prompt import StreamPrompt
from colorama import Fore, Style
from argparse import Namespace

class Model(object):
    api_interval = 0.00 # seconds
    cost_per_1000tokens = 0.02
    gpt_models = {"text-davinci-003", "gpt-3.5-turbo-instruct"}
    palm_models = {"models/text-bison-001"}

    def __init__(self, config: Namespace):
        self._config = config
        if self._config.model in Model.gpt_models:
            self._tokenizer = tiktoken.encoding_for_model(self._config.model)
            self._original_max_tokens = self._config.max_tokens
            self._original_temperature = self._config.temperature
            openai.api_key = os.getenv("OPENAI_API_KEY")
        else:
            palm.configure(api_key=os.getenv("GOOGLE_API_KEY"))

    def retry_with_exponential_backoff(
        func,
        initial_delay: float = 0.5,
        exponential_base: float = 2,
        max_retries: int = 20,
        errors: tuple = (
            openai.error.RateLimitError, openai.error.APIError,
            google.api_core.exceptions.ResourceExhausted, google.api_core.exceptions.ServiceUnavailable, google.api_core.exceptions.GoogleAPIError,
        ),
    ):
        """Retry a function with exponential backoff."""

        def wrapper(*args, **kwargs):
            # Initialize variables
            num_retries = 0
            delay = initial_delay
            # Loop until a successful response or max_retries is hit or an exception is raised
            while True:
                try:
                    return func(*args, **kwargs)
                # Retry on specific errors
                except errors as e:
                    # Increment retries
                    num_retries += 1
                    # Check if max retries has been reached
                    if num_retries > max_retries:
                        raise Exception(
                            Fore.RED + f"Maximum number of retries ({max_retries}) exceeded." + Style.RESET_ALL
                        )
                    # Increment the delay
                    delay *= exponential_base
                    # Sleep for the delay
                    print(Fore.YELLOW + f"Error encountered. Retry ({num_retries}) after {delay} seconds..." + Style.RESET_ALL)
                    time.sleep(delay)
                # Raise exceptions for any errors not specified
                except Exception as e:
                    raise e

        return wrapper

    @retry_with_exponential_backoff
    def complete(self, prompt: str, label_set: Set[str] = set(), temperature: float = None) -> Tuple[str, str]:
        time.sleep(Model.api_interval)
        if self._config.model in Model.gpt_models:
            return self.gpt_complete(prompt, label_set, temperature)
        elif self._config.model in Model.palm_models:
            return self.palm_complete(prompt, label_set, temperature)
        else:
            raise Exception(Fore.RED + f"Model {self._config.model} not supported." + Style.RESET_ALL)

    def gpt_complete(self, prompt: str, label_set: Set[str] = set(), temperature: float = None) -> Tuple[str, str]:
        params = vars(self._config)
        if label_set:
            max_tokens = 0
            for label in label_set:
                tokens = label
                if '(' == tokens[0]:
                    tokens = tokens[1:] # remove parentheses
                    if prompt[-1] != '(':
                        prompt += " ("
                token_ids = self._tokenizer.encode(tokens)
                max_tokens = max(max_tokens, len(token_ids))
            if max_tokens > 0:
                # params["logit_bias"] = logit_bias
                if prompt[-1] != '(':
                    max_tokens += 1 # +1 for whitespace
                params["max_tokens"] = max_tokens
        else:
            params["max_tokens"] = self._original_max_tokens

        if temperature is not None:
            params["temperature"] = temperature
        else:
            params["temperature"] = self._original_temperature
        print(Fore.GREEN + f" (Predicting with temperature: {params['temperature']}, max_tokens: {params['max_tokens']}) " + Style.RESET_ALL, end='')
        return prompt, openai.Completion.create(prompt=prompt, **params)["choices"][0]["text"]

    def palm_complete(self, prompt: str, label_set: Set[str] = set(), temperature: float = None) -> Tuple[str, str]:
        params = {
            "model": self._config.model,
            "temperature": self._config.temperature if (temperature is None) else temperature,
            "candidate_count": 1,
            "top_p": self._config.top_p,
            "max_output_tokens": self._config.max_tokens,
        }
        response = palm.generate_text(
            **params,
            prompt=prompt,
            safety_settings=[  # set all safety settings to 4 (no restrictions)
                {"category": "HARM_CATEGORY_DEROGATORY", "threshold": 4},
                {"category": "HARM_CATEGORY_TOXICITY", "threshold": 4},
                {"category": "HARM_CATEGORY_VIOLENCE", "threshold": 4},
                {"category": "HARM_CATEGORY_SEXUAL", "threshold": 4},
                {"category": "HARM_CATEGORY_MEDICAL", "threshold": 4},
                {"category": "HARM_CATEGORY_DANGEROUS", "threshold": 4}
            ]
        )
        print(Fore.GREEN + f" (Predicting with temperature: {params['temperature']}, max_tokens: {params['max_output_tokens']}) " + Style.RESET_ALL)
        if response.result is None:
            print(Fore.RED + f"Failed to generated new instance." + Style.RESET_ALL)
            raise ValueError("Failed to generated new instance.")
        return prompt, response.result

    def set_api_key(self, key: str) -> None:
        openai.api_key = key

    def count_tokens(self, text: str) -> int:
        return len(self._tokenizer.encode(text))

# for unit testing
if __name__ == "__main__":
    # model
    config = Namespace(
        model="models/text-bison-001",
        max_tokens=1024,
        temperature=0.0,
        top_p=1.0,
    )
    model_api = Model(config)

    # tasks
    from task import TaskGenerator

    task_gen = TaskGenerator(
        task_input_path="./bbh/BIG-Bench-Hard/bbh/",
        task_desc_path="./bbh/bbh_task_description.json",
        batch_size=1
    )
    task_names = ["boolean_expressions", "causal_judgement", "date_understanding", "formal_fallacies", "sports_understanding"]

    # prompt
    for task_name in task_names:
        task = task_gen.get_task(task_name)
        task_inputs = task.get_new_inputs()
        stream_prompt = StreamPrompt(
            task_desc=task.task_desc,
            inputs=task_inputs,
            num_demos=3,
            shots=[]
        )

        pred_prompt = stream_prompt.gen_prediction()
        print(f"Generating prediction from label set: {task.label_set} ->")
        pred_prompt, res_text = model_api.complete(pred_prompt, task.label_set)
        full_text = pred_prompt + res_text
        print(f"full text: {full_text}")


from __future__ import annotations

import asyncio
import datetime
import imghdr
import logging
import mimetypes
from pathlib import Path
from typing import Any, List, Optional, Sequence, Type, Union

import telegram.ext
import tzlocal
import validators
from apscheduler.schedulers.base import BaseScheduler
from telegram import Bot, Chat, InlineKeyboardMarkup, Message, ReplyKeyboardMarkup, Update
from telegram._utils.defaultvalue import DEFAULT_NONE
from telegram._utils.types import ODVInput
from telegram.constants import ChatAction, ParseMode
from telegram.ext import Application, CallbackQueryHandler, CommandHandler, MessageHandler, PicklePersistence
from telegram.ext._callbackcontext import CallbackContext
from telegram.ext._utils.types import BD, BT, CD, UD

from ._version import __raw_url__
from .models import BaseMessage, ButtonType, TypeCallback, call_function_EAFP, emoji_replace

logger = logging.getLogger(__name__)

class NavigationException(Exception):
    """Base exception."""

class TelegramMenuSession:
    """Session manager, send start message to each new user connecting to the bot."""

    # delays in seconds
    READ_TIMEOUT = 6
    CONNECT_TIMEOUT = 7
    START_MESSAGE = "start"

    def __init__(self, api_key: str, start_message: str = START_MESSAGE, persistence_path: str = "") -> None:
        """Initialize the session object.

        Args:
            api_key: Telegram bot API key
            start_message: text used to start a session, e.g. /start
        """
        if not isinstance(api_key, str):
            raise KeyError("API_KEY must be a string!")

        persistence = PicklePersistence(filepath=persistence_path if persistence_path else "arbitrarycallbackdatabot")
        self.application = (
            Application.builder().token(api_key).persistence(persistence).arbitrary_callback_data(True).build()
        )
        self.scheduler = self.application.job_queue.scheduler  # type: ignore

        self._api_key = api_key
        self.sessions: List[NavigationHandler] = []
        self.start_message_class: Optional[Type[BaseMessage]] = None
        self.start_message_args: Optional[List[Any]] = None
        self.navigation_handler_class: Optional[Type[NavigationHandler]] = None

        # on different commands - answer in Telegram
        self.application.add_handler(CommandHandler(start_message, self._send_start_message))
        self.application.add_handler(MessageHandler(telegram.ext.filters.TEXT, self._button_select_callback))
        self.application.add_handler(
            MessageHandler(telegram.ext.filters.StatusUpdate.WEB_APP_DATA, self._button_webapp_callback)
        )
        self.application.add_handler(CallbackQueryHandler(self._button_inline_select_callback))
        self.application.add_handler(telegram.ext.PollAnswerHandler(self._poll_answer))
        self.application.add_error_handler(self._msg_error_handler)
        self.application.add_handler(MessageHandler(telegram.ext.filters.LOCATION, self._get_location_handler))

    def start(
        self,
        start_message_class: Type[BaseMessage],
        start_message_args: Optional[List[Any]] = None,
        polling: bool = True,
        navigation_handler_class: Optional[Type[NavigationHandler]] = None,
        stop_signals: ODVInput[Sequence[int]] = DEFAULT_NONE,
    ) -> None:
        """Set the start message and run the dispatcher.

        Args:
            start_message_class: class used to create start message
            start_message_args: optional arguments passed to the start class
            polling: if True, start polling updates from Telegram
            navigation_handler_class: optional class used to extend the base NavigationHandler

        """
        self.start_message_class = start_message_class
        self.start_message_args = start_message_args
        self.navigation_handler_class = navigation_handler_class or NavigationHandler
        if not issubclass(start_message_class, BaseMessage):
            raise NavigationException("start_message_class must be a BaseMessage!")
        if start_message_args is not None and not isinstance(start_message_args, list):
            raise NavigationException("start_message_args is not a list!")
        if not issubclass(self.navigation_handler_class, NavigationHandler):
            raise NavigationException("navigation_handler_class must be a NavigationHandler!")

        if not self.scheduler.running:
            self.scheduler.start()
        if polling:
            self.application.run_polling(stop_signals=stop_signals)

    async def _send_start_message(self, update: Update, context: CallbackContext[BT, UD, CD, BD]) -> None:
        """Start main message, app choice."""
        chat = update.effective_chat
        if chat is None:
            raise NavigationException("Chat object was not created")
        if self.navigation_handler_class is None:
            raise NavigationException("Navigation Handler class not defined")
        session = self.navigation_handler_class(self.application.bot, chat, self.scheduler)
        self.sessions.append(session)
        if self.start_message_class is None:
            raise NavigationException("Message class not defined")
        if self.start_message_args is not None:
            start_message = self.start_message_class(session, message_args=self.start_message_args)
        else:
            start_message = self.start_message_class(session)
        await session.goto_menu(start_message, context)

    def get_session(self, chat_id: int = 0) -> Optional[NavigationHandler]:
        """Get session from list."""
        sessions = [x for x in self.sessions if chat_id in (x.chat_id, 0)]
        if not sessions:
            return None
        return sessions[0]

    async def _get_location_handler(self, update: Update, context: CallbackContext[BT, UD, CD, BD]) -> None:
        if update.effective_chat is None or update.message is None or update.message.location is None:
            raise NavigationException("Incorrect session, location can't be updated")
        session = self.get_session(update.effective_chat.id)
        if session is not None:
            session.location = update.message.location

    async def _button_select_callback(self, update: Update, context: CallbackContext[BT, UD, CD, BD]) -> None:
        """Menu message main entry point."""
        if update.effective_chat is None or update.message is None:
            raise NavigationException("Chat object was not created")
        session = self.get_session(update.effective_chat.id)
        if session is None:
            await self._send_start_message(update, context)
            return
        if update.message.text:
            await session.select_menu_button(update.message.text, context)

    async def _poll_answer(self, update: Update, _: CallbackContext[BT, UD, CD, BD]) -> None:
        """Entry point for poll selection."""
        if update.effective_user is None or update.poll_answer is None:
            raise NavigationException("User object was not created")
        session = next((x for x in self.sessions if x.user_name == update.effective_user.first_name), None)
        if session:
            await session.poll_answer(update.poll_answer.option_ids[0])

    async def _button_inline_select_callback(self, update: Update, context: CallbackContext[BT, UD, CD, BD]) -> None:
        """Execute inline callback of an BaseMessage."""
        if update.effective_chat is None or update.callback_query is None:
            raise NavigationException("Chat object was not created")
        session = self.get_session(update.effective_chat.id)
        if session is None:
            await self._send_start_message(update, context)
            return
        if update.callback_query.data and update.callback_query.id:
            await session.app_message_button_callback(update.callback_query.data, update.callback_query.id, context)

    async def _button_webapp_callback(self, update: Update, context: CallbackContext[BT, UD, CD, BD]) -> None:
        """Execute webapp callback."""
        if (
            update.effective_chat is None
            or update.effective_message is None
            or update.effective_message.web_app_data is None
        ):
            raise NavigationException("Chat object was not created")
        session = self.get_session(update.effective_chat.id)
        if session is None:
            await self._send_start_message(update, context)
            return
        await session.app_message_webapp_callback(
            update.effective_message.web_app_data.data, update.effective_message.web_app_data.button_text
        )

    @staticmethod
    async def _msg_error_handler(update: object, context: CallbackContext[BT, UD, CD, BD]) -> None:  # type: ignore
        """Log Errors caused by Updates."""
        if not isinstance(update, Update):
            raise NavigationException("Incorrect update object")
        error_message = str(context.error) if update is None else f"Update {update.update_id} - {str(context.error)}"
        logger.error(error_message)

    async def broadcast_message(self, message: str, notification: bool = True) -> List[telegram.Message]:
        """Broadcast simple message without keyboard markup to all sessions."""
        messages = []
        for session in self.sessions:
            msg = await session.send_message(message, notification=notification)
            if msg is not None:
                messages.append(msg)
        return messages

    async def broadcast_picture(self, picture_path: str, notification: bool = True) -> List[telegram.Message]:
        """Broadcast picture to all sessions."""
        messages = []
        for session in self.sessions:
            msg = await session.send_photo(picture_path, notification=notification)
            if msg is not None:
                messages.append(msg)
        return messages

    async def broadcast_sticker(self, sticker_path: str, notification: bool = True) -> List[telegram.Message]:
        """Broadcast picture to all sessions."""
        messages = []
        for session in self.sessions:
            msg = await session.send_sticker(sticker_path, notification=notification)
            if msg is not None:
                messages.append(msg)
        return messages

class NavigationHandler:
    """Navigation handler for Telegram application."""

    POLL_DEADLINE = 10  # seconds
    MESSAGE_CHECK_TIMEOUT = 10  # seconds
    CONNECTION_POOL_SIZE = 8

    def __init__(self, bot: Bot, chat: Chat, scheduler: BaseScheduler) -> None:
        """Init NavigationHandler class."""
        self._bot = bot
        self._poll: Optional[Message] = None
        self._poll_callback: Optional[TypeCallback] = None

        self.scheduler = scheduler
        self.chat_id = chat.id
        self.user_name = chat.first_name
        self.poll_name = f"poll_{self.user_name}"
        self.location: Optional[telegram.Location] = None

        logger.info(f"Opening chat with user {self.user_name}")

        self._menu_queue: List[BaseMessage] = []  # list of menus selected by user
        self._message_queue: List[BaseMessage] = []  # list of application messages sent

        # check if messages have expired every MESSAGE_CHECK_TIMEOUT seconds
        scheduler.add_job(
            self._expiry_date_checker,
            "interval",
            id="state_nav_update",
            seconds=self.MESSAGE_CHECK_TIMEOUT,
            replace_existing=True,
        )

    async def _expiry_date_checker(self) -> None:
        """Check expiry date of message and delete if expired."""
        for message in self._message_queue:
            if message.has_expired():
                await self._delete_queued_message(message)

        # go back to home after sub-menu message has expired
        if len(self._menu_queue) >= 2 and self._menu_queue[-1].has_expired():
            await self.goto_home()

    async def delete_message(self, message_id: int) -> None:
        """Delete a message from its id."""
        await self._bot.delete_message(chat_id=self.chat_id, message_id=message_id)

    async def _delete_queued_message(self, message: BaseMessage) -> None:
        """Delete a message, remove from queue."""
        message.kill_message()
        if message in self._message_queue:
            self._message_queue.remove(message)
            await self.delete_message(message.message_id)
        del message

    async def goto_menu(
        self, menu_message: BaseMessage, context: Optional[CallbackContext[BT, UD, CD, BD]] = None
    ) -> int:
        """Send menu message and add to queue."""
        content = await menu_message.get_updated_content(context)
        logger.info(f"Opening menu {menu_message.label}")
        keyboard = menu_message.gen_keyboard_content()
        if menu_message.picture:
            message = await self.send_photo(
                menu_message.picture, notification=menu_message.notification, keyboard=keyboard, caption=content
            )
        else:
            message = await self.send_message(content, keyboard, notification=menu_message.notification)
        if message is None:
            return -1  # message was not sent, abort
        menu_message.is_alive()
        menu_message.message_id = message.message_id
        self._menu_queue.append(menu_message)
        return message.message_id

    async def goto_home(self, context: Optional[CallbackContext[BT, UD, CD, BD]] = None) -> int:
        """Go to home menu, empty menu_queue."""
        if not self._menu_queue:
            return -1
        if len(self._menu_queue) == 1:
            # already at 'home' level
            return self._menu_queue[0].message_id
        menu_previous = self._menu_queue.pop()
        while self._menu_queue:
            menu_previous = self._menu_queue.pop()
        return await self.goto_menu(menu_previous, context)

    @staticmethod
    def filter_unicode(input_string: str) -> str:
        """Remove non-unicode characters from input string."""
        return input_string.encode("ascii", "ignore").decode("utf-8")

    async def _send_app_message(
        self, message: BaseMessage, label: str, context: Optional[CallbackContext[BT, UD, CD, BD]] = None
    ) -> int:
        """Send an application message."""
        content = await message.get_updated_content(context)
        # if message with this label already exist in message_queue, delete it and replace it
        info_message = self.filter_unicode(f"Send message '{message.label}': '{label}'")
        logger.info(str(info_message))
        if "_" not in message.label:
            message.label = f"{message.label}_{label}"

        # delete message if already displayed
        message_existing = self.get_message(message.label)
        if message_existing is not None:
            await self._delete_queued_message(message)

        message.is_alive()

        keyboard = message.gen_inline_keyboard_content()
        if message.picture:
            msg = await self.send_photo(
                message.picture, notification=message.notification, caption=content, keyboard=keyboard
            )
        else:
            msg = await self.send_message(content, keyboard, message.notification)
        if msg is None:
            return -1  # message was not sent, abort
        message.message_id = msg.message_id
        self._message_queue.append(message)

        message.content_previous = content
        message.keyboard_previous = message.keyboard.copy()
        return message.message_id

    async def send_message(
        self,
        content: str,
        keyboard: Optional[Union[ReplyKeyboardMarkup, InlineKeyboardMarkup]] = None,
        notification: bool = True,
    ) -> telegram.Message:
        """Send a text message with html formatting."""
        return await self._bot.send_message(
            chat_id=self.chat_id,
            text=content,
            parse_mode=ParseMode.HTML,
            reply_markup=keyboard,
            disable_notification=not notification,
        )

    async def edit_message(
        self, message: BaseMessage, context: Optional[CallbackContext[BT, UD, CD, BD]] = None
    ) -> bool:
        """Edit an inline message asynchronously."""
        message_updt = self.get_message(message.label)
        if message_updt is None:
            return False

        # check if content and keyboard have changed since previous message
        content = await message_updt.get_updated_content(context)
        if not self._message_check_changes(message_updt, content):
            return False

        keyboard_format = message_updt.gen_inline_keyboard_content()
        try:
            if message_updt.picture:
                await self._bot.edit_message_caption(
                    caption=content,
                    chat_id=self.chat_id,
                    message_id=message_updt.message_id,
                    parse_mode=ParseMode.HTML,
                    reply_markup=keyboard_format,
                )
            else:
                await self._bot.edit_message_text(
                    text=content,
                    chat_id=self.chat_id,
                    message_id=message_updt.message_id,
                    parse_mode=ParseMode.HTML,
                    reply_markup=keyboard_format,
                )
        except telegram.error.BadRequest as error:
            logger.error(error)
            return False
        return True

    @staticmethod
    def _message_check_changes(message: BaseMessage, content: str) -> bool:
        """Check is message content and keyboard has changed since last edit."""
        content_identical = content == message.content_previous
        keyboard_identical = [y.label for x in message.keyboard_previous for y in x] == [
            y.label for x in message.keyboard for y in x
        ]
        if content_identical and keyboard_identical:
            return False
        message.content_previous = content
        message.keyboard_previous = message.keyboard.copy()
        return True

    async def select_menu_button(
        self, label: str, context: Optional[CallbackContext[BT, UD, CD, BD]] = None
    ) -> Optional[int]:  # noqa: C901
        """Select menu button using label."""
        msg_id = 0
        if label == "Back":
            if len(self._menu_queue) == 1:
                # already at 'home' level
                return self._menu_queue[0].message_id
            menu_previous = self._menu_queue.pop()  # delete actual menu
            if self._menu_queue:
                menu_previous = self._menu_queue.pop()
            return await self.goto_menu(menu_previous, context)
        if label == "Home":
            return await self.goto_home(context)

        for menu_item in self._menu_queue[::-1]:
            btn = menu_item.get_button(label)
            if not btn:
                continue
            if isinstance(btn.callback, BaseMessage):
                if btn.callback.inlined:
                    msg_id = await self._send_app_message(btn.callback, label, context)
                    if btn.callback.home_after:
                        msg_id = await self.goto_home(context)
                else:
                    msg_id = await self.goto_menu(btn.callback, context)
            elif btn.callback is not None and hasattr(btn.callback, "__call__"):
                if btn.args is not None:
                    await call_function_EAFP(btn.callback, context, btn.args)
                else:
                    await call_function_EAFP(btn.callback, context)
            return msg_id

        # label does not match any sub-menu, just process the user input
        await self.capture_user_input(label, context)
        return None

    async def capture_user_input(self, label: str, context: Optional[CallbackContext[BT, UD, CD, BD]] = None) -> None:
        """Process the user input in the last message updated."""
        last_menu_message = self._menu_queue[-1]
        if self._message_queue:
            for last_app_message in self._message_queue[::-1]:
                if last_app_message.time_alive > last_menu_message.time_alive:
                    last_menu_message = last_app_message
        await last_menu_message.text_input(label, context)

    async def app_message_webapp_callback(self, webapp_data: str, button_text: str) -> None:
        """Execute the callback associated to this webapp."""
        last_menu = self._menu_queue[-1]
        webapp_message = next(iter(y for x in last_menu.keyboard for y in x if y.label == button_text), None)
        if webapp_message is not None and callable(webapp_message.callback):
            if asyncio.iscoroutinefunction(webapp_message.callback):
                html_response = await webapp_message.callback(webapp_data)
            else:
                html_response = webapp_message.callback(webapp_data)
            await self.send_message(html_response, notification=webapp_message.notification)

    async def app_message_button_callback(
        self, callback_label: str, callback_id: str, context: Optional[CallbackContext[BT, UD, CD, BD]] = None
    ) -> None:
        """Entry point to execute an action after message button selection."""
        label_message, label_action = callback_label.split(BaseMessage.SEPARATOR)
        log_message = self.filter_unicode(f"Received action request from '{label_message}': '{label_action}'")
        logger.info(log_message)
        message = self.get_message(label_message)
        if message is None:
            logger.error(f"Message with label {label_message} not found, return")
            return
        btn = message.get_button(label_action)

        if btn is None:
            logger.error(f"No button found with label {label_action}, return")
            return

        if btn.btype in [ButtonType.PICTURE, ButtonType.STICKER]:
            # noinspection PyTypeChecker
            await self._bot.send_chat_action(chat_id=self.chat_id, action=ChatAction.UPLOAD_PHOTO)
        elif btn.btype == ButtonType.MESSAGE:
            # noinspection PyTypeChecker
            await self._bot.send_chat_action(chat_id=self.chat_id, action=ChatAction.TYPING)
        elif btn.btype == ButtonType.POLL:
            await self.send_poll(question=btn.args[0], options=btn.args[1])
            self._poll_callback = btn.callback
            await self._bot.answer_callback_query(callback_id, text="Select an answer...")
            return

        if not callable(btn.callback):
            return

        if btn.args is not None:
            action_status = await call_function_EAFP(btn.callback, context, btn.args)
        else:
            action_status = await call_function_EAFP(btn.callback, context)

        # send picture if custom label found
        if btn.btype == ButtonType.PICTURE:
            await self.send_photo(picture_path=action_status, notification=btn.notification)
            await self._bot.answer_callback_query(callback_id, text="Picture sent!")
            return
        if btn.btype == ButtonType.STICKER:
            await self.send_sticker(sticker_path=action_status, notification=btn.notification)
            await self._bot.answer_callback_query(callback_id, text="Sticker sent!")
            return
        if btn.btype == ButtonType.MESSAGE:
            await self.send_message(action_status, notification=btn.notification)
            await self._bot.answer_callback_query(callback_id, text="Message sent!")
            return
        await self._bot.answer_callback_query(callback_id, text=action_status)

        # update expiry period and update
        message.is_alive()
        await self.edit_message(message, context)

    @staticmethod
    def _sticker_check_replace(sticker_path: str) -> Union[str, bytes]:
        """Check if sticker is correct file type."""
        try:
            if not sticker_path.lower().endswith(".webp"):
                raise ValueError("Sticker has no .webp format")
            if validators.url(sticker_path):
                # todo: add check if url exists
                return sticker_path
            if Path(sticker_path).is_file() and imghdr.what(sticker_path):
                with open(sticker_path, "rb") as file_h:
                    return file_h.read()
            raise ValueError("Path is not a picture")
        except ValueError:
            url_default = f"{__raw_url__}/resources/stats_default.webp"
            logger.error(f"Picture path '{sticker_path}' is invalid, replacing with default {url_default}")
            return url_default

    @staticmethod
    def _picture_check_replace(picture_path: str) -> Union[str, bytes]:
        """Check if the given picture path or uri is correct, replace by default if not."""
        try:
            if validators.url(picture_path):
                # check if the url has image format
                mimetype, _ = mimetypes.guess_type(picture_path)
                if mimetype and mimetype.startswith("image"):
                    return picture_path
                raise ValueError("Url is not a picture")
            if Path(picture_path).is_file() and imghdr.what(picture_path):
                with open(picture_path, "rb") as file_h:
                    return file_h.read()
            raise ValueError("Path is not a picture")
        except ValueError:
            url_default = f"{__raw_url__}/resources/stats_default.png"
            logger.error(f"Picture path '{picture_path}' is invalid, replacing with default {url_default}")
            return url_default

    async def send_photo(
        self,
        picture_path: str,
        notification: bool = True,
        caption: str = "",
        keyboard: Optional[Union[ReplyKeyboardMarkup, InlineKeyboardMarkup]] = None,
    ) -> Optional[telegram.Message]:
        """Send a picture."""
        picture_obj = self._picture_check_replace(picture_path=picture_path)
        try:
            return await self._bot.send_photo(
                chat_id=self.chat_id,
                photo=picture_obj,
                caption=caption,
                reply_markup=keyboard,
                disable_notification=not notification,
                parse_mode=ParseMode.HTML,
            )
        except telegram.error.BadRequest as error:
            logger.error(f"Failed to send picture {picture_path}: {error}")
        return None

    async def send_sticker(self, sticker_path: str, notification: bool = True) -> Optional[telegram.Message]:
        """Send a picture."""
        sticker_obj = self._sticker_check_replace(sticker_path=sticker_path)
        try:
            return await self._bot.send_sticker(
                chat_id=self.chat_id, sticker=sticker_obj, disable_notification=not notification
            )
        except telegram.error.BadRequest as error:
            logger.error(f"Failed to send picture {sticker_path}: {error}")
        return None

    def get_message(self, label_message: str) -> Optional[BaseMessage]:
        """Get message from message_queue matching attribute label_message."""
        return next(iter(x for x in self._message_queue if x.label == label_message), None)

    async def send_poll(self, question: str, options: List[str]) -> None:
        """Send poll to user with question and options."""
        if self.scheduler.get_job(self.poll_name) is not None:
            await self.poll_delete()
        options = [emoji_replace(x) for x in options]
        self._poll = await self._bot.send_poll(
            chat_id=self.chat_id,
            question=emoji_replace(question),
            options=options,
            is_anonymous=False,
            open_period=self.POLL_DEADLINE,
        )
        self.scheduler.add_job(
            self.poll_delete,
            "date",
            id=self.poll_name,
            next_run_time=datetime.datetime.now(tz=tzlocal.get_localzone())
            + datetime.timedelta(seconds=self.POLL_DEADLINE + 1),
            replace_existing=True,
        )

    async def poll_delete(self) -> None:
        """Run when poll timeout has expired."""
        if self._poll is not None and self._poll.poll is not None:
            try:
                logger.info(f"Deleting poll '{self._poll.poll.question}'")
                await self._bot.delete_message(chat_id=self.chat_id, message_id=self._poll.message_id)
            except telegram.error.BadRequest:
                logger.error(f"Poll message {self._poll.message_id} already deleted")

    async def poll_answer(self, answer_id: int) -> None:
        """Run when poll message is received."""
        if (
            self._poll is None
            or self._poll.poll is None
            or self._poll_callback is None
            or not callable(self._poll_callback)
        ):
            logger.error("Poll is not defined")
            return

        answer_ascii = self._poll.poll.options[answer_id].text.encode("ascii", "ignore").decode()
        logger.info(f"{self.user_name}'s answer to question '{self._poll.poll.question}' is '{answer_ascii}'")
        if asyncio.iscoroutinefunction(self._poll_callback):
            await self._poll_callback(self._poll.poll.options[answer_id].text)
        else:
            self._poll_callback(self._poll.poll.options[answer_id].text)
        await asyncio.sleep(1)
        await self.poll_delete()

        if self.scheduler.get_job(self.poll_name) is not None:
            self.scheduler.remove_job(self.poll_name)
        self._poll = None


from datetime import datetime, timedelta
import os
from airflow import DAG
from airflow.operators.dummy_operator import DummyOperator
from airflow.operators import (StageToRedshiftOperator, LoadFactOperator,
                               LoadDimensionOperator, DataQualityOperator)
from helpers import SqlQueries

default_args = {
    'owner': 'guido',
    'start_date': datetime(2019, 1, 12),
    'retries': 3,
    'retry_delay': timedelta(minutes=5),
    'email_on_retry': False,
    'depends_on_past': False
}

dag = DAG('udac_example_dag',
          default_args=default_args,
          description='Load and transform data in Redshift with Airflow',
          schedule_interval='0 * * * *',
          catchup=False
          )

# Mark start of pipeline
start_operator = DummyOperator(task_id='Begin_execution',  dag=dag)

# Populate events table with origin data from S3
stage_events_to_redshift = StageToRedshiftOperator(
    task_id="load_events_from_s3_to_redshift",
    aws_credentials_id="aws_credentials",
    dag=dag,
    redshift_conn_id="redshift",
    s3_bucket="udacity-dend",
    s3_key="log_data",
    table="staging_events",
    json_path='s3://udacity-dend/log_json_path.json'
)

# Populate songs table with origin data from S3
stage_songs_to_redshift = StageToRedshiftOperator(
    task_id="load_songs_from_s3_to_redshift",
    aws_credentials_id="aws_credentials",
    dag=dag,
    redshift_conn_id="redshift",
    s3_bucket="udacity-dend",
    s3_key="song_data",
    ignore_headers="0",
    table="staging_songs",
)

# Extract data for songplays table and insert
load_songplays_table = LoadFactOperator(
    task_id='Load_songplays_fact_table',
    dag=dag,
    redshift_conn_id="redshift",
    sql_query=SqlQueries.songplay_table_insert,
    table='songplays',
)

# Extract data for users table and insert
load_user_dimension_table = LoadDimensionOperator(
    task_id='Load_user_dim_table',
    dag=dag,
    redshift_conn_id="redshift",
    sql_query=SqlQueries.user_table_insert,
    table='users',
)

# Extract data for songs table and insert
load_song_dimension_table = LoadDimensionOperator(
    task_id='Load_song_dim_table',
    dag=dag,
    redshift_conn_id="redshift",
    sql_query=SqlQueries.song_table_insert,
    table='songs',
)

# Extract data for artists table and insert
load_artist_dimension_table = LoadDimensionOperator(
    task_id='Load_artist_dim_table',
    dag=dag,
    redshift_conn_id="redshift",
    sql_query=SqlQueries.artist_table_insert,
    table='artists',
)

# Extract data for time table and insert
load_time_dimension_table = LoadDimensionOperator(
    task_id='Load_time_dim_table',
    dag=dag,
    redshift_conn_id="redshift",
    sql_query=SqlQueries.time_table_insert,
    table='time',
)

# Make sure tables contain at least one record
run_quality_checks = DataQualityOperator(
    task_id='Run_data_quality_checks',
    dag=dag,
    redshift_conn_id="redshift",
    tables=['artists', 'songplays', 'songs', 'time', 'users'],
)

# Mark end of pipeline
end_operator = DummyOperator(task_id='Stop_execution',  dag=dag)

# Dag orchestrator
start_operator >> stage_events_to_redshift
start_operator >> stage_songs_to_redshift

stage_events_to_redshift >> load_songplays_table
stage_songs_to_redshift >> load_songplays_table

load_songplays_table >> load_song_dimension_table
load_songplays_table >> load_user_dimension_table
load_songplays_table >> load_artist_dimension_table
load_songplays_table >> load_time_dimension_table

load_user_dimension_table >> run_quality_checks
load_song_dimension_table >> run_quality_checks
load_artist_dimension_table >> run_quality_checks
load_time_dimension_table >> run_quality_checks

run_quality_checks >> end_operator


import matplotlib.pyplot as plt
import basic_file_app
import numpy as np
import os
import px_shift_on_picture_array_rolling

PP = basic_file_app.load_1d_array("AVG20220815_ZnOIIIevenThres/avg20220815_ZnOIIIevenThreshold.txt", 0, 1)
UP = basic_file_app.load_1d_array("AVG20220815_ZnOIIIevenThres/avg20220815_ZnOIIIoddThreshold.txt", 0,1)

shift_avg = px_shift_on_picture_array_rolling.PixelShift( PP,[145], "min")
Pis = shift_avg.evaluate_shift_for_input_array(UP, 11)
#shift_shifted = px_shift_on_picture_array_rolling.PixelShift(pii, [102], "min")
#piii = shift_shifted.evaluate_shift_for_input_array(piii,111)

plt.figure(100)
plt.plot(-np.log(PP/UP), label = "ODD avg shift on single not on avg")
plt.plot(-np.log(PP/Pis), label = "ODD avg shift on single and on avg")
plt.title("ODD ZnO even odd 419 III")
plt.xlabel("px")
plt.ylabel("i.a.u.")
#plt.xlim(50,2000)
plt.ylim(-0.01, 0.01)
plt.legend()

def create_result_directory(name):
    if os.path.isdir(name):
        pass
    else:
        os.mkdir(name)

create_result_directory("Comparison_avg_stack")
save_pic = os.path.join("Comparison_avg_stack", "ODDZnO_419III_vs"+ ".png")
plt.savefig(save_pic, bbox_inches="tight", dpi=500)
plt.show()


from PyQt5 import QtWidgets, QtCore
from PyQt5.QtCore import Qt
from PyQt5.QtWidgets import QWidget, QGridLayout, QPushButton, QLabel, QDockWidget

class PcSendStringView(QWidget):
    def __init__(self,mainwindow):
        super().__init__()
        self.setWindowFlags(QtCore.Qt.WindowStaysOnTopHint)
        self.mainwindow = mainwindow
        self.setWindowTitle("添加电脑发送字符串事件")
        self.setGeometry(200, 40, 380, 150)
        self.IntUi()

    def IntUi(self):
        mainLayout = QGridLayout()

        ok_button = QPushButton('OK')
        cancel_button = QPushButton('Cancel')

        self.string_edit = QtWidgets.QTextEdit('')

        self.dock_mainLayout = QDockWidget(self)
        self.dock_mainLayout.setWidget(self.string_edit)
        self.dock_mainLayout.setTitleBarWidget(QLabel('电脑发送字符串:'))

        mainLayout.addWidget(self.dock_mainLayout, 0, 0, 0, 2, alignment=Qt.AlignTop)
        mainLayout.addWidget(ok_button, 1, 0)
        mainLayout.addWidget(cancel_button, 1, 1)

        cancel_button.clicked.connect(self.cancel_connect)
        ok_button.clicked.connect(self.ok_connect)

        self.setLayout(mainLayout)

    def cancel_connect(self):
        self.destroy()

    def ok_connect(self):
        text = self.string_edit.toPlainText()
        if text:
            buttonItems = ['电脑发送字符串', text, '0', '1']
            writeTable().select_table(self.mainwindow, buttonItems)
            self.destroy()


import pyzed.sl as sl
import math
import numpy as np
import cv2
import os
import sys
import tifffile as tf
import datetime
import open3d as o3d
import pptk

# Option 1 is for image capture, depth, confidence and point cloud
# Option 2 is for video recording
# Option 3 is for playback

exposure_steps = [0.5, 0.75, 1.25, 1.5]
num_images_to_collect = len(exposure_steps) +1 # first image is auto
option = 3

# /home/emma/evaluation-of-depth-cameras/apple-data
# dirname = 'apple-data/{}'.format(datetime.datetime.now().strftime('%Y%m%d_%H%M%S_%f'))
dirname = '{}'.format(sys.argv[1].split('.')[0])
# dirname = str(timestamp.get_milliseconds())
os.makedirs(dirname, exist_ok=True)
# create a folder inside folder with timestamp inside apple-data
os.makedirs(os.path.join(dirname, 'left-images'), exist_ok=True)
# os.makedirs(os.path.join(dirname, 'right-images'), exist_ok=True)
# os.makedirs(os.path.join(dirname, 'confidence-images'), exist_ok=True)
# os.makedirs(os.path.join(dirname, 'pointcloud-images'), exist_ok=True)
# os.makedirs(os.path.join(dirname, 'depth-images'), exist_ok=True)
# os.makedirs(os.path.join(dirname, 'depth-tiff'), exist_ok=True)

def print_camera_information(cam):
    while True:
        res = input("Do you want to display camera information? [y/n]: ")
        if res == "y":
            print()
            print("Distorsion factor of the right cam before calibration: {0}.".format(
                cam.get_camera_information().calibration_parameters_raw.right_cam.disto))
            print("Distorsion factor of the right cam after calibration: {0}.\n".format(
                cam.get_camera_information().calibration_parameters.right_cam.disto))

            print("Confidence threshold: {0}".format(cam.get_runtime_parameters().confidence_threshold))
            print("Depth min and max range values: {0}, {1}".format(cam.get_init_parameters().depth_minimum_distance,
                                                                    cam.get_init_parameters().depth_maximum_distance)
)
            print("Resolution: {0}, {1}.".format(round(cam.get_camera_information().camera_resolution.width, 2), cam.get_camera_information().camera_resolution.height))
            print("Camera FPS: {0}".format(cam.get_camera_information().camera_fps))
            print("Frame count: {0}.\n".format(cam.get_svo_number_of_frames()))
            break
        elif res == "n":
            print("Camera information not displayed.\n")
            break
        else:
            print("Error, please enter [y/n].\n")

def saving_image(key, mat, i):
    img = sl.ERROR_CODE.FAILURE
    while img != sl.ERROR_CODE.SUCCESS:
        filepath = os.path.join(dirname, 'frame_{:02d}.jpg'.format(i))
        img = mat.write(filepath)
        print("Saving image : {0}".format(repr(img)))
        if img == sl.ERROR_CODE.SUCCESS:
            break
        else:
            print("Help: you must enter the filepath + filename + PNG extension.")

def main():
    # Create a Camera object
    zed = sl.Camera()

    # Create a InitParameters object and set configuration parameters
    init_params = sl.InitParameters()

    ######## Camera
    init_params.camera_resolution = sl.RESOLUTION.HD720
    #init_params.camera_resolution = sl.RESOLUTION.HD2K
    #init_params.camera_resolution = sl.RESOLUTION.HD1080
    #init_params.camera_resolution = sl.RESOLUTION.VGA
    init_params.camera_fps = 30  # Set fps at 30 - other frame rates can be 15, 30, 60, 100
    # Set exposure to a certain % of camera framerate
    #zed.set_camera_settings(sl.VIDEO_SETTINGS.EXPOSURE, 50)
    # Set white balance to a certain K
    #zed.set_camera_settings(sl.VIDEO_SETTINGS.WHITE_BALANCE, 4600)
    # Reset to auto exposure
    #zed.set_camera_settings(sl.VIDEO_SETTINGS.EXPOSURE, -1)

    ########  Depth
    init_params.depth_mode = sl.DEPTH_MODE.PERFORMANCE  # Use PERFORMANCE depth mode (not necessary as this is default), there are three different modes PERFORMANCE, MEDIUM, QUALITY
    #init_params.depth_mode = sl.DEPTH_MODE.QUALITY
    #init_params.depth_mode = sl.DEPTH_MODE.MEDIUM
    init_params.coordinate_units = sl.UNIT.METER
    init_params.depth_minimum_distance = 1.0
    init_params.depth_maximum_distance = 10.0

    # Open the camera
    if (option == 1 or option == 2):
        status = zed.open(init_params)
        if status != sl.ERROR_CODE.SUCCESS:
            print(repr(status))
            exit(1)

    if (option == 1):

        # Create and set RuntimeParameters after opening the camera
        runtime_parameters = sl.RuntimeParameters()
        runtime_parameters.sensing_mode = sl.SENSING_MODE.STANDARD  # Use STANDARD sensing mode
        # Setting the depth confidence parameters
        runtime_parameters.confidence_threshold = 100
        runtime_parameters.textureness_confidence_threshold = 100

        # Capture 150 images and depth, then stop
        i = 0
        image_left = sl.Mat() # stores the images
        image_right = sl.Mat() # stores the images
        depth = sl.Mat() # stores the depth map
        point_cloud = sl.Mat()
        confidence_map = sl.Mat()

        mirror_ref = sl.Transform()
        mirror_ref.set_translation(sl.Translation(2.75,4.0,0))
        tr_np = mirror_ref.m

        while i < num_images_to_collect:
            # A new image is available if grab() returns SUCCESS. This transfers GPU -> CPU by default here
            if i ==0:
                exposure = zed.get_camera_settings(sl.VIDEO_SETTINGS.EXPOSURE)
            else:
                newExposure = exposure*exposure_steps[i-1]
                zed.set_camera_settings(sl.VIDEO_SETTINGS.EXPOSURE, newExposure)

            if zed.grab(runtime_parameters) == sl.ERROR_CODE.SUCCESS:
                # Retrieve left image
                zed.retrieve_image(image_left, sl.VIEW.LEFT)
                # Retrieve right image
                zed.retrieve_image(image_right, sl.VIEW.RIGHT)
                # Retrieve depth map. Depth is aligned on the overlapped image
                zed.retrieve_measure(depth, sl.MEASURE.DEPTH)
                # Retrieve colored point cloud. Point cloud is aligned on the overlapped image.
                zed.retrieve_measure(point_cloud, sl.MEASURE.XYZRGBA)
                # Retrieve confidence map
                zed.retrieve_measure(confidence_map, sl.MEASURE.CONFIDENCE)

                timestamp = zed.get_timestamp(sl.TIME_REFERENCE.CURRENT)  # Get the timestamp at the time the image was captured
                print("Image resolution: {0} x {1} || Image timestamp: {2}".format(image_left.get_width(), image_left.get_height(),
                    timestamp.get_milliseconds()))
                print("Image resolution: {0} x {1} || Image timestamp: {2}".format(image_right.get_width(), image_right.get_height(),
                    timestamp.get_milliseconds()))

                # depth_for_display = sl.Mat()
                # zed.retrieve_image(depth_for_display, sl.VIEW.DEPTH)
                # depth_for_display.write('depth_view.jpg')

                depth_data = depth.get_data()
                depth_data[~ np.isfinite(depth_data)] = 0
                depth_img_name = os.path.join(dirname, 'depth_{:02d}.jpg'.format(i))
                cv2.imwrite(depth_img_name, (depth_data/depth_data.max())*255)
                tf.imwrite(depth_img_name.replace('.jpg','.tiff').replace('depth','depth1'), depth_data)

                confidence_data = confidence_map.get_data()
                confidence_img_name = os.path.join(dirname, 'confidence_{:02d}.jpg'.format(i))
                scale = 1./np.max(confidence_data)*255
                cv2.imwrite(confidence_img_name, confidence_data*scale)

                pointcloud_data = point_cloud.get_data()
                pointcloud_name = os.path.join(dirname,'pointcloud_{:02d}.ply'.format(i))
                if not point_cloud.write(pointcloud_name) == sl.ERROR_CODE.SUCCESS:
                    print('point cloud data not saved')

                img_left = image_left.get_data()
                img_name = os.path.join(dirname, 'left_{:02d}.jpg'.format(i))
                cv2.imwrite(img_name, img_left)

                img_right = image_right.get_data()
                img_name = os.path.join(dirname, 'right_{:02d}.jpg'.format(i))
                cv2.imwrite(img_name, img_right)

                np.savez(os.path.join(dirname, 'raw_{:02d}.npz'.format(i)), depth=depth_data, right=img_right, left=img_left, confidence=confidence_data)

                camera_exposure = zed.get_camera_settings(sl.VIDEO_SETTINGS.EXPOSURE)
                camera_whitebalance = zed.get_camera_settings(sl.VIDEO_SETTINGS.WHITEBALANCE_TEMPERATURE)
                print("Exposure time: {:03d}      Whitebalance Temperature: {}".format(camera_exposure, camera_whitebalance))
                i = i + 1

                # for arr in [depth,image_left, image_right, point_cloud, confidence_map ]:
                #     print(arr.get_memory_type())

    elif(option == 2):
        # cd to folder that this script is in
        # type python "datacollection.py" {filename.svo} in terminal
        # go Ctrl-C to stop recording

        if not sys.argv or len(sys.argv) != 2:
            print("Only the path of the output SVO file should be passed as argument.")
            exit(1)

        init_params.depth_mode = sl.DEPTH_MODE.NONE

        path_output = sys.argv[1]
        recording_param = sl.RecordingParameters(path_output, sl.SVO_COMPRESSION_MODE.H264)
        err = zed.enable_recording(recording_param)
        if err != sl.ERROR_CODE.SUCCESS:
            print(repr(status))
            exit(1)

        runtime = sl.RuntimeParameters()
        print("SVO is Recording, use Ctrl-C to stop.")
        frames_recorded = 0

        while True:
            if zed.grab(runtime) == sl.ERROR_CODE.SUCCESS :
                frames_recorded += 1
                print("Frame count: " + str(frames_recorded), end="\r")

    elif(option == 3):
        if len(sys.argv) != 2:
            print("Please specify path to .svo file.")
            exit()

        filepath = sys.argv[1]
        print("Reading SVO file: {0}".format(filepath))

        input_type = sl.InputType()
        input_type.set_from_svo_file(filepath)
        init = sl.InitParameters(input_t=input_type, svo_real_time_mode=False)
        status = zed.open(init)
        if status != sl.ERROR_CODE.SUCCESS:
            print(repr(status))
            exit()

        runtime = sl.RuntimeParameters()
        image_left = sl.Mat() # stores the images
        image_right = sl.Mat() # stores the images
        depth = sl.Mat() # stores the depth map
        point_cloud = sl.Mat()
        confidence_map = sl.Mat()
        mat = sl.Mat()

        key = ''
        print("  Quit the video reading:     q\n")
        img_counter = 0
        while key != 113:  # for 'q' key
            err = zed.grab(runtime)
            if err == sl.ERROR_CODE.SUCCESS:
                if err == sl.ERROR_CODE.SUCCESS:
                    zed.retrieve_image(mat)
                    cv2.imshow("ZED", mat.get_data())
                    key = cv2.waitKey(1)
                else:
                    key = cv2.waitKey(1)
                # Retrieve left image
                zed.retrieve_image(image_left, sl.VIEW.LEFT)
                # Retrieve right image

                # zed.retrieve_image(image_right, sl.VIEW.RIGHT)
                # # Retrieve depth map. Depth is aligned on the overlapped image
                # zed.retrieve_measure(depth, sl.MEASURE.DEPTH)
                # # Retrieve colored point cloud. Point cloud is aligned on the overlapped image.
                #pcd = o3d.geometry.PointCloud()
                # zed.retrieve_measure(point_cloud, sl.MEASURE.XYZRGBA)
                # pc = point_cloud.get_data()
                #pcd.points = o3d.utility.Vector3dVector(pc[500:600, 500:600, :3].reshape(-1, 3))
                # pcd = pc[:, :, :3].reshape(-1, 3)
                # v = pptk.viewer(pcd)
                # v.attributes(pcd)
                # v.set(point_size=0.01)

                # # Retrieve confidence map
                # zed.retrieve_measure(confidence_map, sl.MEASURE.CONFIDENCE)

                timestamp = zed.get_timestamp(sl.TIME_REFERENCE.CURRENT)  # Get the timestamp at the time the image was captured
                print("Image resolution: {0} x {1} || Image timestamp: {2}".format(image_left.get_width(), image_left.get_height(),
                timestamp.get_milliseconds()))
                #print(len(pc))
                # print("Image resolution: {0} x {1} || Image timestamp: {2}".format(image_right.get_width(), image_right.get_height(),
                #     timestamp.get_milliseconds()))

                # depth_for_display = sl.Mat()
                # zed.retrieve_image(depth_for_display, sl.VIEW.DEPTH)
                # depth_for_display.write('depth_view.jpg')

                # depth_data = depth.get_data()
                # depth_data[~ np.isfinite(depth_data)] = 0
                # depth_img_name = os.path.join(dirname, 'depth-images', 'depth_{:02d}.jpg'.format(img_counter))
                # cv2.imwrite(depth_img_name, (depth_data/depth_data.max())*255)
                # depth_tiff_name = os.path.join(dirname, 'depth-tiff', 'depth_{:02d}.tiff'.format(img_counter))
                # tf.imwrite(depth_tiff_name, depth_data)

                # confidence_data = confidence_map.get_data()
                # confidence_img_name = os.path.join(dirname, 'confidence-images','confidence_{:02d}.jpg'.format(img_counter))
                # scale = 1./np.max(confidence_data)*255
                # cv2.imwrite(confidence_img_name, confidence_data*scale)

                # pointcloud_data = point_cloud.get_data()
                # pointcloud_name = os.path.join(dirname, 'pointcloud-images', 'pointcloud_{:02d}.ply'.format(img_counter))
                # if not point_cloud.write(pointcloud_name) == sl.ERROR_CODE.SUCCESS:
                #     print('point cloud data not saved')

                #saving point-cloud
                # pcd = o3d.geometry.PointCloud()
                # output_array = pointcloud_data[:3, :, :]
                # # reshape numpy array from 4D to 3D
                # # Gives a new shape to an array without changing its data
                # pointcloud_data = pointcloud_data[~np.isnan(pointcloud_data)]
                # output_array = output_array.reshape(-1, 3)
                # #output_array = pointcloud_data[:3, :, :].reshape(-1, 3)
                # #output_array = pointcloud_data.reshape(-1, 3)
                # #convert numpy array to open3d
                # pcd.points = o3d.utility.Vector3dVector(output_array)
                # o3d.io.write_point_cloud(pointcloud_name, pcd)
                # o3d.visualization.draw_geometries([pcd])

                img_left = image_left.get_data()
                img_name = os.path.join(dirname, 'left-images', 'left_{:02d}.jpg'.format(img_counter))
                cv2.imwrite(img_name, img_left)

                # img_right = image_right.get_data()
                # img_name = os.path.join(dirname, 'right-images', 'right_{:02d}.jpg'.format(img_counter))
                # cv2.imwrite(img_name, img_right)

                img_counter += 1

            else:
                key = cv2.waitKey(1)
        cv2.destroyAllWindows()

        print_camera_information(zed)
        zed.close()
        print("\nFINISH")

    # Close the camera
    if(option == "1"):
        zed.close()

if __name__ == "__main__":
    main()


from django.conf import settings
import requests
import os
from shortid import short_id
import re
import json

BASE_URL = 'https://graph.qq.com/oauth2.0/authorize'

def _get(url, params):
    r = requests.get(BASE_URL + url, params=params, headers={"Accept": 'application/json'})
    if (r.status_code == 200):
        return r.json()
    else:
        raise RuntimeError

# 获取认证url
def get_auth_url(request):
    domain = "https://" + request.META.get("HTTP_HOST")

    # 防止CSRF攻击
    state = short_id.get_short_id()

    client_id = __get_config('QQ_CLIENT_ID')
    redirect_uri = domain + __get_config('QQ_CLIENT_CALLBACK')
    request.session['state'] = state
    return BASE_URL + "?response_type=code&client_id={}&redirect_uri={}&state={}".format(client_id, redirect_uri, state)

def get_access_token(request, code):
    domain = "https://" + request.META.get("HTTP_HOST")

    url = 'https://graph.qq.com/oauth2.0/token?grant_type=authorization_code&client_id={}&client_secret={}&code={}&redirect_uri={}'.format(
        __get_config('QQ_CLIENT_ID'),
        __get_config('QQ_CLIENT_SECRET'),
        code,
        domain + __get_config('QQ_CLIENT_CALLBACK'))
    r = requests.get(url, headers={"Accept": 'application/json'})
    if (r.status_code == 200):
        str = r.text
        array = str.split("&")
        params = {}
        for item in array:
            items = item.split("=")
            params[items[0]] = items[1]
        return params
    else:
        raise RuntimeError(r.text)

# 获取用户信息
def get_user(access_token):
    rs = requests.get('https://graph.qq.com/oauth2.0/me?access_token={}'.format(access_token),
                      headers={"Accept": 'application/json'})
    if rs.status_code == 200:
        jsonp = rs.text
        p = re.match(r'callback\((.*?)\)', jsonp)
        if p:
            me = json.loads(p.group(1))
            openid = me.get('openid')
            params = {
                'access_token': access_token,
                'oauth_consumer_key': __get_config('QQ_CLIENT_ID'),
                'openid': openid
            }
            rs = requests.get('https://graph.qq.com/user/get_user_info', params=params,
                              headers={"Accept": 'application/json'})
            if rs.status_code == 200:
                user = rs.json()
                member = {}

                member['name'] = user.get('nickname')
                member['node_id'] = openid
                member['avatar_url'] = user.get('figureurl_qq_2')
                member['type'] = 1

                return member;
    pass

def __get_config(name):
    value = os.environ.get(name, getattr(settings, name, None))
    return value


from collections import defaultdict
from datetime import datetime

file=open("input.in","r")
content = file.readlines()

guardsShift = defaultdict(list)
guardTimes = defaultdict(int)

for line in sorted(content):
    time, action = line.split('] ')
    time = datetime.strptime(time[1:], '%Y-%m-%d %H:%M')

    if action.startswith('Guard'):
        idGuard= int(action.split()[1][1:])
    elif 'falls asleep' in action:
        start_time = time
    elif 'wakes up' in action:
        end_time = time
        guardsShift[idGuard].append((start_time.minute,end_time.minute))
        guardTimes[idGuard] += (end_time.minute - start_time.minute)

(idSleepyGuard, highestTime) = max(guardTimes.items(), key=lambda par: par[1])

countTimes = [0 for x in range(60)]
for sleep in guardsShift[idSleepyGuard]:
    for i in range(sleep[0],sleep[1] + 1):
        countTimes[i] += 1

bestMinute = countTimes.index(max(countTimes))

print(idSleepyGuard)
print(bestMinute)
print(idSleepyGuard * bestMinute)


import os, os.path
import pickle
import pytest
from sqlalchemy import inspect
import vedana
import pkdiagram
from pkdiagram import util, Scene, ServerFileManagerModel, Person

from fdserver.extensions import db
from fdserver.models import Diagram

@pytest.fixture
def create_model(request):

    created = []

    def _create_model(session=True):
            model = ServerFileManagerModel()
            _session = pkdiagram.Session()
            if session:
                if session is True:
                    test_session = request.getfixturevalue('test_session')
                else:
                    test_session = session
                _session.init(sessionData=test_session.account_editor_dict(), syncWithServer=False)
            else:
                _session.init(syncWithServer=False)

            model.init()
            model.setSession(_session)
            assert util.wait(model.updateFinished)

            created.append(model)
            return model

    yield _create_model

    for model in created:
        if model.session:
            model.session.deinit()
        model.deinit()

def _grant_ro_access(diagrams, grantee):
    for diagram in diagrams:
        if diagram.user_id != grantee.id:
            diagram.grant_access(grantee, vedana.ACCESS_READ_ONLY)
    db.session.commit()

def test_get_index(test_user_diagrams, test_user, create_model):
    model = create_model()

    db.session.add_all((test_user, *test_user_diagrams))
    _grant_ro_access(test_user_diagrams, test_user)
    diagrams = Diagram.query.all()
    model.update()
    assert util.Condition(model.updateFinished).wait()
    assert model.rowCount() == len(test_user_diagrams) + 1
    # assert model.data(model.index(0, 0), model.IDRole) == diagrams[0].id
    # assert model.data(model.index(1, 0), model.IDRole) == diagrams[1].id
    # assert model.data(model.index(2, 0), model.IDRole) == diagrams[2].id
    # assert model.data(model.index(3, 0), model.IDRole) == diagrams[3].id
    # assert model.data(model.index(4, 0), model.IDRole) == diagrams[4].id
    # assert model.data(model.index(5, 0), model.IDRole) == diagrams[5].id
    # assert model.data(model.index(6, 0), model.IDRole) == diagrams[6].id
    # assert model.data(model.index(7, 0), model.IDRole) == diagrams[7].id
    # assert model.data(model.index(8, 0), model.IDRole) == diagrams[8].id
    # assert model.data(model.index(9, 0), model.IDRole) == diagrams[9].id
    # assert model.data(model.index(0, 0), model.IDRole) == test_user.free_diagram.id
    # assert model.data(model.index(0, 0), model.OwnerRole) == diagrams[0].user.username
    # assert model.data(model.index(1, 0), model.OwnerRole) == diagrams[1].user.username
    # assert model.data(model.index(2, 0), model.OwnerRole) == diagrams[2].user.username
    # assert model.data(model.index(3, 0), model.OwnerRole) == diagrams[3].user.username
    # assert model.data(model.index(4, 0), model.OwnerRole) == diagrams[4].user.username
    # assert model.data(model.index(5, 0), model.OwnerRole) == diagrams[5].user.username
    # assert model.data(model.index(6, 0), model.OwnerRole) == diagrams[6].user.username
    # assert model.data(model.index(7, 0), model.OwnerRole) == diagrams[7].user.username
    # assert model.data(model.index(8, 0), model.OwnerRole) == diagrams[8].user.username
    # assert model.data(model.index(9, 0), model.OwnerRole) == diagrams[9].user.username
    # assert model.data(model.index(10, 0), model.OwnerRole) == test_user.free_diagram.user.username

def test_disk_cache(test_user, test_user_diagrams, create_model):
    model = create_model()

    db.session.add_all((test_user, *test_user_diagrams))
    assert model.rowCount() == len([x for x in test_user_diagrams if x.user_id == test_user.id]) + 1
    model.write()

    model2 = ServerFileManagerModel(dataPath=model.dataPath)
    model2.init()
    model2.setSession(model.session)
    updateFinished = util.Condition(model2.updateFinished)
    indexGETResponse2 = util.Condition(model2.indexGETResponse)
    diagramGETResponse2 = util.Condition(model2.diagramGETResponse)
    assert updateFinished.wait() == True
    assert indexGETResponse2.callCount == 1
    assert diagramGETResponse2.callCount == 0
    assert model2.rowCount() == model.rowCount()
    # assert encryption.isEncrypted(model.index(0, 0).data(role=model.DiagramDataRole)) == False
    # assert encryption.isEncrypted(model.index(1, 0).data(role=model.DiagramDataRole)) == False
    # assert encryption.isEncrypted(model.index(2, 0).data(role=model.DiagramDataRole)) == False
    # assert encryption.isEncrypted(model.index(3, 0).data(role=model.DiagramDataRole)) == False
    # assert encryption.isEncrypted(model.index(4, 0).data(role=model.DiagramDataRole)) == False

# def test_cached_fds_are_encrypted(test_user_diagrams, create_model):
#     model = create_model()
#     assert model.rowCount() == 6
#     def _is_fd_encrypted(row):
#         diagram_id = model.index(row, 0).data(model.IDRole)
#         fpath = model.localPathForID(diagram_id)
#         with open(os.path.join(fpath, 'diagram.pickle'), 'rb') as f:
#             ebdata = f.read()
#             return encryption.isEncrypted(ebdata)
#     assert _is_fd_encrypted(0)
#     assert _is_fd_encrypted(1)
#     assert _is_fd_encrypted(2)
#     assert _is_fd_encrypted(3)
#     assert _is_fd_encrypted(4)

def test_dataChanged_on_poll(test_user, create_model):
    model = create_model()
    dataChanged = util.Condition(model.dataChanged)
    model.syncDiagramFromServer(test_user.free_diagram_id)
    assert dataChanged.callCount == 0
    assert model.rowCount() == 1

    # Simulate save on another machine
    diagram = Diagram.query.get(test_user.free_diagram_id)
    data = {}
    scene = Scene()
    scene.read(pickle.loads(diagram.data))
    scene.addItems(Person(name="Patrick"))
    scene.write(data)
    diagram.update(data=pickle.dumps(data), _commit=True)
    inspect(diagram).session.add(diagram)
    inspect(diagram).session.commit()

    # Polled update on this machine
    model.syncDiagramFromServer(test_user.free_diagram_id)
    assert dataChanged.callCount == 1
    assert dataChanged.callArgs[0] == (model.index(0, 0), model.index(0, 0), [model.DiagramDataRole])
    scene = Scene()
    bdata = model.data(dataChanged.callArgs[0][0], model.DiagramDataRole)
    scene.read(pickle.loads(bdata))
    assert scene.query1(name='Patrick')

    # Test sync idempotence
    model.syncDiagramFromServer(test_user.free_diagram_id)
    assert dataChanged.callCount == 1

def test_add_diagram_file(test_user, create_model):
    data = {}
    Scene().write(data)
    diagram = Diagram(
        user_id=test_user.id,
        data=pickle.dumps(data)
    )
    db.session.add(diagram)
    db.session.commit()

    model = create_model()
    db.session.add(diagram)
    _diagram = pkdiagram.Diagram.create(diagram.as_dict())
    model._addOrUpdateDiagram(_diagram)
    fpath = model.localPathForID(diagram.id)
    assert os.path.isdir(fpath) == True

def test_helpers(test_user_diagrams, create_model):
    ROW = 1

    model = create_model()
    diagram_1 = model.diagramForRow(ROW)
    diagram_id = model.index(ROW, 0).data(model.IDRole)
    assert diagram_id == diagram_1.id

    fpath = model.localPathForID(diagram_id)
    diagram = model.serverDiagramForPath(fpath)
    assert model.rowForDiagramId(diagram.id) == ROW

def test_searchText(test_user_diagrams, test_user, create_model):
    SEARCH_TERM = 'patrickkidd+unittest+2@gmail.com'

    db.session.add_all((test_user, *test_user_diagrams))
    num_matching_search = len([x for x in test_user_diagrams if SEARCH_TERM in x.user.username])
    _grant_ro_access(test_user_diagrams, test_user)

    model = create_model()
    assert model.rowCount() == len(test_user_diagrams) + 1

    model.searchText = SEARCH_TERM
    assert model.rowCount() == num_matching_search

    model.reset('searchText')
    assert model.rowCount() == len(test_user_diagrams) + 1

def test_clear_on_logout(test_user_diagrams, test_user, create_model):
    model = create_model()
    db.session.add_all(test_user_diagrams)
    assert model.rowCount() == len([x for x in test_user_diagrams if x.check_read_access(test_user)]) + 1

    cleared = util.Condition(model.cleared)
    model.session.logout()
    assert cleared.wait() == True
    assert model.rowCount() == 0

def test_dont_clear_cache_on_restart(test_user_diagrams, test_user, test_session, create_model):
    model = create_model()
    db.session.add_all(test_user_diagrams)
    assert model.rowCount() == len([x for x in test_user_diagrams if x.check_read_access(test_user)]) + 1

    model.deinit()
    session = pkdiagram.Session()
    session.init(sessionData=test_session.account_editor_dict(), syncWithServer=False)
    model2 = ServerFileManagerModel()
    model2.init()
    model2.setSession(session)
    assert util.wait(model2.updateFinished)
    assert model2.rowCount() == len([x for x in test_user_diagrams if x.check_read_access(test_user)]) + 1

def test_save_free_diagram_persists(test_session):
    # Read free diagram 1
    model = ServerFileManagerModel()
    session = pkdiagram.Session()
    session.init(sessionData=test_session.account_editor_dict(), syncWithServer=False)
    assert session.hasFeature(vedana.LICENSE_FREE) == True
    model.init()
    model.setSession(session)
    assert util.wait(model.updateFinished)
    assert model.rowCount() == 1

    # Write free diagram + logout
    scene = Scene(items=Person(name='Me'))
    bdata = pickle.dumps(scene.data())
    model.setData(model.index(0, 0), bdata, role=model.DiagramDataRole)
    session.logout()
    assert model.rowCount() == 0

    # Login + Read free diagram
    session.setData(test_session.account_editor_dict())
    assert util.wait(model.updateFinished)
    assert model.rowCount() == 1
    assert model.index(0, 0).data(model.DiagramDataRole) == bdata

def test_delete_file(qtbot, create_model):
    model = create_model()
    # session = Session()
    # session.init(sessionData=test_session.account_editor_dict(), syncWithServer=False)

    qtbot.clickYesAfter(lambda: model.deleteFileAtRow(0))
    assert model.rowCount() == 0


import logging
logging.basicConfig(level=logging.INFO,format='{asctime} {levelname} ({threadName:11s}) {message}', style='{')

""" /// FUNCIONES VARIAS /// """

#funcion para descargar info diaria. Para intra se utilizará alphavange
def getInfoDiaria(ticker, start="2000-01-01",interval="1d",end=None):
    """Recibe un ticjer y devuelve la info diaria historica del papel"""
    logging.info(f"importando info de {ticker} ")
    import yfinance as yf
    df=yf.download(ticker,start=start,end=end,interval=interval,auto_adjust=True)
    return df

def getDataExcel(ticker):
    '''
       SOLO SIRVE PARA EXCEL DESCARGADO DEYAHOO FINANCE
       Busca excel de datos y devuelve DF las columnas: 'Open','High','Low','Close','Volume'.
       La columna 'Close' trae el precio ajustado.
       |
       |_ ticker:  El ticker a buscar
       '''
    import pandas as pd
    try:
        data = pd.read_excel(ticker + '.xlsx').set_index('Date').sort_index()
        #data.index.names=["Date"]
        data.columns = ['Open', 'High', 'Low','Close', 'Volume']
    except:
        try:
            data = pd.read_excel('excels_csvs/ADRs/' + ticker + '.xlsx').set_index('timestamp').sort_index()
            data.columns = ['Open', 'High', 'Low', 'Close', 'AdjClose', 'Volume']
            data['pctChange'] = data.AdjClose.pct_change()
        except:
            data = 'Sorry man no encontre el archivo en tu directorio'
    #data = data.drop(["Borrar"],axis=1)
    return data

def cruceSMA(df,fast=10,slow=20):
    """Recibe un df y una cantidad de ruedas y devuelve una columna con el SMA_ruedas"""
    logging.info(f"Creando cruce SMA, fast: {fast} slow: {fast} ")
    df["cruce"] = round((df.Close.rolling(fast).mean() / df.Close.rolling(slow).mean() -1) *100,2)
    return df

def cruceEMA(df,fast=10,slow=20): #ewm(span=ruedas)
    """Recibe un df y una cantidad de ruedas y devuelve una columna con el SMA_ruedas"""
    logging.info(f"Creando cruce EMA, fast: {fast} slow: {fast} ")
    df["cruce"] = round((df.Close.ewm(span=fast).mean() / df.Close.ewm(span=slow).mean() - 1) * 100, 2)
    return df

def getRSI(df,ruedas=14):
    """Recibe un df y una cantidad de ruedas (por defecto 14) y devuelve una columna con el RSI"""
    logging.info(f"Creando RSI, ruedas: {ruedas}")
    import numpy as np

    df["diferencia"]=df.Close.diff()
    df["dif_pos"] = np.where(df["diferencia"] > 0,df["diferencia"],0)
    df["dif_neg"] = np.where(df["diferencia"] < 0, df["diferencia"], 0)
    df["media_pos"] = df["dif_pos"].ewm(alpha=1/ruedas).mean()
    df["media_neg"] = df["dif_neg"].ewm(alpha=1/ruedas).mean()
    df["rs"] = df["media_pos"] / abs(df["media_neg"])
    df["RSI"] = round(100 - 100 / (1 + df["rs"]),2)
    df=df.drop(["diferencia","dif_pos","dif_neg","media_pos","media_neg","rs"],axis=1)
    return df

def getMACD(df,fast=12,slow=26,signal=9):
    """Recibe un df y un numero para fast (defecto 12), slow (defecto 26) u signal (defecto 9) devuelve una columna con el MACD"""
    logging.info(f"Creando MACD")
    df["ema_fast"] = df.Close.ewm(span=fast).mean()
    df["ema_slow"] = df.Close.ewm(span=slow).mean()
    df["macd"] = df["ema_fast"] - df["ema_slow"]
    df["signal"] = df["macd"].ewm(span=signal).mean()
    df["histograma"] = df["macd"] - df["signal"]
    df = df.drop(["ema_fast", "ema_slow"], axis=1)
    return df

""" /// FUNCIONES DE BACKTEST /// """

def getAlertras(df,media=True,mediaTipo="SMA",fast=5,slow=20,buy_cr=0,sell_cr=0,
                rsi=True,rsi_q=14,buy_rsi=50,sell_rsi=35,
                stopActivado=False,porcentaje=1.0,
                tipo="both"):
    """Recibe un df y los parametros y agrega los indicadores y la señal, por ahora solo SMA, EMA, RSI y STOP LOSS
    Tambien agrega volatilidad y variacion"""
    logging.info(f"Creando Alertas")

    df["Variacion"] = data.Close.pct_change()*100
    df["Volatilidad"] = data.Variacion.rolling(250).std() * 250**0.5
    #Medias
    compraMedia=True
    ventaMedia=True
    if media:
        if mediaTipo=="SMA":
            df=cruceSMA(df,fast,slow).copy()
        else:
            df=cruceEMA(df,fast,slow).copy()
        compraMedia= df.cruce > buy_cr
        ventaMedia= df.cruce < sell_cr

    #RSI
    compraRSI=True
    ventaRSI=True
    if rsi:
        df=getRSI(df,rsi_q).copy()
        compraRSI= df.RSI > buy_rsi
        ventaRSI= df.RSI < sell_rsi

    df= getSenales(df,compraMedia,ventaMedia,compraRSI,ventaRSI,stopActivado=stopActivado,porcentaje=porcentaje,tipo=tipo)
    df.set_index("Date",inplace=True)

    return df

def getSenales(df,compraMedia,ventaMedia,compraRSI,ventaRSI,stopActivado,porcentaje,tipo):
    """ ACÁ VIENE LA LOGICA CON EL STOP LOSS
    Funcion que recibe el df y los parametros y agrega señales al df
    Primero verifico la señal de stop (verifica que esté comprado o vendido antes)
    y luego si no da STOP, pasa a la otra logica que seria si da compra o venta.
    """
    logging.info(f"Añadiendo las señales")
    copia = df.copy()
    copia = copia.reset_index()

    #ACA VIENE LA LOGICA DEL STOP LOSS
    #variable stop es un diccionario que tiene el porcentaje, el precio y el estado actual
    stop={"porcentaje":porcentaje,"pStop":False,"estado":"sin señal"}
    listaSenal=[] #esta lista la hago para luego meterlo en el df

    for i in range(len(copia)):
        precioActual = copia.loc[i]["Close"]

        if stopActivado: #verifica si se toma o no el stop loss

            if tipo != "short" and stop["estado"] == "compra" and precioActual < stop["pStop"]: #stop si es long o both y está comprado
                listaSenal.append("stop")
                stop["pStop"] = False
                stop["estado"] = "stop"
                continue
            elif tipo != "long" and stop["estado"] == "venta" and precioActual > stop["pStop"]: #stop si es short y está vendido
                listaSenal.append("stop")
                stop["pStop"] = False
                stop["estado"] = "stop"
                continue

        if compraRSI[i] and compraMedia[i]: #señal compra
            if stop["estado"] == "compra": #si ya estoy comprado no pone señal
                listaSenal.append("sin señal")

            elif tipo == "short" and stop["estado"] != "venta": #si voy short y no estoy vendido no puede comprar
                listaSenal.append("sin señal")

            else: #si estoy long/both y no estoy comprado o si estoy short y esta vendido
                listaSenal.append("compra")
                stop["pStop"] = precioActual * stop["porcentaje"] #pone el nuevo stop
                stop["estado"]="compra"

        elif ventaMedia[i] and ventaRSI[i]: #señal venta
            if stop["estado"] == "venta": #si ya estoy vendido no pone señal
                listaSenal.append("sin señal")
            elif tipo == "long" and stop["estado"] != "compra": #si voy long y no estoy comprado no puedo vender
                listaSenal.append("sin señal")
            else: #si voy short/both y no estoy vendido, o si voy long y estoy comprado
                listaSenal.append("venta")
                stop["pStop"] = precioActual * (1 + (1- stop["porcentaje"]))
                stop["estado"] = "venta"

        else: #no es ni stop, ni compra ni venta
            listaSenal.append("sin señal")

    copia["Señal"] = listaSenal

    return copia

def getActions(df,tipo="both"):
    """Recibe el df, saca los sin señal y setea si es compra o venta para los primeros y ultimos trades"""
    logging.info(f"Seteando el data frame ")
    actions = df.loc [ df.Señal != "sin señal"] . copy()
    actions = actions.reset_index()

    primera = actions.iloc[0]["Señal"]
    ultima = actions.iloc[-1]["Señal"]

    #Elimina aperturas y cierres que no tienen sentido
    if tipo == "long":
        if primera == "venta":
            actions = actions.iloc[1:]
        if ultima == "compra":
            actions = actions.iloc[:-1]

    if tipo == "short":
        if primera == "compra":
            actions = actions.iloc[1:]
        if ultima == "venta":
            actions = actions.iloc[:-1]

    actions = actions.set_index("Date")
    return actions

def getTrades(df,tipo="both"): #estoy n esto

    """toma las acciones y me arma un libro con cada trade"""
    logging.info(f"Creando libro de trades")
    trades={}
    f_ini=[]
    f_fin=[]
    p_ini=[]
    p_fin=[]
    listaSide=[]
    df=df.reset_index()
    abierto = False

    for i in range(len(df)):

        if tipo != "both": #si es diferente de both, funciona la regla de que los pares abren operacion
            # tomo los pares para empezar por el primero ya que viene seteado desde antes

            if i % 2 == 0 and i != (len(df) - 1):
                f_ini.append(df.iloc[i]["Date"])
                p_ini.append(df.iloc[i]["Close"])
                listaSide.append(tipo) #cuando abre tambien pone el side
            elif i % 2 != 0:
                f_fin.append(df.iloc[i]["Date"])
                p_fin.append(df.iloc[i]["Close"])

        else:
            senal = df.iloc[i]["Señal"] #le agrego el side
            if senal == "compra":
                direccion = "long"
            elif senal == "venta":
                direccion = "short"

            if i == 0:
                #abro el primer trade
                f_ini.append(df.iloc[i]["Date"])
                p_ini.append(df.iloc[i]["Close"])
                abierto = True
                listaSide.append(direccion)

            elif i == (len(df) - 1): #si es el ultimo, lo cierro
                f_fin.append(df.iloc[i]["Date"])
                p_fin.append(df.iloc[i]["Close"])
                break

            elif senal == "stop":
                f_fin.append(df.iloc[i]["Date"])
                p_fin.append(df.iloc[i]["Close"])
                abierto = False

            elif i>0 and senal != "stop":
                if abierto == True: #si viene un trade abierto, lo cierro y abro otro
                    f_fin.append(df.iloc[i]["Date"])
                    p_fin.append(df.iloc[i]["Close"])

                    f_ini.append(df.iloc[i]["Date"])
                    p_ini.append(df.iloc[i]["Close"])
                    listaSide.append(direccion)

                else: #si vengo de un trade cerrado, solo lo abro
                    f_ini.append(df.iloc[i]["Date"])
                    p_ini.append(df.iloc[i]["Close"])
                    abierto=True
                    listaSide.append(direccion)

    trades["Fecha_ini"]=f_ini
    trades["Precio_ini"]=p_ini
    trades["Side"] = listaSide
    trades["Fecha_fin"]=f_fin
    trades["Precio_fin"]=p_fin

    trades=pd.DataFrame(trades) #creo el data frame con el diccionario

    df = df.set_index("Date") #pongo el indice como estaba
    trades = getResult(trades)

    return trades

def getResult(df):
    """Recibe un df con acciones de compra venta y stop y le agrega el Yield y la cantidad de Dias"""
    logging.info(f"Calculando Yield y Dias")
    df["Yield"] = np.where(df["Side"] == "long",(df["Precio_fin"] / df["Precio_ini"] -1)*100,(df["Precio_ini"] / df["Precio_fin"] -1)*100)
    df["Dias"] = (df.Fecha_fin - df.Fecha_ini).dt.days

    return df

def signalMetrics(df,tipo="both"):
    """Recibe un df y un side con señales y entrega variables estadisticas. Toma tiempos comprados, vendidos
    y sin posicion que hubieran operado según la estrategia.
    -Volatilidad media
    -Variacion media
    -Cantidad total de en cada estado
    -Volumen medio de cada estado
    -Tiempo total en cada estado"""

    df_copia = df.copy()
    df_copia = seteoSignal(df_copia,tipo=tipo)

    #COMIENZO CON LAS METRICAS
    #creo las medias en base a las señales
    medias = df_copia.groupby("Señal").mean()

    #genero un df nuevo que será el de las metricas, en este caso agrego volatilidad media
    res = pd.DataFrame(medias.Volatilidad)
    res.loc["Total"] = df_copia.Volatilidad.mean()

    #variacion
    res["Variacion"] = medias.Variacion
    res.loc["Total","Variacion"] = df_copia.Variacion.mean()

    #cantidad
    res["Cantidad"] = df_copia.groupby("Señal").size()
    res.loc["Total","Cantidad"] = df_copia.Señal.count()

    #volumen
    res["Volumen"] = medias.Volume /1000000
    res.loc["Total","Volumen"] = df_copia.Volume.mean() /1000000

    #tiempo
    res["TiempoIn"] = (df_copia.groupby("Señal").size() / df_copia.Señal.count()) *100
    res.loc["Total", "TiempoIn"] = 1

    #agrupado = metricas.drop(["Close","Volume","Variacion","Volatilidad"],axis=1)
    #agrupado= agrupado.groupby("Señal").size()

    return res

def seteoSignal(df,tipo="both"): # SETEO EL DF_COPIA para poder sacar metricas
    """ Recibe el df y side de signamMetrics y lo setea para colocar el estado que hubiera estado, ya sea
    comprado, vendido o sin posicion"""

    #reviso cuando da la primer señal
    df = df.reset_index()
    for i in range(len(df)):
        senal = df.iloc[i]["Señal"]
        if senal != "sin señal":
            fechaIni = i
            break

    # reviso cuando da la ultima señal
    for x in range(len(df)-1,0,-1):
        senalFin = df.iloc[x]["Señal"]
        if senalFin != "sin señal":
            fechaFin = x
            break
    df = df.set_index("Date").sort_index()

    # interpolo para que quede el "estado"
    df = df.reset_index()
    for z in range(fechaIni,fechaFin):
        if df.loc[z,"Señal"] == "sin señal" or df.loc[z,"Señal"] == "stop":
            df.loc[z, "Señal"] = np.nan

    df["Señal"] = df.Señal.fillna(method="ffill")
    df = df.set_index("Date")

    # como tomo precios close y no podría comprar al cierre, tomo las variaciones del dia posterior
    df["Señal"] = df.Señal.shift(-1)

    # seteo si es long o short
    if tipo == "long":
        df["Señal"] = np.where(df.Señal == "venta", "sin señal",
                                     df.Señal)  # saco la posicion vendido
        df["Señal"] = np.where(df.Señal == "stop", "sin señal", df.Señal)  # saco la posicion stop
    if tipo == "short":
        df["Señal"] = np.where(df.Señal == "compra", "sin señal", df.Señal)  # saco posicion compra
        df["Señal"] = np.where(df.Señal == "stop", "sin señal", df.Señal)  # saco posicion stop

    return df

""" /// PONGO EN MARCHA LA ESTRATEGIA /// """
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
pd.options.display.max_rows=50
pd.options.display.max_columns=10

#data=getDataExcel("pruebaY")
data=getInfoDiaria("AAPL")
side = "both"

#Medias
media=True
fast=5
slow=20
mediaTipo="EMA"
buy_cr=0
sell_cr=0

#RSI
rsi=True
rsi_q=14
buy_rsi=50
sell_rsi=35

#STOPLOSS
stopActivar=True
porcen=0.9

data=data.drop(["Open","High","Low"],axis=1)
data=getAlertras(data,media=media,fast=fast, slow=slow, buy_cr=buy_cr, sell_cr=sell_cr,mediaTipo=mediaTipo,
                 rsi_q=rsi_q, buy_rsi=buy_rsi, sell_rsi=sell_rsi,rsi=rsi,
                 stopActivado=stopActivar,porcentaje=porcen,tipo=side)

print(data)

acciones=getActions(data,side)
print(acciones)

trades=getTrades(acciones,side)
print(trades)

Python/PythonTraders-master/JupyterNotebook/Curso1/Clase 11 - Exposición TPs/backtest_stop_loss_NicoPacheco.py


"""Tools for melodies used in phonems"""
from typing import List
import random
from os import path
import json

au_clair_de_la_lune = ["C4"] * 3 + ["D4", "E4", "C4", "D4", "E4", "D4"]

chord_progressions =[
    ["C", "Am", "F", "G"],
    ["D", "Bm", "G", "A"],
    ["G", "Em", "C", "D"],
    ["C", "Am", "D", "G"],
    ["D", "Bm", "Em", "A"],
    ["G", "Em", "Am", "D"],
    ["C", "Em", "F", "G"],
    ["G", "Bm", "C", "D"],
    ["C", "F", "C", "G"],
    ["A", "D", "A", "E"],
]

chords_ratios = {
    "major" : [0, 3, 6],
    "minor" : [0, 2, 6],
}

def get_freqs(fundamental : int, ratios : List[int]):
    return [int(round(fundamental * (1 + r/12))) for r in ratios]

def get_harmonies(note_pitch, chord_type="major") -> List[int]:
    return get_freqs(note_pitch, chords_ratios[chord_type])

pitch_filepath = path.join(path.dirname(path.realpath(__file__)), "data/pitches.json")
with open(pitch_filepath) as pitch_file:
    pitches = json.load(pitch_file)

def get_chord_harmonies(chord_str : str) -> List[int]:
    octave = int(chord_str[-1])
    chord = chord_str[:-1]
    if chord.endswith("m"):
        note, chord_type = chord.strip("m"), "minor"
    else:
        note, chord_type = chord, "major"
    return get_harmonies(pitches[note + str(octave)], chord_type)

def get_progression_freqs(progression : List[str], beats_per_measure) -> List[float]:
    track_frequencies = []
    blinker = False
    for chord in progression:
        chords_freqs = get_chord_harmonies(chord)
        for i in range(beats_per_measure):
            track_frequencies.append(chords_freqs[2] if blinker else chords_freqs[0])
            blinker = not blinker
    return track_frequencies

def get_trinote_progression_freqs(progression : List[str], beats_per_measure) -> List[float]:
    track_frequencies = []
    freqs_index = [0, 2, 0, 1]
    for chord in progression:
        chords_freqs = get_chord_harmonies(chord)
        for i in range(beats_per_measure):
            track_frequencies.append(chords_freqs[freqs_index[i]])
    return track_frequencies

def get_random_progression_freqs(progression : List[str], beats_per_measure) -> List[float]:
    track_frequencies = []
    freqs_index = [0, 2, 0, 1]
    for chord in progression:
        chords_freqs = get_chord_harmonies(chord)
        for i in range(beats_per_measure):
            track_frequencies.append(chords_freqs[random.choice(freqs_index)])
    return track_frequencies

class AudioTrack:
    pass

class MelodicSynthTrack(AudioTrack):

    def __init__(self, beats_per_measures, chords_progression):
        pass

class RawAudioSynthTrack(AudioTrack):

    def __init__(self, audio_file):
        pass


import yaml
import torch
import wandb
from tqdm import tqdm
import torch.nn as nn
import torch.optim as optim
from src.bert_model import FeedbackBERT
from src.utils import calculate_scores, create_loaders
import warnings
warnings.filterwarnings("ignore")

softmax = nn.Softmax(dim=1)
train_step, val_step = 1, 1

def train_single_epoch(loader, model, optim, loss_fn):
    global train_step
    total_loss = 0.0
    all_targets = []
    all_preds = []
    for batch in tqdm(loader):
        optim.zero_grad()
        # Data Unpacking
        input_ids = batch["input_ids"].cuda()
        mask = batch["attention_mask"].cuda()
        token_type_ids = batch["token_type_ids"].cuda()
        targets = batch["target"].cuda()

        # Prediction & Loss Calculation
        pred = model(input_ids, mask, token_type_ids)
        loss = loss_fn(pred, targets)
        pred = torch.argmax(softmax(pred), dim=1)
        targets = targets.cpu().detach().numpy().tolist()
        pred = pred.cpu().detach().numpy().tolist()

        # Loss appending
        all_targets += targets
        all_preds += pred
        total_loss += loss.item()
        batch_scores = calculate_scores(targets, pred)

        # logging
        wandb.log(batch_scores, step=train_step)
        wandb.log({"training_loss" : loss.item()})
        train_step += 1
        loss.backward()
        optim.step()
        epoch_scores = calculate_scores(all_targets, all_preds)

    return model, optim, epoch_scores

def validate_single_epoch(loader, model, loss_fn):
    global val_step
    total_loss = 0.0
    all_targets = []
    all_preds = []
    for batch in tqdm(loader):
        # Data Unpacking
        input_ids = batch["input_ids"].cuda()
        mask = batch["attention_mask"].cuda()
        token_type_ids = batch["token_type_ids"].cuda()
        targets = batch["target"].cuda()

        # Prediction & Loss Calculation
        pred = model(input_ids, mask, token_type_ids)
        loss = loss_fn(pred, targets)
        pred = torch.argmax(softmax(pred), dim=1)
        targets = targets.cpu().detach().numpy().tolist()
        pred = pred.cpu().detach().numpy().tolist()

        # Loss appending
        all_targets += targets
        all_preds += pred
        total_loss += loss.item()
        batch_scores = calculate_scores(targets, pred, "validation")

        # logging
        wandb.log(batch_scores, step=val_step)
        wandb.log({"validation_loss" : loss.item()})
        val_step += 1
        epoch_scores = calculate_scores(all_targets, all_preds, "valdation")
    return total_loss, epoch_scores

def train(params):
    print("Process Initiated...")

    # Initializing objects

    print("Initializing DataLoaders...")
    train_loader, val_loader = create_loaders(params["dataset"])
    print("DataLoaders Generated...")
    print("Initializing Model...")
    model = FeedbackBERT(**params["model"]).cuda()
    print("Model Generated...")
    print(model)
    loss_fn = nn.CrossEntropyLoss()
    if params["optimizer"]["name"] == "adam":
        optimizer = optim.Adam(model.parameters(),**params["optimizer"]["hparams"])
    else:
        raise NotImplementedError

    best_loss = torch.inf
    num_epochs = params["training"]["max_epochs"]
    early_stop_counter = 0
    checkpoint_path = params["training"]["chkp_path"]
    wandb.watch(model, log_freq=10)

    best_train_scores = None
    best_val_scores = None

    # Initializing training iterations
    for epoch in range(num_epochs):
        if early_stop_counter == params["training"]["early_stop_count"]:
            print("Iteration Stopper due to Repeatative Degradation...")
            break
        else:
            print(f"Epoch {epoch + 1}: ")
            model.train()
            model, optimizer, train_scores = train_single_epoch(train_loader, model,
                                                optimizer, loss_fn)
            print(train_scores)
            model.eval()
            current_loss, val_scores = validate_single_epoch(val_loader, model, loss_fn)
            print(val_scores)
            if current_loss < best_loss:
                best_loss = current_loss
                early_stop_counter = 0
                best_train_scores = train_scores
                best_val_scores = val_scores
                chkp = {"model_state": model.state_dict(),
                        "epoch": epoch+1}
                torch.save(chkp, checkpoint_path)
                print("Model Weights Updated...")
            else:
                early_stop_counter += 1
            print("\n")
    print("Process Finished...")
    print(f"Best Training Scores : {best_train_scores}")
    print(f"Best Validation Scores : {best_val_scores}")

if __name__ == '__main__':
    config_path = "config/__base__.yaml"
    with open(config_path, "r") as f:
        params = yaml.safe_load(f)
        f.close()
    print(params)
    wandb.init(config=params)
    train_step, val_step = 1, 1
    train(params)