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()


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"} ]
#     """
# )


# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details
# http://www.gnu.org/licenses/gpl-3.0.txt

from ...client import PeerFilter
from .base import ItemWidgetBase, ListWidgetBase
from .peer import TUICOLUMNS
from .utils import stringify_torrent_filter

class PeerItemWidget(ItemWidgetBase):
    palette_unfocused = 'peerlist'

    @property
    def id(self):
        return self.data['id']

    @property
    def torrent_id(self):
        return self.data['tid']

class PeerListWidget(ListWidgetBase):
    tuicolumns      = TUICOLUMNS
    ListItemClass   = PeerItemWidget
    keymap_context  = 'peer'
    palette_name    = 'peerlist'
    focusable_items = False

    def __init__(self, srvapi, keymap, tfilter=None, pfilter=None, columns=None, sort=None, title=None):
        super().__init__(srvapi, keymap, columns=columns, sort=sort, title=title)
        self._tfilter = tfilter
        self._pfilter = pfilter
        self._secondary_filter = None

        # Create peer filter generator
        if pfilter is not None:
            def filter_peers(peers):
                yield from pfilter.apply(peers)
        else:
            def filter_peers(peers):
                yield from peers
        self._maybe_filter_peers = filter_peers

        self._poller = self._srvapi.create_poller(
            self._srvapi.torrent.torrents, tfilter, keys=('peers', 'name', 'id')
        )
        self._poller.on_response(self._handle_peers)

    def _handle_peers(self, response):
        if response is None or not response.torrents:
            self.clear()
        else:
            # Auto-generate title from our filters if not set
            if self._title_name is None:
                self._title_name = stringify_torrent_filter(self._tfilter, response.torrents)
                if self._pfilter:
                    self._title_name += ' %s' % self._pfilter

            # Create list items our base widget can handle
            def peers_combined(torrents):
                for t in torrents:
                    yield from self._maybe_filter_peers(t['peers'])
            self._data_dict = {p['id']:p for p in peers_combined(response.torrents)}
        self._invalidate()

    def clear(self):
        for w in self._listbox.body:
            w.data.clearcache()
        super().clear()

    def refresh(self):
        self._poller.poll()

    @property
    def sort(self):
        return self._sort

    @sort.setter
    def sort(self, sort):
        ListWidgetBase.sort.fset(self, sort)
        self._poller.poll()

    @property
    def secondary_filter(self):
        return self._secondary_filter

    @secondary_filter.setter
    def secondary_filter(self, peer_filter):
        if peer_filter is None:
            self._secondary_filter = None
        else:
            self._secondary_filter = PeerFilter(peer_filter)
        self._invalidate()

    def _limit_items(self, peer_widgets):
        # Combine primary and secondary peer filters
        pfilter = self._pfilter
        spfilter = self._secondary_filter
        if pfilter is None:
            pfilter = spfilter
        elif spfilter is not None:
            pfilter = pfilter & spfilter

        if pfilter is not None:
            for pw in peer_widgets:
                if not pfilter.match(pw.data):
                    yield pw


# 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 socket
import select
import errno
import os
import time
from threading import Thread
import json
import hashlib
import sys
import hashlib
import datetime
from pathlib import Path
from _thread import *

# get configurations
config = json.load(open(f"{os.path.dirname(os.path.abspath(__file__))}/config.json"))

CLIENT_ID = int(os.path.basename(Path(os.path.realpath(__file__)).parent).split('_')[1])
IP = config['client']['ip_address']
PORT = config['server']['ports'][CLIENT_ID]
HEADER_LENGTH = config['header_length']
META_LENGTH = config['meta_length']
THREAD_PORTS = config['client']['ports']
DOWNLOAD_FOLDER_NAME = config['client']['download_folder_name']
REDOWNLOAD_TIME = config['redownload_times']
LOG = open(f"{os.path.dirname(os.path.abspath(__file__))}/{config['client']['log_file']}", "a")

# Logs messages
def log_this(msg):
    # print(msg)
    LOG.write(f"{datetime.datetime.now()} {msg}\n")
    LOG.flush()

# Verbose function
def help():
    print("\n*** INFO ON FUNCTIONS ***\n")
    print("[function_name] [options] [parameters] - [description]\n")
    print("get_files_list - gets the file names from all the clients\n")
    print("download -p <client_number> <file_name> ... <file_name> - downloads one or more files serially or parallely from the target client. To download serially, use it without the -p option\n")
    print("help - prints verbose for functions\n")
    print("quit - exits client interface\n")

# send updated directory to server
def update_server():
    try:
        list_of_dir = os.listdir(f"{os.path.dirname(os.path.abspath(__file__))}/{DOWNLOAD_FOLDER_NAME}/")
        list_of_dir = '\n'.join(list_of_dir)
        list_of_dir = f"update_list {list_of_dir}".encode('utf-8')
        list_of_dir_header = f"{len(list_of_dir):<{HEADER_LENGTH}}".encode('utf-8')
        meta = f"{f'{CLIENT_ID}':<{META_LENGTH}}".encode('utf-8')
        central_socket.send(list_of_dir_header + meta + list_of_dir)

    except:
        # client closed connection, violently or by user
        return False

# daemon that updates the directory to the server whenever a new file is added or an file is deleted
def folder_watch_daemon(current_file_directory):
    while True:
        temp = os.listdir(f"{os.path.dirname(os.path.abspath(__file__))}/{DOWNLOAD_FOLDER_NAME}/")
        if current_file_directory != temp:
            update_server()
            current_file_directory = temp

# Waiting for a list of directories from the server
def wait_for_list(full_command):
    # Encode command to bytes, prepare header and convert to bytes, like for username above, then send
    full_command = full_command.encode('utf-8')
    command_header = f"{len(full_command):<{HEADER_LENGTH}}".encode('utf-8')
    meta = f"{'':<{META_LENGTH}}".encode('utf-8')
    central_socket.send(command_header + meta + full_command)

    log_this(f"Client sent command: {full_command}")

    start = time.time()

    # Keep trying to recieve until client recieved returns from the server
    while True:
        try:
            # Receive our "header" containing username length, it's size is defined and constant
            header = central_socket.recv(HEADER_LENGTH)

            # If we received no data, server gracefully closed a connection, for example using socket.close() or socket.shutdown(socket.SHUT_RDWR)
            if not len(header):
                log_this(f"Connection closed by the server")
                return

            # Convert header to int value
            header = int(header.decode('utf-8').strip())

            # Get meta data
            meta = central_socket.recv(META_LENGTH)
            meta = meta.decode('utf-8').strip()

            # Receive and decode msg
            dir_list = central_socket.recv(header).decode('utf-8')
            dir_list = json.loads(dir_list)

            # Print List
            for client in dir_list:
                print(f"Client with id {client}:")
                for file in dir_list[client]:
                    print(f"\t{file}")

            end = time.time()

            log_this(f"FileQueryComplete: {(end-start)*1000} ms.")
            # Break out of the loop when list is recieved
            break

        except IOError as e:

            if e.errno != errno.EAGAIN and e.errno != errno.EWOULDBLOCK:
                log_this('Reading error: {}'.format(str(e)))
                return

            # We just did not receive anything
            continue

        except Exception as e:
            # Any other exception - something happened, exit
            log_this('Reading error: '.format(str(e)))
            return

# Handles command receiving
def receive_command(client_socket):

    try:

        # Receive our "header" containing command length, it's size is defined and constant
        command_header = client_socket.recv(HEADER_LENGTH)

        # If we received no data, client gracefully closed a connection, for example using socket.close() or socket.shutdown(socket.SHUT_RDWR)
        if not len(command_header):
            return False

        # Get meta data
        meta = client_socket.recv(META_LENGTH)
        meta = meta.decode('utf-8').strip()

        # Convert header to int value
        command_length = int(command_header.decode('utf-8').strip())

        # Return an object of command header and command data
        return {'header': command_header, 'meta': meta, 'data': client_socket.recv(command_length)}

    except:
        # client closed connection, violently or by user
        return False

# waiting function for parallelized/serial file download
def parallelize_wait_for_file_download(client_socket, files):

    # Encode command to bytes, prepare header and convert to bytes, like for username above, then send
    full_command = f"download {' '.join(files)}".encode('utf-8')
    command_header = f"{len(full_command):<{HEADER_LENGTH}}".encode('utf-8')
    meta = f"{'':<{META_LENGTH}}".encode('utf-8')
    client_socket.send(command_header + meta + full_command)

    log_this(f"Client sent command: {full_command}")

    # open files
    fds = [open(f"{os.path.dirname(os.path.abspath(__file__))}/{DOWNLOAD_FOLDER_NAME}/{files[i]}",'w') for i in range(len(files))]
    files_closed = 0
    redownload_count = 0

    # md5 reconstruction
    m = [hashlib.md5() for _ in range(len(files))]

    # Keep trying to recieve until client recieved returns from the server
    while True:
        try:
            start = time.time()
            header = client_socket.recv(HEADER_LENGTH)

            # If we received no data, server gracefully closed a connection, for example using socket.close() or socket.shutdown(socket.SHUT_RDWR)
            if not len(header):
                log_this('Connection closed by the server')
                return

            # Convert header to int value
            header = int(header.decode('utf-8').strip())

            # Get meta data
            meta = client_socket.recv(META_LENGTH)
            meta = meta.decode('utf-8').strip()
            meta = meta.split(' ')

            # Recieve line and convert to string
            line = client_socket.recv(header).decode('utf-8')

            # if there is any error, remove all files
            if meta[0] == 'ERROR':
                log_this(line)

                for i in range(len(files)):
                    fds[i].flush()
                    fds[i].close()
                    os.remove(f"{os.path.dirname(os.path.abspath(__file__))}/{DOWNLOAD_FOLDER_NAME}/{files[i]}")
                break

            # Flush and close and files is finished recieving
            elif meta[0] == 'END':
                fds[int(meta[1])].flush()
                fds[int(meta[1])].close()
                files_closed += 1

                # if there is contamination in the checksum, log and delete file
                if m[int(meta[1])].hexdigest() != line:
                    log_this(f"Incorrect checksum for file : {files[int(meta[1])]}")
                    log_this(f"Deleting file : {files[int(meta[1])]}")
                    os.remove(f"{os.path.dirname(os.path.abspath(__file__))}/{DOWNLOAD_FOLDER_NAME}/{files[int(meta[1])]}")

            # continue to write and flush to files
            else:
                m[int(meta[0])].update(line.encode('utf-8'))
                fds[int(meta[0])].write(line)
                fds[int(meta[0])].flush()

            # when all files are closed/downloaded sucessfully then we can break from the loop
            if files_closed == len(fds):
                end = time.time()
                log_this(f"DownloadComplete: {(end-start)*1000} ms. Downloaded Files are {' '.join(files)}")
                break

        except IOError as e:

            if e.errno != errno.EAGAIN and e.errno != errno.EWOULDBLOCK:

                if redownload_count < REDOWNLOAD_TIME:
                    redownload_count += 1
                    continue

                log_this('Reading error: {}'.format(str(e)))
                return

            # We just did not receive anything
            continue

        except Exception as e:
            # Any other exception - something happened, exit
            log_this('Reading error: {}'.format(str(e)))
            return

# Waiting for the file contents from the server
def wait_for_file_download(full_command):

    parameters = full_command.split(' ')[1:]

    # check for parallelism option
    if parameters[0] == '-p':
        if len(parameters) <= 2:
            log_this("ParameterError: Too less parameters")
            return
        else:
            parallelize = True
            target_client = int(parameters[1])
            files = parameters[2:]
    else:
        if len(parameters) <= 1:
            log_this("ParameterError: Too less parameters")
            return
        else:
            parallelize = False
            target_client = int(parameters[0])
            files = parameters[1:]

    # if the target client is itself, don't do anything
    if target_client == CLIENT_ID:
        log_this("WrongClient: Target Client is Current Client.")
        return

    # Compute ports that 'this' peer will try to connect
    client_thread_ports = [i+((len(config['client']['ports'])+1)*target_client) for i in THREAD_PORTS]

    # initialize connections with the other peer
    client_sockets = []

    for i in range(len(client_thread_ports)):
        temp = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        temp.connect((IP, client_thread_ports[i]))
        temp.setblocking(False)

        # Initialize conneciton with the server
        temp.send(username_header + meta + username)
        client_sockets.append(temp)

    # starts waiting for file download
    if parallelize:
        for i in range(0, len(files), len(client_thread_ports)):
            thread_idx = 0
            threads = []
            for j in range(i,i+ len(client_thread_ports)):
                if j < len(files):
                    t = Thread(target=parallelize_wait_for_file_download, args=(client_sockets[thread_idx], [files[j]],))
                    t.start()
                    threads.append(t)
                    thread_idx += 1

            for t in threads:
                t.join()
    else:
        # start = time.time()
        t = Thread(target=parallelize_wait_for_file_download, args=(client_sockets[0], files,))
        t.start()

        # end = time.time()
        # log_this(f"DownloadComplete: {(end-start)*1000} ms. Downloaded Files are {' '.join(files)}")

    return

# Sends file to the peer
def send_files(peer_socket, peers, files):
    try:
        fds = [open(f"{os.path.dirname(os.path.abspath(__file__))}/{DOWNLOAD_FOLDER_NAME}/{files[i]}",'r') for i in range(len(files))]

        for i in range(len(files)):
            # using md5 checksum
            m = hashlib.md5()

            while True:

                # read line
                line = fds[i].readline()

                if not line:
                    line = m.hexdigest().encode('utf-8')
                    line_header = f"{len(line):<{HEADER_LENGTH}}".encode('utf-8')
                    meta = f"{f'END {i}':<{META_LENGTH}}".encode('utf-8')
                    peer_socket.send(line_header + meta + line)

                    log_this(f"{files[i]} was sent to {peers[peer_socket]['data']}")
                    break

                line = line.encode('utf-8')

                # update md5 checksum
                m.update(line)

                line_header = f"{len(line):<{HEADER_LENGTH}}".encode('utf-8')
                meta = f"{f'{i}':<{META_LENGTH}}".encode('utf-8')
                peer_socket.send(line_header + meta + line)

            fds[i].close()

    except Exception as e:
        # client closed connection, violently or by user
        error = str(e).encode('utf-8')
        header = f"{len(error):<{HEADER_LENGTH}}".encode('utf-8')
        meta = f"{'ERROR':<{META_LENGTH}}".encode('utf-8')
        peer_socket.send(header + meta + error)
        return False

# A daemon that listens for download requests from any other clients/peers
def server_daemon():
    # Create list of listening server sockets
    server_sockets = []

    # Compute ports that 'this' peer will listen to as a server
    server_thread_ports = [i+((len(config['client']['ports'])+1)*CLIENT_ID) for i in THREAD_PORTS]

    for i in range(len(server_thread_ports)):
        temp = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        temp.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        temp.bind((IP, server_thread_ports[i]))
        temp.listen()
        server_sockets.append(temp)

    # Create list of sockets for select.select()
    sockets_list = [i for i in server_sockets]

    # List of connected clients - socket as a key, user header and name as data
    peers = {}

    for port in server_thread_ports:
        log_this(f'Listening for connections on {IP}:{port}...')

    while True:
        read_sockets, _, exception_sockets = select.select(sockets_list, [], sockets_list)

        for notified_socket in read_sockets:

            # If notified socket is a server socket - new connection, accept it
            if notified_socket in server_sockets:

                client_socket, client_address = server_sockets[server_sockets.index(notified_socket)].accept()

                # Client should send his name right away, receive it
                user = receive_command(client_socket)

                # If False - client disconnected before he sent his name
                if user is False:
                    continue

                # Add accepted socket to select.select() list
                sockets_list.append(client_socket)

                # Also save username and username header
                peers[client_socket] = user

                #logging
                log_msg = 'Accepted new connection from {}:{}, username: {}'.format(*client_address, user['data'].decode('utf-8'))
                log_this(log_msg)

            # Else existing socket is sending a command
            else:

                # Recieve command
                command = receive_command(notified_socket)

                # If False, client disconnected, cleanup
                if command is False:
                    log_msg = '{} Closed connection from: {}'.format(datetime.datetime.now(), peers[notified_socket]['data'].decode('utf-8'))
                    log_this(log_msg)

                    # remove connections
                    sockets_list.remove(notified_socket)
                    del peers[notified_socket]
                    continue

                # Get user by notified socket, so we will know who sent the command
                user = peers[notified_socket]

                # Get command
                command_msg = command["data"].decode("utf-8")
                command_msg = command_msg.split(' ')

                # logging
                log_msg = f'{datetime.datetime.now()} Recieved command from {user["data"].decode("utf-8")}: {command_msg[0]}\n'
                log_this(log_msg)

                # Handle commands
                if command_msg[0] == 'download':
                    start_new_thread(send_files, (notified_socket,peers,command_msg[1:],))

            # handle some socket exceptions just in case
            for notified_socket in exception_sockets:

                # remove connections
                sockets_list.remove(notified_socket)
                del peers[notified_socket]

if __name__ == "__main__":

    # Start the peer's server daemon
    start_new_thread(server_daemon,())

    # allow time for the server daemon to start up
    time.sleep(3)

    # Manual Mode of Client Interface
    if len(sys.argv) == 1:

        # create username to connect to the server
        my_username = input("Username: ")
        username = my_username.encode('utf-8')
        username_header = f"{len(username):<{HEADER_LENGTH}}".encode('utf-8')
        meta = f"{'':<{META_LENGTH}}".encode('utf-8')

        # Initialize connection with the server with central socket
        central_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        central_socket.connect((IP, PORT))
        central_socket.setblocking(False)
        central_socket.send(username_header + meta + username)

        # Initialize file directory to server
        update_server()

        # get current file directories
        current_file_directory = os.listdir(f"{os.path.dirname(os.path.abspath(__file__))}/{DOWNLOAD_FOLDER_NAME}/")

        # Start folder watch daemon to automatically update to server
        start_new_thread(folder_watch_daemon,(current_file_directory,))

        # Print verbose client shell begins
        help()

        # Does Client Things
        while True:

            # Wait for user to input a command
            full_command = input(f'{my_username} > ').strip()
            command = full_command.split(' ')[0]
            parameters = full_command.split(' ')[1:]

            if command == "download":
                if len(parameters) != 0:
                    wait_for_file_download(full_command)
                else:
                    log_this("ParameterError: Too less parameters")

            elif command == "get_files_list":
                if len(parameters) == 0:
                    wait_for_list(full_command)
                else:
                    log_this("ParameterError: Too many parameters")

            elif command == "help":
                help()

            elif command == "quit":
                # Encode command to bytes, prepare header and convert to bytes, like for username above, then send
                full_command = f"unregister".encode('utf-8')
                command_header = f"{len(full_command):<{HEADER_LENGTH}}".encode('utf-8')
                meta = f"{f'{CLIENT_ID}':<{META_LENGTH}}".encode('utf-8')
                central_socket.send(command_header + meta + full_command)

                log_this(f"Client sent command: {full_command}")
                LOG.close()
                sys.exit()

    # Automatic Mode of Client Interface

    #Args
    #python client.py username command1 command2 ... commandn

    else:

        # create username to connect to the server
        my_username = sys.argv[1]
        username = my_username.encode('utf-8')
        username_header = f"{len(username):<{HEADER_LENGTH}}".encode('utf-8')
        meta = f"{'':<{META_LENGTH}}".encode('utf-8')

        # Initialize connection with the server with central socket
        central_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        central_socket.connect((IP, PORT))
        central_socket.setblocking(False)
        central_socket.send(username_header + meta + username)

        # Initialize file directory to server
        update_server()

        # get current file directories
        current_file_directory = os.listdir(f"{os.path.dirname(os.path.abspath(__file__))}/{DOWNLOAD_FOLDER_NAME}/")

        # Start folder watch daemon to automatically update to server
        start_new_thread(folder_watch_daemon,(current_file_directory,))

        # Does Client Things
        for i in sys.argv[2:]:

            # give time to process command
            time.sleep(1)

            # Wait for user to input a command
            full_command = i
            command = full_command.split(' ')[0]
            parameters = full_command.split(' ')[1:]

            if command == "download":
                if len(parameters) != 0:
                    wait_for_file_download(full_command)
                else:
                    log_this("ParameterError: Too less parameters")

            elif command == "get_files_list":
                if len(parameters) == 0:
                    wait_for_list(full_command)
                else:
                    log_this("ParameterError: Too many parameters")

            elif command == "help":
                help()

            elif command == "quit":
                # Encode command to bytes, prepare header and convert to bytes, like for username above, then send
                full_command = f"unregister".encode('utf-8')
                command_header = f"{len(full_command):<{HEADER_LENGTH}}".encode('utf-8')
                meta = f"{f'{CLIENT_ID}':<{META_LENGTH}}".encode('utf-8')
                central_socket.send(command_header + meta + full_command)

                log_this(f"Client sent command: {full_command}")
                LOG.close()
                sys.exit()


import io
import shutil
import tempfile
from http import HTTPStatus
from urllib.parse import urljoin, unquote

from django import forms
from django.conf import settings
from django.test import TestCase, Client, override_settings
from django.shortcuts import reverse
from django.contrib.auth import get_user_model
from django.core.cache import cache
from django.utils.http import urlencode
from django.core.files.uploadedfile import SimpleUploadedFile
from PIL import Image

from posts.models import Post, Group, Comment
from posts.forms import PostForm
from .utils import write_log

User = get_user_model()
# temporary location for media files to be stored at
TEMP_MEDIA_ROOT = tempfile.mkdtemp(dir=settings.BASE_DIR)

class NamedClient(Client):
    def __init__(self, name):
        super().__init__()
        self.name = name

@override_settings(MEDIA_ROOT=TEMP_MEDIA_ROOT)
class PostsFormsTest(TestCase):
    @classmethod
    def setUpClass(cls):
        super().setUpClass()
        cls.new_group = Group.objects.create(
            title='new group',
            slug='new-group',
            description='something descriptive'
        )
        cls.new_group_2 = Group.objects.create(
            title='new group 2',
            slug='new-group-2',
            description='something even more descriptive'
        )
        cls.user_random = User.objects.create(username='wacko wacko')
        cls.user_author = User.objects.create(username='danger-pwrs')
        cls.CREATE = reverse('posts:post_create')
        cls.PROFILE = reverse(
            'posts:profile',
            kwargs={'username': cls.user_author.username}
        )
        cls.EXPECTED_FIELD_TYPES = {
            'text': forms.CharField,
            'group': forms.ModelChoiceField,
            'image': forms.ImageField,
        }
        cls.EXPECTED_HTML_TAGS = {
            'text': (
                '<textarea name="text" cols="40" '
                'rows="10" required id="id_text">'
            ),
            'group': '<option value="" selected> Выберите группу </option>',
            'image': (
                '<input type="file" name="image" '
                'accept="image/*" id="id_image">'
            )
        }

    @classmethod
    def tearDownClass(cls):
        super().tearDownClass()
        cache.clear()
        shutil.rmtree(TEMP_MEDIA_ROOT, ignore_errors=True)

    @staticmethod
    def get_test_img(color: str) -> SimpleUploadedFile:
        image = Image.new('RGB', (10, 10), color)
        byte_arr = io.BytesIO()
        image.save(byte_arr, format='jpeg')
        file = SimpleUploadedFile(
            name=f'{color}.jpeg',
            content=byte_arr.getvalue(),
            content_type='image/jpeg'
        )
        return file

    def get_post_data(self, action):
        post_data = {
            'create': {
                'text': 'Test post 1',
                'group': self.new_group.pk,
                'image': self.get_test_img('red'),
            },
            'update': {
                'text': 'Test post 2',
                'group': self.new_group_2.pk,
                'image': self.get_test_img('blue')
            }
        }
        return post_data[action]

    def create_post(self, test_client):
        data = self.get_post_data('create')
        response = PostsFormsTest.post_to(
            test_client, PostsFormsTest.CREATE, data
        )
        return response

    def update_post(self, test_client, post_id):
        edit_url = reverse('posts:post_edit', args=(post_id,))
        data = self.get_post_data('update')
        response = PostsFormsTest.post_to(
            test_client, edit_url, data
        )
        return response

    def post_to(test_client, to, data):
        response = test_client.post(
            to,
            data,
            follow=True,
            multipart=True,
        )
        write_log(test_client, to, data, TEMP_MEDIA_ROOT, response)
        return response

    def setUp(self):
        self.form = PostForm()
        self.guest_client = NamedClient('Guest client')
        self.auth_client = NamedClient('Authorised client')
        self.auth_client.force_login(self.user_random)
        self.author_client = NamedClient('Author client')
        self.author_client.force_login(self.user_author)

    def test_expected_fields(self):
        """test if the fields are present"""
        for field in ('text', 'group'):
            with self.subTest(field_name=field):
                self.assertIn(field, self.form.fields)

    def test_create_form_is_rendered(self):
        """test if the fields are rendered correctly"""
        for field, html in self.EXPECTED_HTML_TAGS.items():
            with self.subTest(field_name=field):
                self.assertInHTML(html, str(self.form))

    def test_field_types(self):
        """test if the fields are of correct types"""
        for field_name, field_type in self.EXPECTED_FIELD_TYPES.items():
            with self.subTest(field_name=field_name):
                field = self.form.fields.get(field_name)
                self.assertIsInstance(field, field_type)

    def test_count_increment(self):
        """test if post objects count is incremented"""
        post_count = Post.objects.count()
        expected_increment = {
            self.guest_client: post_count,
            self.auth_client: post_count + 1,
        }
        for client, new_count in expected_increment.items():
            with self.subTest(client=client.name):
                self.create_post(client)
                self.assertEqual(
                    Post.objects.count(), new_count,
                )

    def test_create_redirects(self):
        """test if user is redirected to profile page"""
        response = self.create_post(self.author_client)
        self.assertRedirects(response, self.PROFILE)

    def test_edit_form_initial_state(self):
        """test if edit form's fields are correctly pre-populated"""
        self.create_post(self.author_client)
        test_post = Post.objects.latest('pub_date')
        edit_url = reverse(
            'posts:post_edit',
            kwargs={'pk': test_post.id}
        )
        r = self.author_client.get(edit_url)
        form = r.context['form']
        for field, value in form.initial.items():
            with self.subTest(field=field):
                field_name = 'group_id' if field == 'group' else field
                self.assertEqual(test_post.__dict__[field_name], value)

    def test_edit_form_change(self):
        """test if initial post is modified"""
        expected_asserts = {
            self.guest_client: self.assertEqual,
            self.auth_client: self.assertEqual,
            self.author_client: self.assertNotEqual
        }
        for client, func in expected_asserts.items():
            self.create_post(self.author_client)
            test_post = Post.objects.latest('pub_date')
            self.update_post(client, test_post.id)
            edited_post = Post.objects.get(id=test_post.id)
            for field in ('text', 'group', 'image'):
                with self.subTest(client=client.name, field=field):
                    func(
                        getattr(test_post, field), getattr(edited_post, field)
                    )

    def test_edit_redirects(self):
        """test if user is redirected to details page"""
        self.create_post(self.author_client)
        test_post = Post.objects.latest('pub_date')
        response = self.update_post(
            self.author_client, test_post.id
        )
        details_url = reverse('posts:post_detail', args=(test_post.id,))
        self.assertRedirects(response, details_url)

    def test_comment_redirects(self):
        """test if unauthorised user is redirected to login page"""
        self.create_post(self.author_client)
        test_post = Post.objects.latest('pub_date')
        add_comment_url = reverse(
            'posts:add_comment',
            kwargs={'post_id': test_post.id}
        )
        query_string = '?' + unquote(urlencode({'next': add_comment_url}))
        response = self.client.post(
            path=add_comment_url,
            data={'text': 'Hello there'},
            follow=True
        )
        self.assertRedirects(
            response=response,
            expected_url=urljoin(
                reverse(settings.LOGIN_URL),
                query_string
            ),
            status_code=HTTPStatus.FOUND,
            target_status_code=HTTPStatus.OK
        )

    def test_comment_created(self):
        """test if a comment left by an authorised user is created and shown"""
        self.create_post(self.author_client)
        test_post = Post.objects.latest('pub_date')
        test_text = 'Hello there'
        self.auth_client.post(
            path=reverse('posts:add_comment', args=(test_post.id,)),
            data={'text': test_text},
            follow=True
        )
        expected_comment = Comment.objects.latest('pub_date')
        expected_context = self.client.get(
            reverse(
                'posts:post_detail',
                kwargs={'post_id': test_post.id}
            )
        ).context
        self.assertTrue(expected_comment in test_post.comments.all())
        self.assertTrue(expected_comment in expected_context['comments'])
        self.assertEqual(expected_comment.text, test_text)


# Enter your code here. Read input from STDIN. Print output to STDOUT
from itertools import product
from sys import stdin

def getInputs():
    return list(map(int, stdin.readline().strip().split()))

def cartesian_product(A,B):
    result = sorted(list(product(A,B)))
    for pair in result:
        print(pair, end=' ')

# {
# Driver Code starts
if __name__ == "__main__":
    # write your code here
    A = getInputs()
    B = getInputs()
    cartesian_product(A,B)

# } Driver Code ends


import esper
import pygame
from typing import Callable, List

from src.ecs.components.c_surface import CSurface
from src.ecs.components.c_transform import CTransform
from src.ecs.components.tags.c_tag_bullet import CTagBullet
from src.ecs.components.tags.c_tag_enemie import CTagEnemie

def system_collision_bullet_enemy(world:esper.World, explosion_dict:dict, create_explosion:Callable[[esper.World, dict, pygame.Vector2], None], ene_ent:List[int]):
    bullets = world.get_components(CSurface, CTransform, CTagBullet)
    enemies = world.get_components(CSurface, CTransform, CTagEnemie)
    cb_s:CSurface
    cb_t:CTransform
    ce_s:CSurface
    ce_t:CTransform
    points = 0
    for bullet_entity, (cb_s, cb_t, c_b) in bullets:
        bull_rect = CSurface.get_area_relative(cb_s.area, cb_t.pos)
        for enemy_entety, (ce_s, ce_t, c_e) in enemies:
            enemy_rect = CSurface.get_area_relative(ce_s.area, ce_t.pos)
            if bull_rect.colliderect(enemy_rect):
                points += c_e.puntos
                world.delete_entity(bullet_entity)
                world.delete_entity(enemy_entety)
                ene_ent.remove(enemy_entety)
                create_explosion(world, explosion_dict, ce_t.pos)
    return points


import configparser
import os
import sys

from supabase import Client, create_client

config = configparser.ConfigParser()
config.read(os.path.join(os.path.dirname(os.path.dirname(__file__)), 'config.ini'))

parent_module = sys.modules[".".join(__name__.split(".")[:-1]) or "__main__"]
if __name__ == "__main__" or parent_module.__name__ == "__main__":
    from CSVHandler import CSVHandler
    from Skill import Skill
else:
    from .CSVHandler import CSVHandler
    from .Skill import Skill

SUPABASE_URL = config["DATABASE"]["DATABASE_URL"]
SUPABASE_KEY = config["DATABASE"]["DATABASE_API_KEY"]

class DatabaseHandler:
    def __init__(self, database:str = "Skills") -> None:
        self.database = database

    def update_data(self) -> int:
        data = CSVHandler.load_skills_from_csv()
        data = [item.export() for item in data]
        supabase: Client = create_client(SUPABASE_URL, SUPABASE_KEY)

        try:
            return_data = supabase.table(self.database).upsert(data).execute()

            inserted = len(return_data.data)
        except:
            # conn.rollback()
            raise
        print(f"Inserted {inserted} new rows")
        return inserted

    def get_data(self, raw=False) -> list[Skill] or list[tuple]:
        supabase: Client = create_client(SUPABASE_URL, SUPABASE_KEY)

        data = supabase.table(self.database).select("*").execute()

        skills = [tuple(skill.values()) for skill in data.data]
        if raw:
            return skills
        skills = [Skill().import_data(data) for data in skills]
        return skills

    def get_headers(self) -> list[str]:
        supabase: Client = create_client(SUPABASE_URL, SUPABASE_KEY)

        data = supabase.table(self.database).select("*").limit(1).execute()

        headers = [column for column in data.data[0].keys()]

        return headers

if __name__ == "__main__":
    db = DatabaseHandler()

    db.update_data()


# @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


from flask_wtf import FlaskForm
from wtforms import StringField, SubmitField, SelectField, BooleanField
from wtforms.validators import DataRequired, ValidationError

class SubmitBusinessForm(FlaskForm):
    name = StringField("Name", validators=[DataRequired()])
    website = StringField("Website", validators=[DataRequired()])
    category = SelectField("Category", choices=[("Restaurants", "Restaurants"), ("Cafes", "Cafes"), ("Specialty Goods", "Specialty Goods"), ("Apparel", "Apparel"), ("Lifestyle", "Lifestyle")])
    address = StringField("Address", validators=[DataRequired()])
    country = StringField("Country", validators=[DataRequired()])
    submit = SubmitField("Submit")

class FilterForm(FlaskForm):
    restaurants = BooleanField("Restaurants")
    cafes = BooleanField("Cafes")
    specialty_goods = BooleanField("Specialty Goods")
    apparel = BooleanField("Apparel")
    lifestyle = BooleanField("Lifestyle")
    submit = SubmitField("Submit")


from rest_framework.decorators import api_view, permission_classes, authentication_classes
from rest_framework.permissions import IsAuthenticatedOrReadOnly, AllowAny, IsAuthenticated
from rest_framework.authentication import SessionAuthentication
from rest_framework.response import Response
from .models import User
from .serializers import UserSerializer, LoginSerializer
from rest_framework import status
from django.contrib.auth import login
# Create your views here.

@api_view(['POST'])
@permission_classes([AllowAny])
@authentication_classes([SessionAuthentication])
def registerUser(request):
    if request.method == 'POST':
        serializer = UserSerializer(data=request.data)
        if serializer.is_valid():
            serializer.save()
            return Response({'status': 200, 'payload': serializer.data, 'messge': 'your data is saved'})
        return Response(serializer.errors, status=status.HTTP_400_BAD_REQUEST)

@api_view(['POST'])
@permission_classes([AllowAny])
def loginView(request):

    serializer = LoginSerializer(data=request.data, context={ 'request': request })
    serializer.is_valid(raise_exception=True)
    user = serializer.validated_data['user']
    login(request, user)
    return Response(None, status=status.HTTP_202_ACCEPTED)

@api_view(['GET'])
@permission_classes([IsAuthenticatedOrReadOnly])
def getUsers(request):

    if request.method ==  'GET':
        user = User.objects.all()
        serializer = UserSerializer(user, many=True)
        return Response(serializer.data)

@permission_classes([IsAuthenticated])
@api_view(['GET', 'PUT', 'DELETE'])
def getUser(request, pk):

    try:
        user = User.objects.get(pk=pk)
    except User.DoesNotExist:
        return Response(status=status.HTTP_404_NOT_FOUND)

    if request.method == 'GET':
        serializer = UserSerializer(user)
        return Response(serializer.data)

    elif request.method == 'PUT':
        serializer = UserSerializer(user, data=request.data)
        if serializer.is_valid():
            serializer.save()
            return Response(serializer.data, status=status.HTTP_201_CREATED)
        return Response(serializer.data, status=status.HTTP_201_CREATED)

    elif request.method == 'DELETE':
        user.delete()
        return Response(status=status.HTTP_201_CREATED)

# user registration


import src.default_params
import src.sliders
import src.models

import numpy as np

from scipy.integrate import solve_ivp

from bokeh.models import Slider, RangeSlider, ColumnDataSource
from bokeh.io import curdoc, output_file
from bokeh.layouts import column, row
from bokeh.plotting import figure,save
from bokeh.models.widgets import Paragraph
from bokeh.client import push_session

def readme():
    return Paragraph(text="""Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor
                                incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud
                                exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute
                                irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur.
                                Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt
                                mollit anim id est laborum""",
                                width=250, height=100)

def main():

    #retrieve our starting parameters and store them in a dictionary
    params = src.default_params.params
    covid_params = src.default_params.covid_params
    flu_params = src.default_params.influenza_params

    #retrieve our initial conditions, tau parameter and time intervals for coinfection model
    coinfection_inital_conditions, tau, time_interval_1, time_interval_2 = src.default_params.init_params_comb(params)

    #run our coinfection model first, so we have something to
    coinfection_sol_1, coinfection_sol_2 = src.models.run_covid_coinfection_model(src.models.model_covid_influenza,
                                                            time_interval_1, time_interval_2,
                                                            coinfection_inital_conditions, params)

    coin_y = np.hstack([coinfection_sol_1.y, coinfection_sol_2.y])
    coin_t = np.append(coinfection_sol_1.t, coinfection_sol_2.t)

    print(coin_y[2,])
    print(coin_t)

    #put it into a dataframe
    coinfection_src = ColumnDataSource(data=dict(t = coin_t, y1 = coin_y[0,], y2 = coin_y[1,],
                                        y3 = coin_y[2,], y4 = coin_y[3,], y5 = coin_y[4,],
                                        y6 = coin_y[5,], y7 = coin_y[6,], y8 = coin_y[7,]))

    #now do the same for the flu model
    flu_initial_conditions, tint_flu, flu_params_tuple = src.default_params.init_params_flu(params)

    flu_sol = solve_ivp(src.models.model_flu, tint_flu, flu_initial_conditions,
                                atol = [1e-20] * 5, args=flu_params_tuple)

    flu_src = ColumnDataSource(data=dict(t = flu_sol.t, y1 = flu_sol.y[0,],
                               y2 = flu_sol.y[1,], y3 = flu_sol.y[2,],
                               y4 = flu_sol.y[3,], y5 = flu_sol.y[4,]))

    covid_initial_conditions, tint_covid, covid_params_tuple = src.default_params.init_params_covid(params)

    cov_sol = solve_ivp(src.models.model_covid, tint_covid, covid_initial_conditions,
                                atol = [1e-20] * 5, args=covid_params_tuple)

    cov_src = ColumnDataSource(data=dict(t = cov_sol.t, y1 = cov_sol.y[0,],
                               y2 = cov_sol.y[1,], y3 = cov_sol.y[2,],
                               y4 = cov_sol.y[3,], y5 = cov_sol.y[4,]))

    #create our plots
    e_cell_per_ml_vs_time = figure(height=400, width=400, title="E (cells/mL) vs Time (Days)",
                                tools="crosshair,pan,reset,save,wheel_zoom")
    ec_cell_per_ml_vs_time = figure(height=400, width=400, title="E_c (cells/mL) vs Time (Days)",
                                tools="crosshair,pan,reset,save,wheel_zoom")
    c_virions_per_ml_vs_time = figure(height=400, width=400, title="C (virions/mL) vs  Time (Days)",
                                tools="crosshair,pan,reset,save,wheel_zoom")
    ei_cell_per_ml_vs_time = figure(height=400, width=400, title="E_i (cells/mL) vs Time (Days)",
                                tools="crosshair,pan,reset,save,wheel_zoom")
    i_virions_per_ml_vs_time = figure(height=400, width=400, title="I (virions/mL) vs Time (Days)",
                                tools="crosshair,pan,reset,save,wheel_zoom")
    tc_cell_per_ml_vs_time = figure(height=400, width=400, title="T_c (cells/mL) vs Time (Days)",
                                tools="crosshair,pan,reset,save,wheel_zoom")
    ti_cell_per_ml_vs_time = figure(height=400, width=400, title="T_i (cells/mL) vs Time (Days)",
                                tools="crosshair,pan,reset,save,wheel_zoom")
    x_vs_time = figure(height=400, width=400, title="X vs Time (Days)",
                                tools="crosshair,pan,reset,save,wheel_zoom")
    #set up plots
    #todo .line can be layered on top; change colour/dash and it will work

    #plot 1
    e_cell_per_ml_vs_time.line('t', 'y1', source=coinfection_src, line_width=3, line_alpha=0.6,
                               line_color="black")

    #plot 2
    ec_cell_per_ml_vs_time.line('t', 'y2', source=coinfection_src, line_width=3, line_alpha=0.6,
                               line_color="red")

    #plot 3
    c_virions_per_ml_vs_time.line("t", "y3", source=coinfection_src, line_width=3, line_alpha=0.6,
                                   line_color="red")
    #c_virions_per_ml_vs_time.line("t", "y3", source=cov_src, line_width=3, line_alpha=0.6,
    #                               line_color="red", line_dash="dashed")

    #plot 4
    ei_cell_per_ml_vs_time.line("t", "y4", source=coinfection_src, line_width=3, line_alpha=0.6,
                                line_color="blue")

    #plot 5
    i_virions_per_ml_vs_time.line("t", "y5", source=coinfection_src, line_width=3, line_alpha=0.6,
                                  line_color="blue")
    #i_virions_per_ml_vs_time.line("t", "y3", source=cov_src, line_width=3, line_alpha=0.6,
    #                              line_color="blue", line_dash="dashed")

    #plot 6
    tc_cell_per_ml_vs_time.line("t", "y6", source=coinfection_src, line_width=3, line_alpha=0.6,
                                line_color="red")

    #plot 7
    ti_cell_per_ml_vs_time.line("t", "y7", source=coinfection_src, line_width=3, line_alpha=0.6,
                                line_color="blue")

    #plot 8
    x_vs_time.line("t", "y8", source=coinfection_src, line_width=3, line_alpha=0.6, line_color="black")

    #function to update over time
    def update_data(attrname, old, new):
        for parameter in params:
            if parameter in src.sliders.covid_sliders:
                params[parameter] = src.sliders.covid_sliders[parameter].value
            elif parameter in src.sliders.covid_sliders:
                params[parameter] = src.sliders.flu_sliders[parameter].value

        coinfection_sol_1, coinfection_sol_2 = src.models.run_covid_coinfection_model(src.models.model_covid_influenza,
                                                        time_interval_1, time_interval_2, coinfection_inital_conditions, params)

        coin_y = np.hstack([coinfection_sol_1.y, coinfection_sol_2.y])
        coin_t = np.append(coinfection_sol_1.t, coinfection_sol_2.t)

        coinfection_src.data = dict(t = coin_t, y1 = coin_y[0,], y2 = coin_y[1,],
                        y3 = coin_y[2,], y4 = coin_y[3,], y5 = coin_y[4,],
                        y6 = coin_y[5,], y7 = coin_y[6,], y8 = coin_y[7,])

        #now reinit flu model with parameters as they take a different tuple
        flu_initial_conditions, tint_flu, flu_params_tuple = src.default_params.init_params_flu(params)

        flu_sol = solve_ivp(src.models.model_flu, tint_flu, flu_initial_conditions,
                            atol = [1e-20] * 5, args=flu_params_tuple)

        flu_src.data = ColumnDataSource(data=dict(t = flu_sol.t, y1 = flu_sol.y[0,],
                                        y2 = flu_sol.y[1,], y3 = flu_sol.y[2,],
                                        y4 = flu_sol.y[3,], y5 = flu_sol.y[4,]))

        #and for covid
        covid_initial_conditions, tint_covid, covid_params_tuple = src.default_params.init_params_covid(params)

        cov_sol = solve_ivp(src.models.model_covid, tint_covid, covid_initial_conditions,
                                atol = [1e-20] * 5, args=covid_params_tuple)

        cov_src = ColumnDataSource(data=dict(t = cov_sol.t, y1 = cov_sol.y[0,],
                                   y2 = cov_sol.y[1,], y3 = cov_sol.y[2,],
                                   y4 = cov_sol.y[3,], y5 = cov_sol.y[4,]))

    for w in list(src.sliders.covid_sliders.values()) + list(src.sliders.flu_sliders.values()):
        w.on_change('value', update_data)

    inputs = column(list(src.sliders.covid_sliders.values()))

    plot_column_left = column(e_cell_per_ml_vs_time, ec_cell_per_ml_vs_time, c_virions_per_ml_vs_time,
                              ei_cell_per_ml_vs_time)
    plot_column_right = column(i_virions_per_ml_vs_time, tc_cell_per_ml_vs_time, ti_cell_per_ml_vs_time,
                              x_vs_time)

    #make the document
    curdoc().add_root(row(inputs, plot_column_left, plot_column_right))
    curdoc().title = "Coinfection Model of COVID-19 and Influenza"
    output_file('plot.html')

main()