"""
Appdynamics watcher agent for providing job file templates and variables required to ship logs to
appdynamics controller.
"""
import os
import logging

from kube_log_watcher.agents.base import BaseWatcher
from kube_log_watcher.template_loader import load_template

TPL_NAME = 'appdynamics.job.jinja2'

logger = logging.getLogger(__name__)

class AppDynamicsAgent(BaseWatcher):
    def __init__(self, configuration):
        self.dest_path = os.environ.get('WATCHER_APPDYNAMICS_DEST_PATH')

        if not self.dest_path:
            raise RuntimeError('AppDyanmics watcher agent initialization failed. Env variable '
                               'WATCHER_APPDYNAMICS_DEST_PATH must be set.')

        self.cluster_id = configuration['cluster_id']
        self.tpl = load_template(TPL_NAME)

        self.logs = {}
        self._first_run = True

        logger.info('AppDynamics watcher agent initialization complete!')

    @property
    def name(self):
        return 'AppDynamics'

    @property
    def first_run(self):
        return self._first_run

    def add_log_target(self, target: dict):
        """
        Update our log targets, and pick relevant log fields from ``target['kwargs']``
        """
        log = {}
        log['kwargs'] = target['kwargs']
        pod_labels = target['pod_labels']
        container_id = target['id']

        log['kwargs']['app_name'] = pod_labels.get('appdynamics_app')
        log['kwargs']['app_tier'] = pod_labels.get('appdynamics_tier')

        log['job_file_path'] = self._get_job_file_path(container_id)

        self.logs[target['id']] = log

    def remove_log_target(self, container_id):
        job_file = self._get_job_file_path(container_id)

        try:
            del self.logs[container_id]
        except KeyError:
            logger.exception('Failed to remove log target: %s', container_id)

        try:
            os.remove(job_file)
            logger.debug('AppDynamics watcher agent Removed container(%s) job file', container_id)
        except OSError:
            logger.exception('AppDynamics watcher agent Failed to remove job file: %s', job_file)

    def flush(self):
        for log in self.logs.values():
            job_file = log['job_file_path']
            if not os.path.exists(job_file) or self._first_run:
                try:
                    job = self.tpl.render(**log['kwargs'])

                    with open(job_file, 'w') as fp:
                        fp.write(job)

                except Exception:
                    logger.exception('AppDynamics watcher agent failed to write job file %s', job_file)
                else:
                    logger.debug('AppDynamics watcher agent updated job file %s', job_file)

        self._first_run = False

    def _get_job_file_path(self, container_id):
        return os.path.join(self.dest_path, 'container-{}-jobfile.job'.format(container_id))


import asyncio
from contextlib import suppress
from asgiref.sync import sync_to_async

class Periodic:
    def __init__(self, obj, func, time):
        self.func = func
        self.obj=obj
        self.time = time
        self.is_started = False
        self._task = None

    async def start(self):
        if not self.is_started:
            self.is_started = True
            # Start task to call func periodically:
            self._task = asyncio.ensure_future(self._run())

    async def stop(self):
        if self.is_started:
            self.is_started = False
            # Stop task and await it stopped:
            self._task.cancel()
            with suppress(asyncio.CancelledError):
                await self._task

    async def _run(self):
        while True:
            await self.func()
            await asyncio.sleep(self.time)


'''
Timestamps:
Understand - 5:00
Design and Verify - 10:00
Write Code

    1. Understand the problem
        3 other test cases:
            1. s = "(" - False - Correct
            2. s = "([)]" - False - Correct
            3. s = "([])" - True - Correct

    2. Design solution and verify solution
        Create a dictionary of our close parenthesis and their matching opens
        Stack of open parenthesis

        For each parenthesis check if it is a close parenthesis
        If it is:
            Peek the stack
                If stack is empty return false
                If it is the matching open pop the open parenthesis
                If it is not matching return false
        If it is not:
            Add it to the stack

        If stack is not empty:
            return false

        return true
'''

class Solution:
    def isValid(self, s: str) -> bool:
        pairs = {
            ")": "(",
            "]": "[",
            "}": "{"
        }

        stack = []

        # Check each token
        for paren in s:

            # Check if parenthesis is a closing parenthesis
            if paren in pairs.keys():
                # If it's a closing parenthesis and there is no opening it is invalid
                if not stack:
                    return False

                # Check if the open parenthesis is the proper one, if it is remove it, if not it is invalid
                if stack[-1] == pairs[paren]:
                    stack.pop()
                else:
                    return False

            # If it's an open parenthesis just push it to the stack
            else:
                stack.append(paren)

        # If the stack isn't empty there were not enough valid closes
        if stack:
            return False

        # If it has not been invalidated by oone of the prev tests it is valid
        return True


class Solution:
    def exist(self, board: List[List[str]], word: str) -> bool:
        m=len(board)
        n=len(board[0])
        N=len(word)
        dir=[(0,1),(0,-1),(1,0),(-1,0)]
        def dfs(i,j,index):
            if index==N-1:
                return True
            visited[i][j]=True
            for dx,dy in dir:
                nx=i+dx
                ny=j+dy
                if 0<=nx<m and 0<=ny<n and not visited[nx][ny]:
                    if board[nx][ny]==word[index+1] and dfs(nx,ny,index+1):
                        return True
            visited[i][j]=False
            return False

        for i in range(m):
            for j in range(n):
                if board[i][j]==word[0]:
                    visited=[[False]*n for _ in range(m)]
                    if dfs(i,j,0):
                        return True
        return False

LC/Backing Track/79. Word Search/backingTrack.py


""" Given an array of integers nums containing n + 1 integers where each integer is in the range [1, n] inclusive.

There is only one repeated number in nums, return this repeated number.

You must solve the problem without modifying the array nums and uses only constant extra space.

Example 1:

Input: nums = [1,3,4,2,2]
Output: 2
Example 2:

Input: nums = [3,1,3,4,2]
Output: 3 """

from typing import List

_ = """
this can be thought of as a linked list cycle detection problem.
Considering each index of the array as the pointer to the next node, there will
be one index that will be pointed to more than once (as there is a repeated number and we are
Considering indicies as pointers)

 0  1  2  3  4
[1, 3, 4, 2, 2]

(index, value)
(0, 1) -> (1, 3) -> (3, 2) -> (2, 4) -> (4, 2) ->
                                ^               |
                                |               v
                                ----------------

"""

class Solution:
    def findDuplicate(self, nums: List[int]) -> int:
        # similar to the linked list cycle detection problem
        p1 = 0
        p2 = 0

        while 1:
            p1 = nums[p1]
            p2 = nums[nums[p2]]

            if p1 == p2:
                break

        slow = 0

        while p1 != slow:
            slow = nums[slow]
            p2 = nums[p2]

        return nums[slow]

s = Solution()
s.findDuplicate([1, 3, 4, 2, 2])


#!/usr/bin/env python
"""
Part of the World Generator project.

author:  Bret Curtis
license: LGPL v2

This program is free software; you can redistribute it and/or
modify it under the terms of the GNU General Public License
version 2 as published by the Free Software Foundation.

This program is distributed in the hope that it will be useful, but
WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program; if not, write to the Free Software
Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA
02110-1301 USA
"""
#
# Constants for world generation that typically are used across multiple modules
# the main code itself.
#
#

# Direction
NORTH       = [0, -1]
NORTH_EAST  = [1, -1]
EAST        = [1, 0]
SOUTH_EAST  = [1, 1]
SOUTH       = [0, 1]
SOUTH_WEST  = [-1, 1]
WEST        = [-1, 0]
NORTH_WEST  = [-1, -1]
CENTER      = [0, 0]

DIR_NEIGHBORS           = [NORTH, EAST, SOUTH, WEST]
DIR_NEIGHBORS_CENTER    = [CENTER, NORTH, EAST, SOUTH, WEST]
DIR_ALL                 = [NORTH, NORTH_EAST, EAST, SOUTH_EAST, SOUTH, SOUTH_WEST, WEST, NORTH_WEST]
DIR_ALL_CENTER          = [CENTER, NORTH, NORTH_EAST, EAST, SOUTH_EAST, SOUTH, SOUTH_WEST, WEST, NORTH_WEST]

# Heightmap methods
HM_MDA      = 0
HM_DSA      = 1
HM_SPH      = 2
HM_PERLIN   = 3

#Viewer types
VIEWER_HEIGHTMAP    = 0
VIEWER_HEATMAP      = 1
VIEWER_RAINFALL     = 2
VIEWER_WIND         = 3
VIEWER_DRAINAGE     = 4
VIEWER_BIOMES       = 5
VIEWER_RIVERS       = 6
VIEWER_EROSION      = 7
VIEWER_EROSIONAPP   = 8

#Biomes
BIOME_TYPE_UNDEFINED = 0
BIOME_TYPE_WATER = 1
BIOME_TYPE_GRASSLAND = 2
BIOME_TYPE_FOREST = 3
BIOME_TYPE_DESERT_SAND = 4
BIOME_TYPE_DESERT_ROCK = 5
BIOME_TYPE_MOUNTAIN_LOW = 6
BIOME_TYPE_MOUNTAIN_HIGH = 7
BIOME_TYPE_SAVANNA = 8
BIOME_TYPE_MARSH = 9
BIOME_TYPE_SHRUBLAND = 10
BIOME_TYPE_HILLS = 11
BIOME_TYPE_SWAMP = 12
BIOME_TYPE_DESERT_BADLANDS = 13
BIOME_TYPE_MOUNTAIN = 14

# BIOME_ELEVATION
BIOME_ELEVATION_HILLS_LOW = 0.60
BIOME_ELEVATION_HILLS = 0.75
BIOME_ELEVATION_MOUNTAIN_LOW = 0.85
BIOME_ELEVATION_MOUNTAIN = 0.90
BIOME_ELEVATION_MOUNTAIN_HIGH = 0.95

# MOVEMENT
DIAGONAL_COST = 1.41

# World contants
WGEN_HEMISPHERE_NORTH = 1
WGEN_HEMISPHERE_EQUATOR = 2
WGEN_HEMISPHERE_SOUTH = 3
#WGEN_SEA_LEVEL = 0.24
WIND_OFFSET = 180
WIND_PARITY = -1 # -1 or 1
WGEN_WIND_RESOLUTION = 4 # 1 is perfect, higher = rougher
WGEN_RAIN_FALLOFF = 0.2 # Default 0.2 - less for less rain, more for more rain
WGEN_WIND_GRAVITY = 0.975
TEMPERATURE_BAND_RESOLUTION = 2 # 1 is perfect, higher = rougher

#Colour contant
# http://df.magmawiki.com/index.php/Colour (as reference)
COLOR_BLACK = 0x000000
COLOR_CLEAR = 0x808080
COLOR_GRAY = 0x808080
COLOR_SILVER = 0xC0C0C0
COLOR_WHITE = 0xFFFFFF
COLOR_TAUPE_ROSE = 0x905D5D
COLOR_CHESTNUT = 0xCD5C5C
COLOR_MAROON = 0x800000
COLOR_RED = 0xFF0000
COLOR_VERMILION = 0xE34234
COLOR_RUSSET = 0x755A57
COLOR_SCARLET = 0xFF2400
COLOR_BURNT_UMBER = 0x8A3324
COLOR_TAUPE_MEDIUM = 0x674C47
COLOR_DARK_CHESTNUT = 0x986960
COLOR_BURNT_SIENNA = 0xE97451
COLOR_RUST = 0xB7410E
COLOR_AUBURN = 0x6F351A
COLOR_MAHOGANY = 0xC04000
COLOR_PUMPKIN = 0xFF7518
COLOR_CHOCOLATE = 0xD2691E
COLOR_TAUPE_PALE = 0xBC987E
COLOR_TAUPE_DARK = 0x483C32
COLOR_DARK_PEACH = 0xFFDAB9
COLOR_COPPER = 0xB87333
COLOR_LIGHT_BROWN = 0xCD853F
COLOR_BRONZE = 0xCD7F32
COLOR_PALE_BROWN = 0x987654
COLOR_DARK_BROWN = 0x654321
COLOR_SEPIA = 0x704214
COLOR_OCHRE = 0xCC7722
COLOR_BROWN = 0x964B00
COLOR_CINNAMON = 0x7B3F00
COLOR_TAN = 0xD2B48C
COLOR_RAW_UMBER = 0x734A12
COLOR_ORANGE = 0xFFA500
COLOR_PEACH = 0xFFE5B4
COLOR_TAUPE_SANDY = 0x967117
COLOR_GOLDENROD = 0xDAA520
COLOR_AMBER = 0xFFBF00
COLOR_DARK_TAN = 0x918151
COLOR_SAFFRON = 0xF4C430
COLOR_ECRU = 0xC2B280
COLOR_GOLD = 0xD4AF37
COLOR_PEARL = 0xF0EAD6
COLOR_BUFF = 0xF0DC82
COLOR_FLAX = 0xEEDC82
COLOR_BRASS = 0xB5A642
COLOR_GOLDEN_YELLOW = 0xFFDF00
COLOR_LEMON = 0xFDE910
COLOR_CREAM = 0xFFFDD0
COLOR_BEIGE = 0xF5F5DC
COLOR_OLIVE = 0x808000
COLOR_YELLOW = 0xFFFF00
COLOR_IVORY = 0xFFFFF0
COLOR_LIME = 0xCCFF00
COLOR_YELLOW_GREEN = 0x9ACD32
COLOR_DARK_OLIVE = 0x556832
COLOR_GREEN_YELLOW = 0xADFF2F
COLOR_CHARTREUSE = 0x7FFF00
COLOR_FERN_GREEN = 0x4F7942
COLOR_MOSS_GREEN = 0xADDFAD
COLOR_GREEN = 0x00FF00
COLOR_MINT_GREEN = 0x98FF98
COLOR_ASH_GRAY = 0xB2BEB5
COLOR_EMERALD = 0x50C878
COLOR_SEA_GREEN = 0x2E8B57
COLOR_SPRING_GREEN = 0x00FF7F
COLOR_DARK_GREEN = 0x013220
COLOR_JADE = 0x00A86B
COLOR_AQUAMARINE = 0x7FFFD4
COLOR_PINE_GREEN = 0x01796F
COLOR_TURQUOISE = 0x30D5C8
COLOR_PALE_BLUE = 0xAFEEEE
COLOR_TEAL = 0x008080
COLOR_AQUA = 0x00FFFF
COLOR_LIGHT_BLUE = 0xADD8E6
COLOR_CERULEAN = 0x007BA7
COLOR_SKY_BLUE = 0x87CEEB
COLOR_CHARCOAL = 0x36454F
COLOR_SLATE_GRAY = 0x708090
COLOR_MIDNIGHT_BLUE = 0x003366
COLOR_AZURE = 0x007FFF
COLOR_COBALT = 0x0047AB
COLOR_LAVENDER = 0xE6E6FA
COLOR_DARK_BLUE = 0x00008B
COLOR_BLUE = 0x0000FF
COLOR_PERIWINKLE = 0xCCCCFF
COLOR_DARK_VIOLET = 0x423189
COLOR_AMETHYST = 0x9966CC
COLOR_DARK_INDIGO = 0x310062
COLOR_VIOLET = 0x8B00FF
COLOR_INDIGO = 0x4B0082
COLOR_PURPLE = 0x660099
COLOR_HELIOTROPE = 0xDF73FF
COLOR_LILAC = 0xC8A2C8
COLOR_PLUM = 0x660066
COLOR_TAUPE_PURPLE = 0x50404D
COLOR_TAUPE_GRAY = 0x8B8589
COLOR_FUCHSIA = 0xF400A1
COLOR_MAUVE = 0x993366
COLOR_LAVENDER_BLUSH = 0xFFF0F5
COLOR_DARK_PINK = 0xE75480
COLOR_MAUVE_TAUPE = 0x915F6D
COLOR_DARK_SCARLET = 0x560319
COLOR_PUCE = 0xCC8899
COLOR_CRIMSON = 0xDC143C
COLOR_PINK = 0xFFC0CB
COLOR_CARDINAL = 0xC41E3A
COLOR_CARMINE = 0x960018
COLOR_PALE_PINK = 0xFADADD
COLOR_PALE_CHESTNUT = 0xDDADAF
COLOR_DEEPSEA = 0x00003F
COLOR_SEA = 0x00007F
COLOR_GRASSLAND = 0x80FF00
COLOR_HILLS = 0x5A805A


import os
import pytz
from dotenv import load_dotenv
from datetime import date, datetime, timedelta
from json import dumps, loads
from db import fetch_tasks_formatted, get_user, mark_task_as_added
from google_cal import get_calendar_events
from langchain.chat_models import ChatOpenAI
from langchain.schema import HumanMessage
from langchain.prompts.chat import ChatPromptTemplate
from google_cal import GoogleCalendarEvent, add_calendar_event, get_calendar_events

load_dotenv()

OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")

if not OPENAI_API_KEY:
    raise Exception("OPENAI_API_KEY not found in .env file")

chat_model = ChatOpenAI(openai_api_key=OPENAI_API_KEY)

def time_slot_to_datetime(time_slot):
    # Create a datetime object for tomorrow's date
    tomorrow = date.today() + timedelta(days=1)

    # Extract the start time and end time from the time_slot string
    start_time_str, end_time_str = time_slot.split(" - ")

    # Convert the start time and end time strings to datetime.time objects
    if len(start_time_str) == 6:
        start_time_str = "0" + start_time_str

    if len(end_time_str) == 6:
        end_time_str = "0" + end_time_str

    start_time = datetime.strptime(start_time_str, "%I:%M%p").time()
    end_time = datetime.strptime(end_time_str, "%I:%M%p").time()

    # Combine the date object and the time objects to create datetime.datetime objects
    start_datetime = datetime.combine(tomorrow, start_time)
    end_datetime = datetime.combine(tomorrow, end_time)

    return start_datetime, end_datetime

def generate_schedule_from_hobby(hobby, number_of_sessions):
    # Improve the templates in the future -- not mvp
    template = (
        "You are a helpful personal secretary that helps to schedule daily tasks."
    )
    human_template = "I want to learn {hobby}, help me schedule {number_of_sessions} sessions for the week."
    chat_prompt = ChatPromptTemplate.from_messages(
        [
            ("system", template),
            ("human", human_template),
        ]
    )
    chain = chat_prompt | chat_model
    return chain.invoke(
        {"hobby": hobby, "number_of_sessions": number_of_sessions}
    ).content

def plan_tasks(user_id: int):
    user = get_user(user_id)
    refresh_token = user.get("google_refresh_token", "")
    # Get events from tomorrow 12am to tomorrow 11:59pm
    tomorrow = datetime.now(tz=pytz.timezone("America/New_York")) + timedelta(days=1)
    print(tomorrow)
    tomorrow_midnight = datetime(
        tomorrow.year, tomorrow.month, tomorrow.day, 0, 0
    )
    print(tomorrow_midnight)
    events = get_calendar_events(
        refresh_token=refresh_token,
        timeMin=tomorrow_midnight.isoformat() + "Z",
        timeMax=(tomorrow_midnight + timedelta(days=1)).isoformat() + "Z",
        k=20,
    )

    formatted_events = [
        {
            "start_time": datetime.fromisoformat(event["start"]["dateTime"]).strftime(
                "%I:%M%p"
            ),
            "end_time": datetime.fromisoformat(
                event.get("end")["dateTime"]
                if not event.get("endTimeUnspecified", False)
                else datetime.utcnow().isoformat()
            ).strftime("%I:%M%p"),
            "description": event.get("summary", ""),
            "id": event["id"],
        }
        for event in events
    ]

    print("events", events)
    formatted_events_str = dumps(formatted_events)
    print("formatted_events_str", formatted_events_str)

    template = (
        "You are a helpful personal secretary that helps to schedule daily tasks."
    )
    tasks = fetch_tasks_formatted(user_id)
    print(tasks)
    if not tasks:
        return None
    human_template = "This is my schedule for tomorrow (which cannot be changed): \n{formatted_events_str}\nThese are extra tasks I want to do tomorrow, with the lowest id as the most important:\n{tasks}\nHelp me add to my schedule with as many of these tasks as possible. Give a reasonable estimate and you do not have to add all if there is no time. Do NOT create new tasks. Do not return the events or tasks in my existing schedule. Return it in parsable JSON format that can be immediately parsed. The return json should have keys id, task_name and time_slot and nothing else, no other sentences ONLY."
    # It returns in this format:
    # {
    #     "tasks": [
    #         {
    #         "id": 1,
    #         "task_name": "run",
    #         "time_slot": "06:00 AM - 7:00 AM"
    #         },
    #         {
    #         "id": 2,
    #         "task_name": "eat",
    #         "time_slot": "07:30 AM - 8:00 AM"
    #         },
    #         {
    #         "id": 3,
    #         "task_name": "gym",
    #         "time_slot": "08:30 AM - 9:30 AM"
    #         },
    #         {
    #         "id": 4,
    #         "task_name": "school",
    #         "time_slot": "10:00 AM - 4:00 PM"
    #         }
    #     ]
    # }
    chat_prompt = ChatPromptTemplate.from_messages(
        [
            ("system", template),
            ("human", human_template),
        ]
    )
    chain = chat_prompt | chat_model
    res = chain.invoke(
        {"formatted_events_str": formatted_events_str, "tasks": tasks}
    ).content
    print("AI RESPONSE", res)
    resParsed: dict = loads(res)
    print("AI RESPONSE PARSED", resParsed)
    tasks = (
        resParsed["tasks"]
        if "tasks" in resParsed
        else (
            resParsed["events"]
            if "events" in resParsed
            else resParsed[[k for k in resParsed][0]]  # Take the first key
        )
    )

    if tasks:
        # add to calendar
        for task in tasks:
            task_desc = task.get(
                "task", task.get("description", task.get("task_name", ""))
            )
            print(f"Adding {task_desc} to calendar")
            # Split the time_slot into start and end times
            time_slot = task["time_slot"]
            start_time, end_time = time_slot_to_datetime(time_slot=time_slot)
            add_calendar_event(
                refresh_token=refresh_token,
                start_time=start_time,
                end_time=end_time,
                summary=task_desc,
            )
        # for loop to mark tasks as added
        for task in tasks:
            mark_task_as_added(task["id"])
            print(f"Marked added {task['id']}")

    return tasks

def storing_new_input_prompt(user_input: str):
    return """

"""
    template = (
        """
        You're the productivity copilot, tasked with taking note of the user's tasks/events/habits and organizing the user's day in a JSON format.

        You have already come up with the user's schedule for today the previous night. However, the user just sent a new input.

        Sort new inputs as "Event", "Task", or "Habit".

        For New Events: Return a JSON object with the following structure:

        {{
        "type": "Event",
        "command": "/event",
        "description": "<event description>",
        "start_time": "<start time>",
        "end_time": "<end time>"
        }}

        For New Tasks: Return a JSON object with the following structure:

        {{
        "type": "Task",
        "command": "/task",
        "description": "<task description>",
        "deadline": "<deadline>",
        "estimated_duration": "<estimated duration>"
        }}

        For New Habits: Return a JSON object with the following structure:

        {{
        "type": "Habit",
        "command": "/habit",
        "description": "<habit description>",
        "frequency": "<frequency>",
        "estimated_duration": "<estimated duration>"
        }}

        If the user didn't provide any of the required info, represent it with "none".
        """
    )

    human_template = "{user_input}"

    chat_prompt = ChatPromptTemplate.from_messages(
        [
            ("system", template),
            ("human", human_template),
        ]
    )
    chain = chat_prompt | chat_model
    res = chain.invoke({"user_input": user_input}).content
    print("AI RESPONSE", res)

def adjusting_existing_schedule_prompt(user_input: str):
    template = (
        """
        You are the Productivity Co-Pilot, in charge of fine-tuning the user's schedule based on incoming events, tasks, and habits. You've crafted today's schedule, but the user has just sent a new task requiring adjustments.

        Input Formats

        For New Task:
        {{
        "type": "Task",
        "command": "/task",
        "description": "<task description>",
        "deadline": "<deadline>",
        "estimated_duration": "<estimated duration>"
        }}

        For New Habit:
        {{
        "type": "Habit",
        "command": "/habit",
        "description": "<habit description>",
        "frequency": "<frequency>",
        "estimated_duration": "<estimated duration>"
        }}

        For New Event:
        {{
        "type": "Event",
        "command": "/event",
        "description": "<event description>",
        "start_time": "<start time>",
        "end_time": "<end time>"
        }}

        Scheduling Rules
        On receiving a new event, seamlessly integrate it into the existing schedule for the day. If a conflict occurs, confirm with the user before reshuffling.
        Prioritize tasks based on their deadlines. If the new task must be completed today, revise the schedule for the day. If not, add it to an "Uncompleted Tasks" list.
        Aim to leave the current schedule unaltered. If unavoidable, reassign tasks or habits rather than meetings or appointments.
        Mandatory 30-minute lunch slot between 12-2pm and a 30-minute dinner slot between 6-8pm.
        Decisions on schedule adjustments should be made autonomously without consulting the user.
        Ensure zero conflicts in the adjusted schedule.

        Today's Schedule Format
        Today's schedule will be in JSON format like so:
        [ {{"start_time": "09:00AM", "end_time": "10:00AM", "description": "Team Meeting", "id": "xyz1"}}, {{"start_time": "11:00AM", "end_time": "12:00PM", "description": "Work on Project X", "id": "xyz2"}}, ... ]

        Response Format
        Your response should be in the following JSON format:

        {{
            "tasks": [
                {{
                "id": 1,
                "task_name": "Team Meeting",
                "time_slot": "09:00 AM - 10:00 AM"
                }},
                ...
            ]
        }}

        Example Scenario
        Today's Schedule:
        [
            {{"start_time": "09:00AM", "end_time": "10:00AM", "description": "Team Meeting", "id": "xyz1"}},
            {{"start_time": "11:00AM", "end_time": "12:00PM", "description": "Work on Project X", "id": "xyz2"}},
            ...
        ]

        New User Input:
        {{
        "type": "Task",
        "command": "/task",
        "description": "<task description>",
        "deadline": "<deadline>",
        "estimated_duration": "<estimated duration>"
        }}

        Your Response:
        {{ "tasks": [ {{ "id": 1, "task_name": "Submit case study report", "time_slot": "10:30 AM - 01:30 PM" }}, {{ "id": 2, "task_name": "Work on Project X", "time_slot": "01:30 PM - 02:30 PM" }}, {{ "id": 3, "task_name": "Gym", "time_slot": "03:00 PM - 04:30 PM" }} ] }}
        """
    )

    human_template = "{user_input}"

    chat_prompt = ChatPromptTemplate.from_messages(
        [
            ("system", template),
            ("human", human_template),
        ]
    )
    chain = chat_prompt | chat_model
    res = chain.invoke({"user_input": user_input}).content
    print("AI RESPONSE", res)

def night_new_schedule_prompt(user_input: str):
    template = (
        """
       You are the Night-time Productivity Co-Pilot, responsible for crafting the user's schedule for the next day. Take into consideration the user's uncompleted tasks, existing events on the calendar, and habit goals. Generate the most optimized schedule, while adhering to any constraints like meal times.

        Input Formats
        Uncompleted Tasks:
        A list in JSON format like so:
        [ {{"id": 1, "description": "Finish presentation slides", "deadline": "next Tuesday", "estimated_duration": "2hrs"}}, {{"id": 2, "description": "Review marketing proposal", "deadline": "this Friday", "estimated_duration": "1hr"}}, ... ]

        Existing Events on Calendar:
        A list in JSON format similar to the one provided for today's schedule:

        [ {{"start_time": "09:00AM", "end_time": "10:00AM", "description": "Team Meeting", "id": "xyz1"}}, ... ]

        Habit Goals:
        A list in JSON format:
        [ {{"id": 1, "description": "Gym", "frequency": "thrice a week", "estimated_duration": "1.5 hrs"}}, {{"id": 2, "description": "Meditate", "frequency": "daily", "estimated_duration": "0.5 hr"}}, ... ]

        Scheduling Rules
        On receiving a new event, seamlessly integrate it into the existing schedule for the day. If a conflict occurs, confirm with the user before reshuffling.
        Prioritize tasks based on their deadlines. If the new task must be completed today, revise the schedule for the day. If not, add it to an "Uncompleted Tasks" list.
        Aim to leave the current schedule unaltered. If unavoidable, reassign tasks or habits rather than meetings or appointments.
        Mandatory 30-minute lunch slot between 12-2pm and a 30-minute dinner slot between 6-8pm.
        Decisions on schedule adjustments should be made autonomously without consulting the user.
        Ensure zero conflicts in the adjusted schedule.

        Response Format
        Your response should be in the following JSON format:

        {{ "tasks": [ {{ "id": 1, "task_name": "Team Meeting", "time_slot": "09:00 AM - 10:00 AM" }}, ... ] }}

        Example Scenario
        Uncompleted Tasks:
        [ {{"id": 1, "description": "Finish presentation slides", "deadline": "next Tuesday", "estimated_duration": "2hrs"}}, {{"id": 2, "description": "Review marketing proposal", "deadline": "this Friday", "estimated_duration": "1hr"}} ]

        Existing Events:
        [ {{"start_time": "09:00AM", "end_time": "10:00AM", "description": "Team Meeting", "id": "xyz1"}} ]

        Habit Goals:
        [ {{"id": 1, "description": "Gym", "frequency": "thrice a week", "estimated_duration": "1.5 hrs"}}, {{"id": 2, "description": "Meditate", "frequency": "daily", "estimated_duration": "0.5 hr"}} ]

        Your Response:
        {{ "tasks": [ {{ "id": 1, "task_name": "Team Meeting", "time_slot": "09:00 AM - 10:00 AM" }}, {{ "id": 2, "task_name": "Finish presentation slides", "time_slot": "10:30 AM - 12:30 PM" }}, {{ "id": 3, "task_name": "Gym", "time_slot": "01:00 PM - 02:30 PM" }}, {{ "id": 4, "task_name": "Meditate", "time_slot": "03:00 PM - 03:30 PM" }} ] }}
        """
    )

    human_template = "{user_input}"

    chat_prompt = ChatPromptTemplate.from_messages(
        [
            ("system", template),
            ("human", human_template),
        ]
    )
    chain = chat_prompt | chat_model
    res = chain.invoke({"user_input": user_input}).content
    print("AI RESPONSE", res)

# night_new_schedule_prompt(
#     """
#     Uncompleted Tasks:
#     [ {"id": 1, "description": "Finish presentation slides", "deadline": "next Tuesday", "estimated_duration": "2hrs"}, {"id": 2, "description": "Review marketing proposal", "deadline": "this Friday", "estimated_duration": "1hr"} ]

#     Existing Events on Calendar:
#     [ {"start_time": "09:00AM", "end_time": "10:00AM", "description": "Team Meeting", "id": "xyz1"}]

#     Habit Goals:
#     [ {"id": 1, "description": "Gym", "frequency": "thrice a week", "estimated_duration": "1.5 hrs"}, {"id": 2, "description": "Meditate", "frequency": "daily", "estimated_duration": "0.5 hr"} ]
#     """
# )


import numpy as np
from scipy.integrate import quad
from f_P import *

# Challenge A_1_1 --------------------
"""
See P2() in the file f_P.py
"""

# Challenge A_1_2 --------------------
"""
P(a ≤ x ≤ b) stands for the total probability that x will be in the interval [a, b]
"""

# Challenge A_1_3 --------------------
mu = 20
sigma = 2
a1 = mu - sigma
b1 = mu + sigma
a2 = mu - (2*sigma)
b2 = mu + (2*sigma)
a3 = mu - (3*sigma)
b3 = mu + (3*sigma)

p1 = P(a1, b1, mu, sigma)  # Probability for -1*sigma ≤ x-mu ≤ 1*sigma
p2 = P(a2, b2, mu, sigma)  # Probability for -2*sigma ≤ x-mu ≤ 2*sigma
p3 = P(a3, b3, mu, sigma)  # Probability for -3*sigma ≤ x-mu ≤ 3*sigma
print(p1, p2, p3)
"""
p1 = 0.682689492137086
p2 = 0.9544997361036417
p3 = 0.9973002039367398
"""

# Challenge A_1_4 --------------------
"""
FWHM is an acronym for Full Width at Full Maximum
This means the width of the "bump" between the points mu-x1 and mu+x1 where the curve is at half ist maximum value
(which for the gaussian distribution is at mu).
We will have to find x1 so that

f(mu + x1) = f(mu - x1) = f(mu)/2
1/(sigma*sqrt(2*pi)) * exp((-(mu+x1-mu)**2)/(2*sigma**2) = 1/2 * 1/(sigma*sqrt(2*pi)) * exp((-(mu-mu)**2)/(2*sigma**2)
exp((-(x1-mu)**2)/(2*sigma**2) = 1/2 * exp((-(mu-mu)**2)/(2*sigma**2)
exp(-(x1**2)/(2*sigma**2) = 1/2 * exp(0)
2 = exp((x1**2)/(2*sigma**2)
ln(2) = (x1**2)/(2*sigma**2)
2*ln(2)*sigma**2 = x1**2
x1 = sqrt(2*ln(2)*sigma**2)

Since this is the distance to one side, and the bump is symmetric, the distance from the left point where
f(x) = f(mu)/2 to the right point where f(x) = f(mu)/2 is equal to 2*sqrt(2*ln(2)*sigma**2)
This is FWHM
"""


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sat Apr 22 00:35:11 2017

@author: shivamchaturvedi
"""

#A list of numbers can be very unsmooth, meaning very high numbers can be right next to very low numbers.
#This list may represent a smooth path in reality that is masked with random noise (for example,
#satellite trajectories with inaccurate transmission). One way to smooth the values in the list is to
#replace each value with the average of each value's neighbors, including the value itself.

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

#3a

#Let's make a function moving_window_average(x, n_neighbors) that takes a list x and the number of neighbors
# n_neighbors on either side to consider. For each value, the function computes the average of each value's
#neighbors, including themselves. Have the function return a list of these averaged values that is the same
#length as the original list. If there are not enough neighbors (for cases near the edge), substitute the
#original value as many times as there are missing neighbors.
#Use your function to find the moving window sum of x=[0,10,5,3,1,5] and n_neighbors=1.

import random

random.seed(1) #Initialize the basic random number generator.

def moving_window_average(x, n_neighbors=1):
    n = len(x)
    width = n_neighbors*2 + 1
    x = [x[0]]*n_neighbors + x + [x[-1]]*n_neighbors
    # To complete the function,
    # return a list of the mean of values from i to i+width for all values i from 0 to n-1.
    return [sum(x[i:(i+width)]) / width for i in range(n)]

x=[0,10,5,3,1,5]
print(moving_window_average(x, 1))

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

#3b
#Compute and store R=1000 random values from 0-1 as x.
#moving_window_average(x, n_neighbors) is pre-loaded into memory from 3a.
#Compute the moving window average for x for the range of n_neighbors 1-9.
#Store x as well as each of these averages as consecutive lists in a list called Y.

random.seed(1) # This line fixes the value called by your function,
               # and is used for answer-checking.

# write your code here!

R = 1000
x = [random.uniform(0, 1) for i in range(0, 1000)] #Return a random floating point number N such that
   #a <= N <= b for a <= b
   #and b <= N <= a for b < a.
Y = [x] + [moving_window_average(x, n_neighbors) for n_neighbors in range(1, 10)]
print(len(Y))
#-------------------------------------------------

#3c
#moving_window_average(x, n_neighbors=2) and Y are already loaded into memory.
#For each list in Y, calculate and store the range (the maximum minus the minimum) in a new list ranges.
#Print your answer. As the window width increases, does the range of each list increase or decrease?
#Why do you think that is?

ranges = [max(x)-min(x) for x in Y]
print(ranges) #The range decreases, because the average smooths a larger number of neighbors.
#Because the numbers in the original list are just random, we expect the average of many of them to be
#roughly 1 / 2, and more averaging means more smoothness in this value.


import argparse
import multiprocessing
import subprocess
import sys
from datetime import datetime

mutex = multiprocessing.Lock()

def build(wheel):
    print("{}: building {}".format(datetime.now(), wheel), file=sys.stderr)
    process = subprocess.run(
        [sys.executable, "-m", "build", "--wheel", "--no-isolation"], cwd=wheel
    )
    if process.returncode != 0:
        # If the build without isolation fails, try to build it again with
        # isolation. We need to protect this following logic in two ways:
        #   1. Only build one at a time in this process.
        #   2. Retry a few times to get around flakiness.
        success = False
        for _ in range(0, 3):
            with mutex:
                process = subprocess.run(
                    [sys.executable, "-m", "build", "--wheel"],
                    encoding="utf-8",
                    cwd=wheel,
                    capture_output=True,
                )
                if process.returncode != 0:
                    continue
                success = True
                break
        if not success:
            print("STDOUT:", file=sys.stderr)
            print(process.stdout, file=sys.stderr)
            print("STDERR:", file=sys.stderr)
            print(process.stderr, file=sys.stderr)
            raise RuntimeError(
                "{}: ERROR: failed to build {}".format(datetime.now(), wheel)
            )

def main(jobs, wheels):
    with multiprocessing.Pool(jobs) as pool:
        results = []
        for wheel in wheels:
            result = pool.apply_async(build, args=(wheel,))
            results.append(result)
        pool.close()
        for result in results:
            result.get()

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Builds Python wheels.")
    parser.add_argument(
        "wheels",
        metavar="wheel",
        type=str,
        nargs="+",
        help="A path to the extracted wheel directory.",
    )
    parser.add_argument(
        "--jobs",
        type=int,
        default=8,
        help="The number of concurrent build jobs to be executed.",
    )
    args = parser.parse_args()
    exit(main(args.jobs, args.wheels))


# global
from hypothesis import assume, strategies as st
import numpy as np

# local
import ivy
from ivy.functional.ivy.layers import _get_embed_dim
from ivy.functional.frontends.torch.nn.functional.layer_functions import (
    _pack_padded_sequence,
)
from ivy_tests.test_ivy import helpers
from ivy_tests.test_ivy.helpers import handle_frontend_test
from ivy_tests.test_ivy.test_functional.test_nn.test_layers import _mha_helper

# --- Helpers --- #
# --------------- #

@st.composite
def _lstm_helper(draw):
    dtype = draw(helpers.get_dtypes("valid", full=False))

    has_biases = draw(st.booleans())
    bidirectional = draw(st.booleans())
    dropout = draw(st.floats(min_value=0, max_value=0.99))
    train = (
        draw(st.booleans()) and not dropout
    )  # not yet supported by original function
    packed = draw(st.booleans())

    batch_first = draw(st.booleans()) and not packed
    num_batches = draw(st.integers(min_value=1, max_value=5))
    num_layers = draw(st.integers(min_value=1, max_value=3))
    num_directions = 2 if bidirectional else 1
    seq_size = draw(st.integers(min_value=1, max_value=5))
    in_size = draw(st.integers(min_value=1, max_value=3))
    hidden_size = draw(st.integers(min_value=1, max_value=3))

    input = draw(
        helpers.array_values(
            dtype=dtype[0],
            shape=(
                (num_batches, seq_size, in_size)
                if batch_first
                else (seq_size, num_batches, in_size)
            ),
            min_value=0,
            max_value=1,
        )
    )

    init_h = draw(
        helpers.array_values(
            dtype=dtype[0],
            shape=(num_directions * num_layers, num_batches, hidden_size),
            min_value=0,
            max_value=1,
        )
    )
    init_c = draw(
        helpers.array_values(
            dtype=dtype[0],
            shape=(num_directions * num_layers, num_batches, hidden_size),
            min_value=0,
            max_value=1,
        )
    )

    all_weights = []
    for k in range(num_layers):
        for _ in range(num_directions):
            weight_ih = draw(
                helpers.array_values(
                    dtype=dtype[0],
                    shape=(
                        (4 * hidden_size, in_size)
                        if k == 0
                        else (4 * hidden_size, num_directions * hidden_size)
                    ),
                    min_value=0,
                    max_value=1,
                )
            )
            weight_hh = draw(
                helpers.array_values(
                    dtype=dtype[0],
                    shape=(4 * hidden_size, hidden_size),
                    min_value=0,
                    max_value=1,
                )
            )
            all_weights += [weight_ih, weight_hh]
            if has_biases:
                bias_ih = draw(
                    helpers.array_values(
                        dtype=dtype[0],
                        shape=(4 * hidden_size,),
                        min_value=0,
                        max_value=1,
                    )
                )
                bias_hh = draw(
                    helpers.array_values(
                        dtype=dtype[0],
                        shape=(4 * hidden_size,),
                        min_value=0,
                        max_value=1,
                    )
                )
                all_weights += [bias_ih, bias_hh]

    if packed:
        batch_sizes = [seq_size]
        batch_sizes += draw(
            st.lists(
                st.integers(min_value=1, max_value=seq_size),
                min_size=num_batches - 1,
                max_size=num_batches - 1,
            )
        )
        batch_sizes = np.array(draw(st.permutations(batch_sizes)))
        input, batch_sizes = [
            ivy.to_numpy(p) for p in _pack_padded_sequence(input, batch_sizes)
        ]
    else:
        batch_sizes = None

    initial_states = init_h, init_c
    all_weights = tuple(all_weights)
    if batch_sizes is not None:
        dtypes = dtype + ["int64"]
        kwargs = {
            "data": input,
            "batch_sizes": batch_sizes,
            "hx": initial_states,
            "params": all_weights,
            "has_biases": has_biases,
            "num_layers": num_layers,
            "dropout": dropout,
            "train": train,
            "bidirectional": bidirectional,
        }
    else:
        dtypes = dtype
        kwargs = {
            "input": input,
            "hx": initial_states,
            "params": all_weights,
            "has_biases": has_biases,
            "num_layers": num_layers,
            "dropout": dropout,
            "train": train,
            "bidirectional": bidirectional,
            "batch_first": batch_first,
        }
    return dtypes, kwargs

# --- Main --- #
# ------------ #

# lstm
@handle_frontend_test(
    fn_tree="torch.lstm",
    dtypes_kwargs=_lstm_helper(),
    test_with_out=st.just(False),
)
def test_torch_lstm(
    *,
    dtypes_kwargs,
    on_device,
    fn_tree,
    frontend,
    test_flags,
    backend_fw,
):
    dtypes, kwargs = dtypes_kwargs
    # Todo: Debug the function to have this case passing as well
    assume("batch_sizes" not in kwargs or not kwargs["bidirectional"])
    helpers.test_frontend_function(
        input_dtypes=dtypes,
        backend_to_test=backend_fw,
        frontend=frontend,
        test_flags=test_flags,
        fn_tree=fn_tree,
        on_device=on_device,
        **kwargs,
    )

# multi_head_attention_forward
@handle_frontend_test(
    fn_tree="torch.nn.functional.multi_head_attention_forward",
    dtype_mha_args=_mha_helper(same_pre_embed_dim=True, batch_second=True).filter(
        lambda args: args[10] is not None
        and (not args[22] or args[5] is not None)
        and len(set(_get_embed_dim(*args[6:10], args[1]))) == 1
    ),
    test_with_out=st.just(False),
)
def test_torch_multi_head_attention_forward(
    *,
    on_device,
    fn_tree,
    frontend,
    test_flags,
    dtype_mha_args,
    backend_fw,
):
    (
        dtype,
        q,
        k,
        v,
        heads,
        attn_mask,
        in_proj_weight,
        q_proj_weight,
        k_proj_weight,
        v_proj_weight,
        out_proj_weight,
        in_proj_bias,
        out_proj_bias,
        key_padding_mask,
        bias_k,
        bias_v,
        static_k,
        static_v,
        _,
        add_zero_attn,
        dropout_p,
        training,
        is_causal,
        need_weights,
        average_attn_weights,
        batch_first,
    ) = dtype_mha_args
    if k is None and v is None:
        k = v = q
    # re-order the dtypes to match the order of the frontend arguments, not the order
    # of ivy.multi_head_attention's arguments given by _mha_helper
    kwargs = {
        "query": q,
        "key": k,
        "value": v,
        "embed_dim_to_check": q.shape[-1],
        "num_heads": heads,
        "in_proj_weight": in_proj_weight,
        "in_proj_bias": in_proj_bias,
        "bias_k": bias_k,
        "bias_v": bias_v,
        "add_zero_attn": add_zero_attn,
        "dropout_p": dropout_p,
        "out_proj_weight": out_proj_weight,
        "out_proj_bias": out_proj_bias,
        "training": training,
        "key_padding_mask": key_padding_mask,
        "need_weights": need_weights,
        "attn_mask": attn_mask,
        "use_separate_proj_weight": in_proj_weight is None,
        "q_proj_weight": q_proj_weight,
        "k_proj_weight": k_proj_weight,
        "v_proj_weight": v_proj_weight,
        "static_k": static_k,
        "static_v": static_v,
        "average_attn_weights": average_attn_weights,
        "is_causal": is_causal,
    }
    helpers.test_frontend_function(
        input_dtypes=[str(r.dtype) for r in kwargs.values() if ivy.is_array(r)],
        backend_to_test=backend_fw,
        frontend=frontend,
        test_flags=test_flags,
        fn_tree=fn_tree,
        atol=1e-03,
        on_device=on_device,
        test_values=not training or dropout_p == 0.0,
        **kwargs,
    )


import puncher
from enum import Enum
import ctypes as ct
import logging
import commons
import math
import random
import pickle
import tracker
import protocol

class P2pCode(Enum):
    REQUEST = 1
    DATA = 2

CHUNK_SIZE = 1000

class P2pPacket(ct.Structure):
    _fields_ = [("action", ct.c_ushort),
                ("chunk_id", ct.c_ushort),
                ("md5_digest", ct.c_ubyte * 16),
                ("data", ct.c_ubyte * CHUNK_SIZE)]
    def set_action(self, action):
        self.action = action.value
    def get_action(self):
        return P2pCode(self.action)
    def __str__(self):
        return "%s:%s" % (self.get_action(),
            self.chunk_id)
    def set_md5(self, bytes):
        self.md5_digest = protocol.b2cb(bytes, 16)
    def get_md5(self):
        return protocol.cb2b(self.md5_digest)
    def compute_md5(self):
        self.md5_digest = commons.md5sum_bytes(self.data)
    def check_md5(self):
        return self.md5_digest == commons.md5sum_bytes(self.data)
    def set_data(self, bytes):
        self.data = protocol.b2cb(bytes, CHUNK_SIZE)
    def get_data(self):
        return protocol.cb2b(self.data)

class P2pEndpointClient(puncher.PuncherConnClient):
    def __init__(self, r, _id, server_addr=(puncher.PUB_IP, puncher.PORT)):
        puncher.PuncherConnClient.__init__(self, r, _id, server_addr)
        self.set_type(self.addr, P2pPacket)
    def incoming_endpoint(self, p):
        logging.fatal("Endpoint, processing new packet.")
        if p.get_action() != P2pCode.REQUEST:
            logging.fatal("Cannot process anything else but request messages!")

        p.set_data(self.fs.get_chunk(p.chunk_id))
        p.set_action(P2pCode.DATA)

        return p
    @classmethod
    def init_file_service(cls, fs):
        cls.fs = fs
    def request(self, chunk_id):
        p = P2pPacket()
        p.set_action(P2pCode.REQUEST)
        p.chunk_id = chunk_id
        self.send(p)
        p = self.recv()
        self.fs.save_chunk(p.data, p.chunk_id)

class P2pClient(object):
    def __init__(self, dodgy_p, server_ip, in_fd=None):
        self.dodgy_p = dodgy_p
        self.fs = FileService(dodgy_p, in_fd)
        self.t_client = tracker.TrackerClient(dodgy_p.get_addr())
        self.my_id = self.t_client.join()
        logging.info("Client joined with id=%d" % self.my_id)
        P2pEndpointClient.init_file_service(self.fs)
        # Now we start the puncher
        self.p_client = puncher.PuncherClient(self.my_id, P2pEndpointClient,
            (server_ip, puncher.PORT))
        self.p_client.join()
    def download(self, out_fd):

        if out_fd == None:
            out_fd = open(self.dodgy_p.get_name(), "wb")

        download_indices = self.fs.get_incomplete_chunks()
        while len(download_indices) > 0:

            for chunk_id in download_indices:

                logging.debug("Want chunk_id=%d" % chunk_id)
                peer_l = self.t_client.want(chunk_id)
                peer_id = random.choice(peer_l)

                # Now we connect to that peer
                peer = self.p_client.connect_cached(peer_id)
                peer.request(chunk_id)

                # Now we announce
                self.t_client.announce(chunk_id)

            download_indices = self.fs.get_incomplete_chunks()
        self.fs.mem_to_file(out_fd)
    def upload(self):
        n_chunks = self.fs.num_chunks
        logging.info("There are %d chunks in this file, announcing all of them." % n_chunks)
        # Now we announce the chunk_ids
        for chunk_id in range(1, n_chunks+1):
            self.t_client.announce(chunk_id)
    def shutdown(self):
        self.t_client.bye()
        self.p_client.bye()

class FileService(object):
    def __init__(self, dodgy_p, fd=None):
        self.file_name = dodgy_p.get_name()
        self.file_size = dodgy_p.file_size
        self.num_chunks = int(math.ceil(self.file_size/CHUNK_SIZE))
        self.total_chunks = []
        self.is_completed_chunks = []
        if fd==None:
            self.init_mem(dodgy_p)
        else:
            self.file_to_mem(dodgy_p, fd)
    def file_to_mem(self, dodgy_p, fd):
        self.file_mem = []
        for count in range(self.num_chunks):
            self.total_chunks.append(count)
            self.is_completed_chunks.append(True)
            fd.seek(count * CHUNK_SIZE)
            self.file_mem.append(fd.read(CHUNK_SIZE))
    def init_mem(self, dodgy_p):
        for count in range(1, self.num_chunks+1):
              self.total_chunks.append(count)
              self.is_completed_chunks.append(False)

        self.file_mem = [bytes()] * self.num_chunks
    def mem_to_file(self, fd):
        assert(self.get_incomplete_chunks() == [])
        logging.info("Saving to file: %s" % fd.name)
        file_mem_contiguous = bytes()
        for ea_b in self.file_mem:
            file_mem_contiguous+= ea_b
        file_mem_contiguous = file_mem_contiguous[:self.file_size]
        fd.write(file_mem_contiguous)
        fd.close()
    def save_chunk(self, data, chunk_num):
        logging.debug("Save chunk_id=%d: %s" % (chunk_num, str(data)))
        self.file_mem[chunk_num-1] = data
        self.is_completed_chunks[chunk_num-1] = True
    def get_chunk(self, chunk_id):
        data = self.file_mem[chunk_id-1]
        logging.debug("Get chunk_id=%d: %s" % (chunk_id, str(data)))
        return data
    def get_incomplete_chunks(self):
        incomplete_indices = []
        for i,is_complete in zip(self.total_chunks, self.is_completed_chunks):
            if not is_complete:
                incomplete_indices.append(i)
        return incomplete_indices
# vim: expandtab shiftwidth=4 softtabstop=4 textwidth=80:


from PyQt5.QtCore import *
from PyQt5.QtGui import *
from PyQt5.QtWidgets import *

import numpy as np
from numpy2qimage import *

# cv2.color
class ImageCanvas(QLabel):
    rectChanged = pyqtSignal(QRect)

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.rubberBand = QRubberBand(QRubberBand.Rectangle, self)
        self.setMouseTracking(True)
        self.origin = QPoint()
        self.changeRubberBand = False
        self.allowRubberBand = False
        self.setAlignment(Qt.AlignCenter)
        self.viewRect = None

    def setRubberBand(self, state:bool):
        self.allowRubberBand = state

    def mousePressEvent(self, event):
        if self.viewRect and self.viewRect.contains(event.pos()) and self.allowRubberBand and event.button() == Qt.LeftButton:
            self.origin = event.pos()
            self.rubberBand.setGeometry(QRect(self.origin, QSize()))
            #self.rectChanged.emit(self.rubberBand.geometry())
            self.rubberBand.show()
            self.changeRubberBand = True
        super().mousePressEvent(event)

    def mouseMoveEvent(self, event):
        if self.viewRect and self.viewRect.contains(event.pos()) and self.changeRubberBand:
            self.rubberBand.setGeometry(QRect(self.origin, event.pos()).normalized())
            #self.rectChanged.emit(self.rubberBand.geometry())
        super().mouseMoveEvent(event)

    def mouseReleaseEvent(self, event):
        if self.changeRubberBand and event.button() == Qt.LeftButton:
            #self.image2viewTransform.inverted().mapRect(self.rubberBand.geometry())
            view2image = self.image2viewTransform.inverted()[0]
            #print(view2image)
            self.rectChanged.emit(view2image.mapRect(self.rubberBand.geometry()))
            self.rubberBand.hide()
            self.changeRubberBand = False
        super().mouseReleaseEvent(event)

    @pyqtSlot(np.ndarray)
    def setImage(self, image:np.ndarray, rects:list=None, colors:list=None):
        # print("canvas size:", self.width(), self.height())
        qimage = numpy2qimage(image).scaled(self.width(), self.height(), Qt.KeepAspectRatio)
        self.viewRect = QRect(QPoint(int((self.width()-qimage.width())/2.0), int((self.height()-qimage.height())/2.0)), qimage.size())
        # print("self.viewRect:", self.viewRect)
        #self.imageRatio = self.width() / image.shape[1]
        self.image2viewTransform = QTransform(
            qimage.width()/image.shape[1], 0.0,
            0.0, qimage.height()/image.shape[0],
            self.viewRect.topLeft().x(), self.viewRect.topLeft().y()
        )

        if rects and colors:
            for rect,color in zip(rects,colors):
                rt = QRect(rect[1], rect[0], rect[3], rect[2])
                painter = QPainter(qimage)
                painter.setPen(QPen(QColor(color[0],color[1],color[2],100), 3))
                painter.drawRect(self.image2viewTransform.mapRect(rt).translated(-self.viewRect.topLeft()))
                painter.end()
        super().setPixmap(QPixmap.fromImage(qimage))


import numpy as np
import torch
import torch.nn as nn

class fixedDilatedConv(nn.Module):
    def __init__(self, in_channel=4, dil_start=5,
                 dil_end=24,
                 trainable=False, bulge_size=2,
                 binarize_fd=False):
        super(fixedDilatedConv, self).__init__()
        self.binarize_fd = binarize_fd
        self.in_channel = in_channel
        self.out_channel = int(
            (dil_end - dil_start + 1) * 6 * bulge_size * 3)
        self.out_channel = self.out_channel - (
            (dil_end - dil_start + 1) * 6)
        self.dil_start = dil_start
        self.dil_end = dil_end
        self.bulge_size = bulge_size
        self.kernel_size = int(
            dil_end - dil_start + 1) * 2
        self.mid = int(self.kernel_size / 2)
        weights = torch.Tensor(
            torch.zeros(
                self.out_channel,
                self.in_channel,
                self.kernel_size))
        weights = self.inject_weights(weights)
        self.weights = torch.nn.Parameter(weights)
        if not trainable:
            self.weights.requires_grad = False

    def inject_weights(self, weights):
        A, C, G, U = 0, 1, 2, 3
        weight_codes = [
            (A, U, 1), (U, A, 1.25), (C, G, 1.5),
            (G, C, 1.75), (G, U, .5), (U, G, .75)]
        if self.binarize_fd:
            weight_codes = [
                (A, U, 1), (U, A, 1), (C, G, 1),
                (G, C, 1), (G, U, 1), (U, G, 1)]
        self.weight_names = []
        kernel_idx = 0
        for nuc_tuple in weight_codes:
            nuc1, nuc2, adweight = nuc_tuple
            for i in range(self.dil_start, self.dil_end + 1):
                part_1 = round(i / 2)
                part_2 = i - part_1 + 1
                for each_bulge in range(1, self.bulge_size + 1):
                    if each_bulge == 1:
                        curbulges = [-each_bulge, 0, each_bulge]
                    else:
                        curbulges = [-each_bulge, each_bulge]
                    for bulge in curbulges:
                        idx_1 = np.array(
                            [int(self.kernel_size / 2) - part_1,
                             int(self.kernel_size / 2) + part_2])
                        if bulge > 0:
                            idx_1[idx_1 > self.mid] += bulge
                        elif bulge < 0:
                            idx_1[idx_1 < self.mid] -= bulge
                        idx_1 = idx_1[idx_1 < self.kernel_size]
                        idx_1 = idx_1[idx_1 > 0]
                        idx_0 = np.setdiff1d(
                            np.arange(self.kernel_size), idx_1)
                        weights[kernel_idx, :, idx_0] = 0
                        weights[kernel_idx, nuc1, idx_1[0]] = adweight
                        weights[kernel_idx, nuc2, idx_1[-1]] = adweight
                        adname = "{}_{} dil {} bulge {}".format(
                            nuc1, nuc2, i, bulge)
                        self.weight_names.append(adname)
                        kernel_idx += 1
        return weights[:kernel_idx, :, :]

    def forward(self, x):
        out = torch.nn.functional.conv1d(
            x, self.weights)
        if self.binarize_fd:
            out[out < torch.max(out)] = 0
            out[out > 0] = 1
        return out

class stackedDilatedConv(nn.Module):
    def __init__(self, dil_start=5, dil_end=24,
                 kernel_size=16, width=30,
                 inputsize=50):
        super(stackedDilatedConv, self).__init__()
        # self.padding = padding
        self.dil_start = dil_start
        self.dil_end = dil_end
        self.kernel_size = kernel_size
        self.width = width
        self.inputsize = inputsize
        self.padding = 0
        self.padding = self.get_lout(
            dil_end)
        if self.padding < 0:
            self.padding = int(abs(self.padding) + (inputsize * 2))
        self.pad1 = int(self.padding / 2)
        self.pad2 = self.padding - self.pad1
        conv_id = 0
        for i in range(dil_start, dil_end + 1):
            adlayer = nn.Conv1d(
                4, self.width,
                kernel_size=self.kernel_size,
                dilation=i)
            setattr(self, "Dilation.{}".format(conv_id), adlayer)
            conv_id += 1
        self.num_layers = conv_id
        self.width_out = int(self.num_layers * self.width)
        self.kernel_out = self.get_lout(self.dil_start)

    def get_lout(self, dil):
        lout = self.inputsize + 2 * self.padding -\
            dil * (self.kernel_size - 1)
        return lout

    def forward(self, x):
        x_padded = nn.functional.pad(
            x, (self.pad1, self.pad2))
        out = torch.zeros(
            x.shape[0],
            self.width_out,
            self.kernel_out).to(x.device)
        for i in range(self.num_layers):
            width_st = int(i * self.width)
            width_end = width_st + self.width
            adname = "Dilation.{}".format(i)
            temp = getattr(self, adname)(x_padded)
            dim_diff = int(
                (self.kernel_out - temp.shape[-1]) / 2)
            dim_diff_2 = int(self.kernel_out - dim_diff - temp.shape[-1])
            out[:, width_st:width_end, ] = nn.functional.pad(
                temp, (dim_diff, dim_diff_2))
        return out

class ConvArm1(nn.Module):
    def __init__(self, inputsize=50, kernels=[16, 10, 5],
                 widths=[64, 64, 64],
                 pools=[2, 2, 4], dropout=0.5,
                 disable_conv_dp=False):
        super(ConvArm1, self).__init__()
        self.inputsize = inputsize
        self.kernels = kernels
        self.pools = pools
        self.widths = widths
        self.disable_conv_dp = disable_conv_dp
        list1 = [
            nn.Conv1d(
                self.inputsize,
                widths[0], kernel_size=self.kernels[0]),
            nn.BatchNorm1d(widths[0]),
            nn.ReLU()]
        list2 = [
            nn.Conv1d(
                widths[0],
                widths[1], kernel_size=self.kernels[1]),
            nn.BatchNorm1d(widths[1]),
            nn.ReLU()]
        list3 = [
            nn.Conv1d(
                widths[1],
                widths[2], kernel_size=self.kernels[2]),
            nn.BatchNorm1d(widths[2]),
            nn.ReLU()]
        if self.disable_conv_dp:
            list1.append(nn.Dropout(p=dropout))
            list2.append(nn.Dropout(p=dropout))
            list3.append(nn.Dropout(p=dropout))
        list1.append(nn.MaxPool1d(pools[0]))
        list2.append(nn.MaxPool1d(pools[1]))
        list3.append(nn.MaxPool1d(pools[2]))
        self.a1conv1 = nn.Sequential(*list1)
        self.a1conv2 = nn.Sequential(*list2)
        self.a1conv3 = nn.Sequential(*list3)

    def forward(self, x):
        out = self.a1conv1(x)
        out = self.a1conv2(out)
        out = self.a1conv3(out)
        return out

def get_n_params(model):
    pp = 0
    pp2 = 0
    for p in list(model.parameters()):
        nn = 1
        for s in list(p.size()):
            nn = nn*s
        if p.requires_grad:
            pp += nn
        else:
            pp2 += nn
    print("{} trainable and {} non-trainable".format(pp, pp2))
    return pp, pp2

class DeepBind(nn.Module):
    def __init__(self, num_motifs=16, motif_len=24,
                 lin_dim=32, dp=0.5):
        super(DeepBind, self).__init__()
        self.channel_out = num_motifs
        self.kernel_size = motif_len
        self.lin_dim = lin_dim
        self.dp = dp
        self.conv = nn.Sequential(
            nn.Conv1d(
                4, self.channel_out,
                kernel_size=self.kernel_size),
            nn.ReLU(),
            nn.AdaptiveMaxPool1d(self.lin_dim))
        self.feedForward = nn.Sequential(
            nn.Linear(int(self.lin_dim * self.channel_out),
                      self.lin_dim),
            nn.ReLU(),
            nn.Dropout(self.dp),
            nn.Linear(self.lin_dim, 2))

        for m in self.modules():
            if isinstance(m, nn.Conv1d):
                nn.init.kaiming_normal_(
                    m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, nn.Linear):
                nn.init.xavier_uniform_(m.weight)
                m.bias.data.fill_(0.01)

    def forward(self, x):
        x_padded = nn.functional.pad(
            x, (self.kernel_size, self.kernel_size))
        out = self.conv(x_padded)
        out = out.reshape(out.shape[0], -1)
        out = self.feedForward(out)
        return out

class PythiaModel(nn.Module):
    def __init__(self, init_channels=64, kernel_size=16,
                 inputsize=50,
                 use_fixedConv=True, dp=0.25, trainable=False,
                 dil_start=5, dil_end=24, bulge_size=2,
                 binarize_fd=False, disable_conv_dp=False,
                 num_outputs=2):
        super(PythiaModel, self).__init__()
        self.disable_conv_dp = disable_conv_dp
        self.use_fixedConv = use_fixedConv
        self.init_channels = init_channels
        self.kernel_size = kernel_size
        self.inputsize = inputsize
        self.num_outputs = num_outputs
        self.out_channel_fd = int(
            (dil_end - dil_start + 1) * 6 * bulge_size * 3)
        # Dilating layers
        if self.use_fixedConv:
            if trainable:
                self.conv1_dil = fixedDilatedConv(
                    dil_start=dil_start, dil_end=dil_end,
                    bulge_size=bulge_size, trainable=True,
                    binarize_fd=binarize_fd)
            else:
                self.conv1_dil = fixedDilatedConv(
                    dil_start=dil_start, dil_end=dil_end,
                    bulge_size=bulge_size, trainable=False,
                    binarize_fd=binarize_fd)
        else:
            out_channel = int(
                (dil_end - dil_start + 1) * 6 * bulge_size * 3)
            self.dil_dim = out_channel - ((dil_end - dil_start + 1) * 6)
            self.conv1_reg = stackedDilatedConv(
                inputsize=self.inputsize,
                dil_start=dil_start, dil_end=dil_end)
        self.pool_dil = nn.AdaptiveMaxPool1d(60)
        self.process_dil = nn.Sequential(
            nn.BatchNorm1d(60),
            nn.ReLU())
        # arm 1 dilations
        self.arm1_dil = ConvArm1(
            60,
            kernels=[16, 10, 5],
            pools=[2, 2, 2],
            dropout=dp,
            disable_conv_dp=self.disable_conv_dp)
        self.arm2_reg = ConvArm1(
            4,
            widths=[128, 64, 64],
            kernels=[12, 7, 5],
            pools=[2, 2, 2],
            dropout=dp,
            disable_conv_dp=self.disable_conv_dp)
        self.lin_dim = self.get_lindim()
        self.dense_1 = nn.Linear(
            self.lin_dim,
            64)
        self.dense_relu_1 = nn.ReLU()
        self.dense_2 = nn.Linear(
            64, 32)
        self.dense_relu_2 = nn.ReLU()
        self.dense_3 = nn.Linear(
            32, self.num_outputs)

        for m in self.modules():
            if isinstance(m, nn.Conv1d):
                nn.init.kaiming_normal_(
                    m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, (nn.BatchNorm1d, nn.GroupNorm)):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                nn.init.xavier_uniform_(m.weight)
                m.bias.data.fill_(0.01)

    def get_lindim(self):
        x = torch.rand(1, 4, self.inputsize)
        out = self.forward_dil(x)
        return out.shape[1]

    def get_bn_dim(self):
        x = torch.rand(
            1, 4, self.inputsize)
        x_padded = nn.functional.pad(
            x, (25, 24))
        if self.use_fixedConv:
            out = self.conv1_dil(x_padded)
        else:
            out = self.conv1_reg(x_padded)
        return out.shape[2]

    def forward_dil(self, x):
        x_padded = nn.functional.pad(
            x, (x.shape[2], x.shape[2]))
        if self.use_fixedConv:
            out = self.conv1_dil(x_padded)
        else:
            out = self.conv1_reg(x_padded)
        out = self.pool_dil(out)
        out = out.reshape(
            out.shape[0], out.shape[2], -1)
        out = self.process_dil(out)
        out1 = self.arm1_dil(out)
        out2 = self.arm2_reg(x)
        out_merged = torch.cat(
            (out1, out2), 2)
        out = out_merged.view(
            out_merged.shape[0], -1)
        return out

    def forward(self, x):
        out = self.forward_dil(x)
        out_mfe = self.dense_1(out)
        out_mfe = self.dense_relu_1(out_mfe)
        out_mfe = self.dense_2(out_mfe)
        out_mfe = self.dense_relu_2(out_mfe)
        out_mfe = self.dense_3(out_mfe)
        return out_mfe

def conv1_dim(lin, kernel, pad=1, dil=1, stride=1):
    lout = (
        ((lin + 2 * pad - dil * (kernel - 1) - 1) / stride) + 1)
    return lout

def test_model(inputsize, use_fixedConv=True,
               kernel_size=16, trainable=False):
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    print("i {} fixed {} k {} trainable {}".format(
            inputsize, use_fixedConv, kernel_size,
            trainable))
    net = PythiaModel(
        inputsize=inputsize,
        kernel_size=kernel_size,
        use_fixedConv=use_fixedConv,
        trainable=trainable)
    net.to(device)
    print(net)
    in_tensor = torch.rand(
        (10, 4, inputsize)).to(device)
    out_mfe = net(in_tensor)
    print("Input: {}\nOutput: {}".format(
            in_tensor.shape, out_mfe.shape))

def to_one_hot(a, num_classes):
    return np.squeeze(np.eye(num_classes)[a.reshape(-1)])

def get_x():
    np.random.seed(42)
    x = np.zeros((32, 4, 50))
    for i in range(32):
        temp_seq = np.random.choice(
            np.array([0, 1, 2, 3]),
            50, replace=True)
        ad_hot = to_one_hot(temp_seq, 4)
        x[i] = ad_hot.reshape(4, -1)
    return x

def test_fixed():
    np.random.seed(42)
    addict = {"A": 0, "C": 1, "G": 2, "U": 3}
    str_seq = np.random.choice(
        np.array(["A", "C", "G", "U"]),
        50, replace=True)
    str_seq[15:32] = list("ACUUAGGUAACCUAGAA")
    temp_seq = np.array(
        [addict[each] for each in str_seq])
    ad_hot = to_one_hot(temp_seq, 4)
    x = torch.from_numpy(ad_hot.T).reshape(1, 4, -1)
    x = torch.zeros(1, 4, 50)
    for i in range(50):
        x[0, :, i] = torch.from_numpy(ad_hot[i])
    self = fixedDilatedConv(dil_start=2, dil_end=21)
    dict_data = {"Seq": str_seq[:40]}
    for i in np.arange(self.weights.shape[0]):
        adname = self.weight_names[i]
        cur_weight = self.weights[i]
        res_val = torch.mul(cur_weight, x[0, :, :40])
        dict_data[adname] = sum(res_val).detach().numpy()
    import pandas as pd
    outdf = pd.DataFrame(dict_data)
    outdf.to_csv("Example_data_stemloop.tsv", sep="\t")

def test_fixed_2():
    np.random.seed(42)
    addict = {"A": 0, "C": 1, "G": 2, "U": 3}
    str_seq = "ACGCUUACAAAGUAAGCGU"
    filters = ["NNNNCNNNGNNNN",
               "NNNANNNNNUNNN",
               "NNUNNNNNNNANN",
               "NUNNNNNNNNNAN",
               "CNNNNNNNNNNNG"]
    temp_seq = np.array(
        [addict[each] for each in str_seq])
    ad_hot = to_one_hot(temp_seq, 4)
    x = torch.from_numpy(ad_hot.T).reshape(1, 4, -1)
    filters_temp = torch.zeros(5, 4, 19)
    i = 0
    for each_f in filters:
        j = 0
        for each_nuc in list(each_f):
            if each_nuc != "N":
                idx_nuc = addict[each_nuc]
                filters_temp[i, idx_nuc, j] = 1
            j += 1
        i += 1
    x_padded = nn.functional.pad(x, (9, 9))
    outval = torch.nn.functional.conv1d(x_padded.float(), filters_temp.float())
    print(outval[0].T)

if __name__ == "__main__":
    test_model(50, use_fixedConv=False, kernel_size=16)
    test_model(50, use_fixedConv=True, kernel_size=16)
    print("Script passed basic tests")


import time
from configparser import ConfigParser, SectionProxy
from contextlib import suppress
from enum import Enum
from typing import Optional

from builder_const import Building, Config, BuildingCategory, BUILDING_BASE_URL
from selenium.common import NoSuchElementException
from selenium.webdriver.chrome.webdriver import WebDriver
from selenium.webdriver.common.by import By
from selenium.webdriver.remote.webelement import WebElement
from termcolor import cprint
from utils import do_click, get_table_rows

class ExpansionStatus(Enum):
    to_build = 0
    done = 1
    in_progress = 2
    waiting_other = 3
    other = 4

class Expansions(str, Enum):
    rsd = "Rozbudowa o stanowisko dowodzenia"
    prisoner_van = "Rozbudowa o policyjne więźniarki"
    opp = "Rozbudowa o Oddział Prewencji Policji"
    wrd = "Rozbudowa o Wydział Ruchu Drogowego"
    prison_cells = "Cele"
    prison_cell_1 = "Cela więzienna"
    prison_cell_n = "Dodatkowa cela"
    pr = "Rozbudowa dla pojazdów proszkowych"
    containers = "Rozbudowa dla kontenerów"

def build_expansions(driver: WebDriver, building: Building, config: Config) -> None:
    driver.get(f"{BUILDING_BASE_URL}{building.id}")
    expansions_target_raw = config.ini[building.category.name]
    expansions_target = {k: _get_expansion_target_value(v) for k, v in expansions_target_raw.items()}

    do_click(driver, driver.find_element(By.XPATH, '//*[@id="tabs"]/li[2]/a'))
    time.sleep(0.3)

    current_expansions = get_expansions_status(driver)
    cprint(f'Current: {current_expansions}, target: {expansions_target}', 'yellow')
    to_build = {k: max(0, v - current_expansions.get(k, 0)) for k, v in expansions_target.items()}
    cprint(f'To build: {to_build}', 'yellow')

    while sum(to_build.values()) > 0:
        result = queue_expansions(driver, to_build)
        if not result:
            break

def _get_expansion_target_value(value: str) -> int:
    if value == 'False':
        return 0
    if value == 'True':
        return 1
    return int(value)

def _get_expansion(name: WebElement) -> Optional[Expansions]:
    with suppress(ValueError):
        parsed = name.find_element(By.TAG_NAME, 'b').get_attribute('innerHTML').strip()
        expansions = Expansions(parsed)
        if expansions == Expansions.prison_cell_n or expansions == expansions.prison_cell_1:
            return Expansions.prison_cells
    return None

def _get_status(actions: WebElement) -> ExpansionStatus:
    try:
        action_href = actions.find_element(By.TAG_NAME, 'a').get_attribute('href')
        if 'extension_cancel' in action_href:
            return ExpansionStatus.in_progress
        if 'extension_finish' in action_href:
            return ExpansionStatus.in_progress
        if 'extension_ready' in action_href:
            return ExpansionStatus.done
        if 'extension/credits' in action_href:
            return ExpansionStatus.to_build
    except NoSuchElementException:
        if "Pozostały czas" in actions.text:
            return ExpansionStatus.in_progress

    try:
        if "Pozostały czas" in actions.get_attribute('text'):
            return ExpansionStatus.in_progress
    except Exception as err:
        pass
    return ExpansionStatus.waiting_other

def get_expansions_status(driver: WebDriver) -> dict:
    rows = get_table_rows(driver, class_name="table")
    already_built = {}
    for row in rows:
        name, _, _, actions = row.find_elements(By.TAG_NAME, 'td')
        expansion = _get_expansion(name)
        if not expansion:
            continue

        expansion_status = _get_status(actions)
        if expansion_status in [ExpansionStatus.done, ExpansionStatus.in_progress]:
            already_built[expansion.name] = already_built.get(expansion.name, 0) + 1

    return already_built

def queue_expansions(driver: WebDriver, to_build: dict) -> bool:
    rows = get_table_rows(driver, class_name="table")
    for row in rows:
        name, _, _, actions = row.find_elements(By.TAG_NAME, 'td')
        expansion = _get_expansion(name)
        print(name, expansion)
        if not expansion:
            continue

        expansion_status = _get_status(actions)
        print(expansion_status, to_build.get(expansion.name), actions.text)
        if expansion_status is ExpansionStatus.to_build and to_build.get(expansion.name, 0) > 0:
            to_build[expansion.name] = to_build[expansion.name] - 1
            do_click(driver, actions.find_element(By.TAG_NAME, 'a'))
            cprint(f'Expansion was queued {expansion}', 'green')
            return True

    return False


# -*- coding: utf-8 -*-
"""
Created on Thu Feb  8 21:51:26 2018

@author: ROLLE
"""

def tccCheck(d1, d2):
        #Using set to check whether two dictionaries are match on Title, Country, City, regardless Score
        shared_items = set(d1.items()) & set(d2 .items())
        country_check = False
        title_check = False
        city_check = False
        for e in shared_items:
            if 'country' in e:
                country_check = True
            if 'title' in e:
                title_check = True
            if 'city' in e:
                city_check = True
        if (country_check and title_check and city_check):
            return True
        else:
            return False
def cateCheck(d1, d2):
    shared_items = set(d1.items()) & set(d2 .items())
    category_check = False
    score_check = False
    for e in shared_items:
        if 'category' in e:
            category_check = True
        if 'score' in e:
            score_check = True
    if (category_check and score_check):
        return True
    else:
        return False


import numpy as np
import matplotlib.pyplot as plt
from integrate_ode import solve_ode
from numerical_shooting import orbit_calc
from ode_functions import predator_prey

# Predator-Prey equation for b<0.26
def predator_prey_below(t, u, b=0.22):
    a = 1
    d = 0.1
    x, y = u
    dx = x * (1 - x) - (a * x * y) / (d + x)
    dy = b * y * (1 - (y / x))
    return np.array((dx, dy))

# Predator-Prey equation for b>0.26
def predator_prey_above(t, u, b=0.3):
    a = 1
    d = 0.1
    x, y = u
    dx = x * (1 - x) - (a * x * y) / (d + x)
    dy = b * y * (1 - (y / x))
    return np.array((dx, dy))

# Function plots
def simulate_plot(f):
    t = np.linspace(0, 200, 1000)
    x0 = [0.24, 0.24]
    sol = solve_ode(f, x0, t, 'rk4', 0.01)
    plt.plot(t, sol[0], label='x(t)')
    plt.plot(t, sol[1], label='y(t)')
    plt.xlabel('t')
    plt.ylabel('x(t) and y(t)')
    plt.legend(loc='best')
    plt.show()

# Predator-Prey shooting example
def predator_prey_shooting():
    shoot_result = orbit_calc(predator_prey, [0.2, 0.3, 22], var=0)
    [x, y] = shoot_result[:-1]
    T = shoot_result[-1]
    print('Calculated Predator-Prey coordinates using numerical shooting: ' + str(x) + ', ' + str(y))
    print('Calculated Predator-Prey time period using numerical shooting: ' + str(T))
    solve_t = np.linspace(0, T, 100)
    realsol = solve_ode(predator_prey, [x, y], solve_t, 'rk4', 0.01)
    orbit_x = realsol[0]
    orbit_y = realsol[1]
    plt.plot(orbit_x, orbit_y, 'g', label="Shooting Orbit")
    plt.plot(x, y, 'rx', label='Shooting Point')
    plt.xlabel('x(t)')
    plt.xlabel('y(t)')
    plt.title('Periodic Orbit calculated using Numerical Shooting')
    plt.legend()
    plt.show()

if __name__ == '__main__':
    # Simulate Predator-Prey for b<0.26
    simulate_plot(lambda t, u: predator_prey_below(t, u))
    # Simulate Predator-Prey for b>0.26
    simulate_plot(lambda t, u: predator_prey_above(t, u))
    # Predator-Prey numerical shooting for b=0.26
    predator_prey_shooting()


# @Author: Pieter Blok
# @Date:   2021-03-22 09:43:07
# @Last Modified by:   Pieter Blok
# @Last Modified time: 2021-11-08 15:09:29

#!/usr/bin/env python

## based on: detectron2.modeling.roi_heads.box_head
## based on: detectron2.modeling.roi_heads.fast_rcnn

import torch
from torch import nn
import numpy as np
import logging
from typing import Dict, List, Tuple, Union
from fvcore.nn import giou_loss, smooth_l1_loss
import fvcore.nn.weight_init as weight_init
from torch import nn
from torch.nn import functional as F
from detectron2.config import configurable, get_cfg
from detectron2.layers import Conv2d, ConvTranspose2d, ShapeSpec, batched_nms, cat, cross_entropy, nonzero_tuple, get_norm
from detectron2.modeling.box_regression import Box2BoxTransform
from detectron2.structures import Boxes, Instances
from detectron2.utils.events import get_event_storage
from detectron2.utils.registry import Registry

## these specific libraries are needed to alter the network architecture
from detectron2.modeling.roi_heads import ROI_HEADS_REGISTRY, StandardROIHeads, Res5ROIHeads
from detectron2.modeling.roi_heads.box_head import ROI_BOX_HEAD_REGISTRY
from detectron2.modeling.roi_heads.fast_rcnn import fast_rcnn_inference, _log_classification_stats
from detectron2.modeling.roi_heads.mask_head import ROI_MASK_HEAD_REGISTRY, BaseMaskRCNNHead

@ROI_BOX_HEAD_REGISTRY.register()
class FastRCNNConvFCHeadDropout(nn.Sequential):
    """
    A head with several 3x3 conv layers (each followed by norm & relu) and then
    several fc layers (each followed by relu).
    """

    @configurable
    def __init__(
        self, input_shape: ShapeSpec, *, conv_dims: List[int], fc_dims: List[int], conv_norm="",
        dropout_probability: float = 0.5,
    ):
        """
        NOTE: this interface is experimental.

        Args:
            input_shape (ShapeSpec): shape of the input feature.
            conv_dims (list[int]): the output dimensions of the conv layers
            fc_dims (list[int]): the output dimensions of the fc layers
            conv_norm (str or callable): normalization for the conv layers.
                See :func:`detectron2.layers.get_norm` for supported types.
        """
        super().__init__()
        assert len(conv_dims) + len(fc_dims) > 0

        self._output_size = (input_shape.channels, input_shape.height, input_shape.width)

        self.conv_norm_relus = []
        for k, conv_dim in enumerate(conv_dims):
            conv = Conv2d(
                self._output_size[0],
                conv_dim,
                kernel_size=3,
                padding=1,
                bias=not conv_norm,
                norm=get_norm(conv_norm, conv_dim),
                activation=nn.ReLU(),
            )
            self.add_module("conv{}".format(k + 1), conv)
            self.conv_norm_relus.append(conv)
            self._output_size = (conv_dim, self._output_size[1], self._output_size[2])

        self.fcs = []
        for k, fc_dim in enumerate(fc_dims):
            if k == 0:
                self.add_module("flatten", nn.Flatten())
            dropout = nn.Dropout(p=dropout_probability)
            self.add_module("dropout{}".format(k + 1), dropout)
            fc = nn.Linear(int(np.prod(self._output_size)), fc_dim)
            self.add_module("fc{}".format(k + 1), fc)
            self.add_module("fc_relu{}".format(k + 1), nn.ReLU())
            self.fcs.append(fc)
            self._output_size = fc_dim

        for layer in self.conv_norm_relus:
            weight_init.c2_msra_fill(layer)
        for layer in self.fcs:
            weight_init.c2_xavier_fill(layer)

    @classmethod
    def from_config(cls, cfg, input_shape):
        num_conv = cfg.MODEL.ROI_BOX_HEAD.NUM_CONV
        conv_dim = cfg.MODEL.ROI_BOX_HEAD.CONV_DIM
        num_fc = cfg.MODEL.ROI_BOX_HEAD.NUM_FC
        fc_dim = cfg.MODEL.ROI_BOX_HEAD.FC_DIM
        return {
            "input_shape": input_shape,
            "conv_dims": [conv_dim] * num_conv,
            "fc_dims": [fc_dim] * num_fc,
            "conv_norm": cfg.MODEL.ROI_BOX_HEAD.NORM,
            "dropout_probability": cfg.MODEL.ROI_BOX_HEAD.DROPOUT_PROBABILITY,
        }

    def forward(self, x):
        for layer in self:
            x = layer(x)
        return x

    @property
    @torch.jit.unused
    def output_shape(self):
        """
        Returns:
            ShapeSpec: the output feature shape
        """
        o = self._output_size
        if isinstance(o, int):
            return ShapeSpec(channels=o)
        else:
            return ShapeSpec(channels=o[0], height=o[1], width=o[2])

class FastRCNNOutputLayersDropout(nn.Module):
    """
    Two linear layers for predicting Fast R-CNN outputs:

    1. proposal-to-detection box regression deltas
    2. classification scores
    """

    @configurable
    def __init__(
        self,
        input_shape: ShapeSpec,
        *,
        box2box_transform,
        num_classes: int,
        test_score_thresh: float = 0.0,
        test_nms_thresh: float = 0.5,
        test_topk_per_image: int = 100,
        cls_agnostic_bbox_reg: bool = False,
        smooth_l1_beta: float = 0.0,
        box_reg_loss_type: str = "smooth_l1",
        loss_weight: Union[float, Dict[str, float]] = 1.0,
        softmaxes: bool = False,
        dropout_probability: float = 0.5,
    ):
        """
        NOTE: this interface is experimental.

        Args:
            input_shape (ShapeSpec): shape of the input feature to this module
            box2box_transform (Box2BoxTransform or Box2BoxTransformRotated):
            num_classes (int): number of foreground classes
            test_score_thresh (float): threshold to filter predictions results.
            test_nms_thresh (float): NMS threshold for prediction results.
            test_topk_per_image (int): number of top predictions to produce per image.
            cls_agnostic_bbox_reg (bool): whether to use class agnostic for bbox regression
            smooth_l1_beta (float): transition point from L1 to L2 loss. Only used if
                `box_reg_loss_type` is "smooth_l1"
            box_reg_loss_type (str): Box regression loss type. One of: "smooth_l1", "giou"
            loss_weight (float|dict): weights to use for losses. Can be single float for weighting
                all losses, or a dict of individual weightings. Valid dict keys are:
                    * "loss_cls": applied to classification loss
                    * "loss_box_reg": applied to box regression loss
        """
        super().__init__()
        if isinstance(input_shape, int):  # some backward compatibility
            input_shape = ShapeSpec(channels=input_shape)
        self.num_classes = num_classes
        input_size = input_shape.channels * (input_shape.width or 1) * (input_shape.height or 1)
        # prediction layer for num_classes foreground classes and one background class (hence + 1)
        self.dropout1 = nn.Dropout(p=dropout_probability)
        self.cls_score = nn.Linear(input_size, num_classes + 1)
        num_bbox_reg_classes = 1 if cls_agnostic_bbox_reg else num_classes
        box_dim = len(box2box_transform.weights)
        self.dropout2 = nn.Dropout(p=dropout_probability)
        self.bbox_pred = nn.Linear(input_size, num_bbox_reg_classes * box_dim)

        nn.init.normal_(self.cls_score.weight, std=0.01)
        nn.init.normal_(self.bbox_pred.weight, std=0.001)
        for l in [self.cls_score, self.bbox_pred]:
            nn.init.constant_(l.bias, 0)

        self.box2box_transform = box2box_transform
        self.smooth_l1_beta = smooth_l1_beta
        self.test_score_thresh = test_score_thresh
        self.test_nms_thresh = test_nms_thresh
        self.test_topk_per_image = test_topk_per_image
        self.box_reg_loss_type = box_reg_loss_type
        if isinstance(loss_weight, float):
            loss_weight = {"loss_cls": loss_weight, "loss_box_reg": loss_weight}
        self.loss_weight = loss_weight
        self.softmaxes = softmaxes

    @classmethod
    def from_config(cls, cfg, input_shape):
        return {
            "input_shape": input_shape,
            "box2box_transform": Box2BoxTransform(weights=cfg.MODEL.ROI_BOX_HEAD.BBOX_REG_WEIGHTS),
            # fmt: off
            "num_classes"           : cfg.MODEL.ROI_HEADS.NUM_CLASSES,
            "cls_agnostic_bbox_reg" : cfg.MODEL.ROI_BOX_HEAD.CLS_AGNOSTIC_BBOX_REG,
            "smooth_l1_beta"        : cfg.MODEL.ROI_BOX_HEAD.SMOOTH_L1_BETA,
            "test_score_thresh"     : cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST,
            "test_nms_thresh"       : cfg.MODEL.ROI_HEADS.NMS_THRESH_TEST,
            "test_topk_per_image"   : cfg.TEST.DETECTIONS_PER_IMAGE,
            "box_reg_loss_type"     : cfg.MODEL.ROI_BOX_HEAD.BBOX_REG_LOSS_TYPE,
            "loss_weight"           : {"loss_box_reg": cfg.MODEL.ROI_BOX_HEAD.BBOX_REG_LOSS_WEIGHT},
            "softmaxes"             : cfg.MODEL.ROI_HEADS.SOFTMAXES,
            "dropout_probability"   : cfg.MODEL.ROI_BOX_HEAD.DROPOUT_PROBABILITY,
            # fmt: on
        }

    def forward(self, x):
        """
        Args:
            x: per-region features of shape (N, ...) for N bounding boxes to predict.

        Returns:
            (Tensor, Tensor):
            First tensor: shape (N,K+1), scores for each of the N box. Each row contains the
            scores for K object categories and 1 background class.

            Second tensor: bounding box regression deltas for each box. Shape is shape (N,Kx4),
            or (N,4) for class-agnostic regression.
        """
        if x.dim() > 2:
            x = torch.flatten(x, start_dim=1)
        x = self.dropout1(x)
        scores = self.cls_score(x)
        x = self.dropout2(x)
        proposal_deltas = self.bbox_pred(x)
        return scores, proposal_deltas

    def losses(self, predictions, proposals):
        """
        Args:
            predictions: return values of :meth:`forward()`.
            proposals (list[Instances]): proposals that match the features that were used
                to compute predictions. The fields ``proposal_boxes``, ``gt_boxes``,
                ``gt_classes`` are expected.

        Returns:
            Dict[str, Tensor]: dict of losses
        """
        scores, proposal_deltas = predictions

        # parse classification outputs
        gt_classes = (
            cat([p.gt_classes for p in proposals], dim=0) if len(proposals) else torch.empty(0)
        )
        _log_classification_stats(scores, gt_classes)

        # parse box regression outputs
        if len(proposals):
            proposal_boxes = cat([p.proposal_boxes.tensor for p in proposals], dim=0)  # Nx4
            assert not proposal_boxes.requires_grad, "Proposals should not require gradients!"
            # If "gt_boxes" does not exist, the proposals must be all negative and
            # should not be included in regression loss computation.
            # Here we just use proposal_boxes as an arbitrary placeholder because its
            # value won't be used in self.box_reg_loss().
            gt_boxes = cat(
                [(p.gt_boxes if p.has("gt_boxes") else p.proposal_boxes).tensor for p in proposals],
                dim=0,
            )
        else:
            proposal_boxes = gt_boxes = torch.empty((0, 4), device=proposal_deltas.device)

        losses = {
            "loss_cls": cross_entropy(scores, gt_classes, reduction="mean"),
            "loss_box_reg": self.box_reg_loss(
                proposal_boxes, gt_boxes, proposal_deltas, gt_classes
            ),
        }
        return {k: v * self.loss_weight.get(k, 1.0) for k, v in losses.items()}

    def box_reg_loss(self, proposal_boxes, gt_boxes, pred_deltas, gt_classes):
        """
        Args:
            All boxes are tensors with the same shape Rx(4 or 5).
            gt_classes is a long tensor of shape R, the gt class label of each proposal.
            R shall be the number of proposals.
        """
        box_dim = proposal_boxes.shape[1]  # 4 or 5
        # Regression loss is only computed for foreground proposals (those matched to a GT)
        fg_inds = nonzero_tuple((gt_classes >= 0) & (gt_classes < self.num_classes))[0]
        if pred_deltas.shape[1] == box_dim:  # cls-agnostic regression
            fg_pred_deltas = pred_deltas[fg_inds]
        else:
            fg_pred_deltas = pred_deltas.view(-1, self.num_classes, box_dim)[
                fg_inds, gt_classes[fg_inds]
            ]

        if self.box_reg_loss_type == "smooth_l1":
            gt_pred_deltas = self.box2box_transform.get_deltas(
                proposal_boxes[fg_inds],
                gt_boxes[fg_inds],
            )
            loss_box_reg = smooth_l1_loss(
                fg_pred_deltas, gt_pred_deltas, self.smooth_l1_beta, reduction="sum"
            )
        elif self.box_reg_loss_type == "giou":
            fg_pred_boxes = self.box2box_transform.apply_deltas(
                fg_pred_deltas, proposal_boxes[fg_inds]
            )
            loss_box_reg = giou_loss(fg_pred_boxes, gt_boxes[fg_inds], reduction="sum")
        else:
            raise ValueError(f"Invalid bbox reg loss type '{self.box_reg_loss_type}'")
        # The reg loss is normalized using the total number of regions (R), not the number
        # of foreground regions even though the box regression loss is only defined on
        # foreground regions. Why? Because doing so gives equal training influence to
        # each foreground example. To see how, consider two different minibatches:
        #  (1) Contains a single foreground region
        #  (2) Contains 100 foreground regions
        # If we normalize by the number of foreground regions, the single example in
        # minibatch (1) will be given 100 times as much influence as each foreground
        # example in minibatch (2). Normalizing by the total number of regions, R,
        # means that the single example in minibatch (1) and each of the 100 examples
        # in minibatch (2) are given equal influence.
        return loss_box_reg / max(gt_classes.numel(), 1.0)  # return 0 if empty

    def inference(self, predictions: Tuple[torch.Tensor, torch.Tensor], proposals: List[Instances]):
        """
        Args:
            predictions: return values of :meth:`forward()`.
            proposals (list[Instances]): proposals that match the features that were
                used to compute predictions. The ``proposal_boxes`` field is expected.

        Returns:
            list[Instances]: same as `fast_rcnn_inference`.
            list[Tensor]: same as `fast_rcnn_inference`.
        """
        boxes = self.predict_boxes(predictions, proposals)
        scores = self.predict_probs(predictions, proposals)
        image_shapes = [x.image_size for x in proposals]

        return fast_rcnn_inference(
            boxes,
            scores,
            image_shapes,
            self.test_score_thresh,
            self.test_nms_thresh,
            self.test_topk_per_image,
            self.softmaxes,
        )

    def predict_boxes_for_gt_classes(self, predictions, proposals):
        """
        Args:
            predictions: return values of :meth:`forward()`.
            proposals (list[Instances]): proposals that match the features that were used
                to compute predictions. The fields ``proposal_boxes``, ``gt_classes`` are expected.

        Returns:
            list[Tensor]:
                A list of Tensors of predicted boxes for GT classes in case of
                class-specific box head. Element i of the list has shape (Ri, B), where Ri is
                the number of proposals for image i and B is the box dimension (4 or 5)
        """
        if not len(proposals):
            return []
        scores, proposal_deltas = predictions
        proposal_boxes = cat([p.proposal_boxes.tensor for p in proposals], dim=0)
        N, B = proposal_boxes.shape
        predict_boxes = self.box2box_transform.apply_deltas(
            proposal_deltas, proposal_boxes
        )  # Nx(KxB)

        K = predict_boxes.shape[1] // B
        if K > 1:
            gt_classes = torch.cat([p.gt_classes for p in proposals], dim=0)
            # Some proposals are ignored or have a background class. Their gt_classes
            # cannot be used as index.
            gt_classes = gt_classes.clamp_(0, K - 1)

            predict_boxes = predict_boxes.view(N, K, B)[
                torch.arange(N, dtype=torch.long, device=predict_boxes.device), gt_classes
            ]
        num_prop_per_image = [len(p) for p in proposals]
        return predict_boxes.split(num_prop_per_image)

    def predict_boxes(
        self, predictions: Tuple[torch.Tensor, torch.Tensor], proposals: List[Instances]
    ):
        """
        Args:
            predictions: return values of :meth:`forward()`.
            proposals (list[Instances]): proposals that match the features that were
                used to compute predictions. The ``proposal_boxes`` field is expected.

        Returns:
            list[Tensor]:
                A list of Tensors of predicted class-specific or class-agnostic boxes
                for each image. Element i has shape (Ri, K * B) or (Ri, B), where Ri is
                the number of proposals for image i and B is the box dimension (4 or 5)
        """
        if not len(proposals):
            return []
        _, proposal_deltas = predictions
        num_prop_per_image = [len(p) for p in proposals]
        proposal_boxes = cat([p.proposal_boxes.tensor for p in proposals], dim=0)
        predict_boxes = self.box2box_transform.apply_deltas(
            proposal_deltas,
            proposal_boxes,
        )  # Nx(KxB)
        return predict_boxes.split(num_prop_per_image)

    def predict_probs(
        self, predictions: Tuple[torch.Tensor, torch.Tensor], proposals: List[Instances]
    ):
        """
        Args:
            predictions: return values of :meth:`forward()`.
            proposals (list[Instances]): proposals that match the features that were
                used to compute predictions.

        Returns:
            list[Tensor]:
                A list of Tensors of predicted class probabilities for each image.
                Element i has shape (Ri, K + 1), where Ri is the number of proposals for image i.
        """
        scores, _ = predictions
        num_inst_per_image = [len(p) for p in proposals]
        probs = F.softmax(scores, dim=-1)
        return probs.split(num_inst_per_image, dim=0)

@ROI_HEADS_REGISTRY.register()
class Res5ROIHeadsDropout(Res5ROIHeads):
    def __init__(self, cfg, input_shape):
        super().__init__(cfg, input_shape, box_predictor=FastRCNNOutputLayersDropout(cfg, 2048))

@ROI_HEADS_REGISTRY.register()
class StandardROIHeadsDropout(StandardROIHeads):
    def __init__(self, cfg, input_shape):
        super().__init__(cfg, input_shape, box_predictor=FastRCNNOutputLayersDropout(cfg, 1024))

@ROI_MASK_HEAD_REGISTRY.register()
class MaskRCNNConvUpsampleHeadDropout(BaseMaskRCNNHead, nn.Sequential):
    """
    A mask head with several conv layers, plus an upsample layer (with `ConvTranspose2d`).
    Predictions are made with a final 1x1 conv layer.
    """

    @configurable
    def __init__(self, input_shape: ShapeSpec, *, num_classes, conv_dims, conv_norm="",
            dropout_probability: float = 0.5, **kwargs):
        """
        NOTE: this interface is experimental.

        Args:
            input_shape (ShapeSpec): shape of the input feature
            num_classes (int): the number of foreground classes (i.e. background is not
                included). 1 if using class agnostic prediction.
            conv_dims (list[int]): a list of N>0 integers representing the output dimensions
                of N-1 conv layers and the last upsample layer.
            conv_norm (str or callable): normalization for the conv layers.
                See :func:`detectron2.layers.get_norm` for supported types.
        """
        super().__init__(**kwargs)
        assert len(conv_dims) >= 1, "conv_dims have to be non-empty!"

        self.conv_norm_relus = []

        cur_channels = input_shape.channels
        for k, conv_dim in enumerate(conv_dims[:-1]):
            conv = Conv2d(
                cur_channels,
                conv_dim,
                kernel_size=3,
                stride=1,
                padding=1,
                bias=not conv_norm,
                norm=get_norm(conv_norm, conv_dim),
                activation=nn.ReLU(),
            )
            self.add_module("mask_fcn{}".format(k + 1), conv)
            self.conv_norm_relus.append(conv)
            cur_channels = conv_dim

        self.dropout1 = nn.Dropout(p=dropout_probability)

        self.deconv = ConvTranspose2d(
            cur_channels, conv_dims[-1], kernel_size=2, stride=2, padding=0
        )
        self.add_module("deconv_relu", nn.ReLU())
        cur_channels = conv_dims[-1]

        self.dropout2 = nn.Dropout(p=dropout_probability)

        self.predictor = Conv2d(cur_channels, num_classes, kernel_size=1, stride=1, padding=0)

        for layer in self.conv_norm_relus + [self.deconv]:
            weight_init.c2_msra_fill(layer)
        # use normal distribution initialization for mask prediction layer
        nn.init.normal_(self.predictor.weight, std=0.001)
        if self.predictor.bias is not None:
            nn.init.constant_(self.predictor.bias, 0)

    @classmethod
    def from_config(cls, cfg, input_shape):
        ret = super().from_config(cfg, input_shape)
        conv_dim = cfg.MODEL.ROI_MASK_HEAD.CONV_DIM
        num_conv = cfg.MODEL.ROI_MASK_HEAD.NUM_CONV
        ret.update(
            conv_dims=[conv_dim] * (num_conv + 1),  # +1 for ConvTranspose
            conv_norm=cfg.MODEL.ROI_MASK_HEAD.NORM,
            input_shape=input_shape,
        )
        if cfg.MODEL.ROI_MASK_HEAD.CLS_AGNOSTIC_MASK:
            ret["num_classes"] = 1
        else:
            ret["num_classes"] = cfg.MODEL.ROI_HEADS.NUM_CLASSES
        ret["dropout_probability"] = cfg.MODEL.ROI_MASK_HEAD.DROPOUT_PROBABILITY
        return ret

    def layers(self, x):
        for layer in self:
            x = layer(x)
        return x