from .QUICPacketParser import parse_packet_bytes, PacketParserError
from .QUICPacket import *
from .QUICConnection import ConnectionContext, create_connection_id
from .QUICEncryption import EncryptionContext
from socket import socket, AF_INET, SOCK_DGRAM, SOL_SOCKET, SO_REUSEADDR
import math
from time import time
import logging

# Initialize logging facility.
def setup_logger(name: str, filepath: str, level=logging.DEBUG):
    file_handler = logging.FileHandler(filepath)
    logger = logging.getLogger(name)
    logger.setLevel(level)
    logger.addHandler(file_handler)
    return logger

log = setup_logger("log", "packet.log")

# Congestion Controller States
SLOW_START = 1
RECOVERY = 2
CONGESTION_AVOIDANCE = 3

# Congestion Controller Data
INFINITY = math.inf
MAX_DATAGRAM_SIZE = 1200
SAFE_DATAGRAM_PAYLOAD_SIZE = 512 # bytes
INITIAL_CONGESTION_WINDOW = MAX_DATAGRAM_SIZE*10 # Initial window is 10 times max datagram size RFC 9002
MINIMUM_CONGESTION_WINDOW = MAX_DATAGRAM_SIZE*2  # Minimum window is 2 times max datagram size RFC 9002

# This means we have ended the connection.
DISCONNECTED = 1
# This means we have completed the handshake and are currently connected.
CONNECTED = 2

# This means we have sent an INITIAL packet and are waiting
# for a response.
INITIALIZING = 3

# This means that we are listening for connections.
LISTENING_INITIAL = 4
LISTENING_HANDSHAKE = 5

# This means that we have closed the connection.
CLOSED = 6

def contains_long_header(pkt: Packet):
    return pkt.header.type in [HT_RETRY, HT_HANDSHAKE, HT_INITIAL]

class PacketSentInfo:

    def __init__(self, in_flight=False, sent_bytes=0, time_sent=0.0, packet_number=0, ack_eliciting=False, packet=None):
        self.in_flight: bool = in_flight
        self.sent_bytes: int = sent_bytes
        self.time_sent: float = time_sent
        self.ack_eliciting: bool = ack_eliciting
        self.packet_number: int = packet_number
        self.packet: Packet = packet

class PacketReceivedInfo:

    def __init__(self, packet_number=0, ack_packet=False):
        self.packet_number: int = packet_number
        self.ack_packet: bool = ack_packet

class QUICPacketizer:
    """
        Manages transmission order of packets:
        1. Manages which packet number to use next. --> This is just an incrementing counter.
        2. Constructs data packets containing stream frames.
    """

    def __init__(self):
        self._next_packet_number = 0

    def get_next_packet_number(self):
        next = self._next_packet_number
        self._next_packet_number +=1
        return next

    def create_header(self, header_type: int, connection_context: ConnectionContext) -> LongHeader or ShortHeader:
        if header_type in [HT_INITIAL, HT_HANDSHAKE, HT_RETRY]:
            return LongHeader(
                    type=header_type,
                    destination_connection_id=connection_context.get_peer_connection_id(),
                    source_connection_id=connection_context.get_local_connection_id(),
                    packet_number=self.get_next_packet_number()
                )
        if header_type in [HT_DATA]:
            return ShortHeader(
                destination_connection_id=connection_context.get_peer_connection_id(),
                packet_number=self.get_next_packet_number()
            )
        return None

    def packetize_retransmissions(self, lost_packets: list[PacketSentInfo]) -> list[Packet]:
        pkts: list[Packet] = []
        for info in lost_packets:
            hdr: LongHeader or ShortHeader = info.packet.header
            hdr.packet_number = self.get_next_packet_number()
            pkts.append(Packet(header=hdr, frames=info.packet.frames))
        return pkts

    def packetize_initial_packet(self, connection_context: ConnectionContext) -> Packet:
        hdr = self.create_header(HT_INITIAL, connection_context)
        frames = [] # TODO Add crypto frames.
        return Packet(header=hdr, frames=frames)

    def packetize_connection_close_packet(self, connection_context: ConnectionContext) -> Packet:
        hdr = self.create_header(HT_DATA, connection_context)
        reason = b"Normal Connection Termination"
        frames = [ConnectionCloseFrame(error_code=1, reason_phrase_len=len(reason), reason_phrase=reason)]
        return Packet(header=hdr, frames=frames)

    def packetize_handshake_packet(self, connection_context: ConnectionContext) -> Packet:
        hdr = self.create_header(HT_HANDSHAKE, connection_context)
        frames = []
        return Packet(header=hdr, frames=frames)

    def packetize_connection_response_packets(self, connection_context: ConnectionContext, encryption_context: EncryptionContext) -> list[Packet]:
        hdr1 = self.create_header(HT_INITIAL, connection_context)
        hdr2 = self.create_header(HT_HANDSHAKE, connection_context)
        frames1 = []
        frames2 = [CryptoFrame(offset=0, length=len(encryption_context.key), data=encryption_context.key)]
        initial, handshake = Packet(header=hdr1, frames=frames1), Packet(header=hdr2, frames=frames2)
        return [initial, handshake]

    def create_ack_frame(self, packet_numbers_received: list[int]) -> AckFrame:
        # If the received packets list is length 1
        # Create a simple ack frame.
        if len(packet_numbers_received) == 1:
            return AckFrame(largest_acknowledged=packet_numbers_received[0], ack_delay=0, ack_range_count=0, ack_range=[])
        # If the received packets list is greater than 1.
        # We sort the packet number list.
        # Then we incrementally create AckRanges.
        # In the resulting AckRanges, an AckRange with gap 0 is the most recent AckRange,
        # from this AckRange we can get first ack range value.
        packet_numbers_received.sort()
        ranges = []
        count = 0
        for i in range(0, len(packet_numbers_received)):
            count += 1
            if i == len(packet_numbers_received)-1:
                ranges.append(AckRange(gap=0, ack_range_length=count))
                break
            if packet_numbers_received[i+1] != packet_numbers_received[i]+1:
                ranges.append(AckRange(gap=packet_numbers_received[i+1] - packet_numbers_received[i] - 1, ack_range_length=count))
                count = 0
        first_ack_range = ranges[-1].ack_range_length - 1
        largest_acknowledged = packet_numbers_received[-1]
        ranges.remove(ranges[-1])
        return AckFrame(largest_acknowledged=largest_acknowledged, ack_delay=0, ack_range_count=len(ranges), first_ack_range=first_ack_range, ack_range=ranges)

    def packetize_acknowledgement(self, connection_context: ConnectionContext, packet_numbers_received: list[int]) -> Packet:

        # If the received packets list is length 0
        # We cannot create an ack if we haven't received any packets.
        if not packet_numbers_received:
            return None

        hdr = self.create_header(HT_DATA, connection_context)
        frames = [self.create_ack_frame(packet_numbers_received)]
        pkt = Packet(header=hdr, frames=frames)
        return pkt

    def packetize_stream_data(self, stream_id: int, data: bytes, connection_context: ConnectionContext, send_streams: dict) -> list[Packet]:
        packets = []
        MAX_ALLOWED = SAFE_DATAGRAM_PAYLOAD_SIZE-LONG_HEADER_SIZE-STREAM_FRAME_SIZE
        # If the length of this data would go over the safe datagram size - long header size - stream frame size,
        # then we need to split this data into multiple QUIC packets.
        datasize = len(data)
        if datasize > MAX_ALLOWED:
            bytes_written = 0
            while bytes_written < datasize:
                data_chunk = data[0:MAX_ALLOWED]
                hdr = self.create_header(HT_DATA, connection_context)
                frames = [StreamFrame(stream_id=stream_id, offset=send_streams[stream_id].get_offset(), length=len(data_chunk), data=data_chunk)]
                packets.append(Packet(header=hdr, frames=frames))
                send_streams[stream_id].update_offset(len(data_chunk))
                data = data[MAX_ALLOWED:]
                bytes_written += len(data_chunk)
        else:
            hdr = self.create_header(HT_DATA, connection_context)
            frames = [StreamFrame(stream_id=stream_id, offset=send_streams[stream_id].get_offset(), length=len(data), data=data)]
            packets.append(Packet(header=hdr, frames=frames))
            send_streams[stream_id].update_offset(len(data))
        return packets

class SendStream:

    def __init__(self, stream_id: int):
        self.stream_id = stream_id
        self.offset = 0

    def get_offset(self) -> int:
        return self.offset

    def update_offset(self, data_length: int) -> None:
        self.offset += data_length

class ReceiveStream:

    def __init__(self, stream_id: int):
        self.stream_id = stream_id
        self.data = b""
        self.offset = 0
        self.buffered_frames: list[StreamFrame] = []

    def buffer(self, frame: StreamFrame):
        self.buffered_frames.append(frame)

    def write(self, new_data: bytes):
        self.offset += len(new_data)
        self.data += new_data
        self.process_buffered_frames()

    def process_buffered_frames(self):
        frames_to_remove = []
        frames = ""
        for frame in self.buffered_frames:
            frames += frame.__repr__() + '\n'
        for frame in self.buffered_frames:
            if self.offset == frame.offset:
                self.data += frame.data
                self.offset += len(frame.data)
                frames_to_remove.append(frame)
        for frame in frames_to_remove:
            self.buffered_frames.remove(frame)

    def read(self, num_bytes: int) -> bytes:
        data = self.data[0:num_bytes]
        self.data = self.data[len(data):]
        return data

class QUICNetworkController:
    """
        TODO: docstring
    """

    def __init__(self):

        self._connection_context: ConnectionContext = ConnectionContext()
        self._encryption_context: EncryptionContext = None # This gets set when a connection is made.
        self.temp_encryption_context = None
        self._sender_side_controller = QUICSenderSideController()
        self._packetizer = QUICPacketizer()
        self._receive_streams = dict() # Key: Stream ID (int) | Value: Stream object
        self._send_streams = dict()
        self.buffered_packets = []
        self.new_socket = None
        self.peer_issued_connection_closed = False

        # ---- Handshake Data ----
        self.handshake_complete = False
        self.server_initial_received = False
        self.server_handshake_received = False
        # self.client_initial_received = False
        # self.client_handshake_received = False
        self.state = DISCONNECTED
        self.last_peer_address_received = None

        # ---- Acknowledgement Data ----
        self.largest_acknowledged = -1
        self.largest_packet_number_received: int = 0
        self.unacked_packet_numbers_received: list[int] = []

    def initiate_connection_termination(self, udp_socket: socket) -> None:
        """
            1. Send short header packet with ConnectionClose Frame.
                1. Use the packetizer to create a ConnectionClose packet.
                2. transmit the packet using the regular functions.
                3. Close the underlying UDP socket.
                4. Update the connection context information.
        """
        connection_close_packet = self._packetizer.packetize_connection_close_packet(self._connection_context)
        self.send_packets(udp_socket=udp_socket, packets=[connection_close_packet])
        udp_socket.close()
        self._connection_context.connected = False
        self.state = CLOSED

    def respond_to_connection_termination(self, udp_socket: socket):
        udp_socket.close()
        self._connection_context.connected = False
        self.state = CLOSED

    def is_client_handshake_complete(self) -> bool:
        return self.server_handshake_received and self.server_initial_received

    # def is_server_handshake_complete(self) -> bool:
    #     return self.client_handshake_received and self.client_initial_received

    def set_buffered_packets(self, buffered_packets: list):
        self.buffered_packets = buffered_packets

    def get_buffered_packets(self) -> list:
        return self.buffered_packets

    def set_packetizer(self, packetizer: QUICPacketizer):
        self._packetizer = packetizer

    def get_packetizer(self) -> QUICPacketizer:
        return self._packetizer

    def set_state(self, state: int) -> None:
        self.state = state

    def get_state(self) -> int:
        return self.state

    def set_receive_streams(self, recv_streams: dict):
        self._receive_streams = recv_streams

    def get_receive_streams(self) -> dict:
        return self._receive_streams

    def set_send_streams(self, send_streams: dict):
        self._send_streams = send_streams

    def get_send_streams(self) -> dict:
        return self._send_streams

    def set_encryption_context(self, encryption_context: EncryptionContext) -> None:
        self._encryption_context = encryption_context

    def get_encryption_context(self) -> EncryptionContext | None:
        return self._encryption_context

    def set_connection_context(self, connection_context: ConnectionContext) -> None:
        self._connection_context = connection_context

    def get_connection_context(self) -> ConnectionContext | None:
        return self._connection_context

    def create_stream(self, stream_id: int) -> None:
        self._receive_streams[stream_id] = ReceiveStream(stream_id=stream_id)
        self._send_streams[stream_id] = SendStream(stream_id=stream_id)

    def listen(self, udp_socket: socket):
        # When a socket is listening, we need to bind it to the wildcard
        # address so that it doesn't associate itself with incoming connections.
        self.state = LISTENING_INITIAL
        udp_socket.bind(("", self._connection_context.get_local_port()))

    def create_connection(self, udp_socket: socket, server_address: tuple[str, int]):

        if self.state != DISCONNECTED:
            print("Socket must be DISCONNECTED to create a connection.")
            exit(1)

        # ---- UPDATE 5-TUPLE ----
        udp_socket.connect(server_address)

        # ---- INITIALIZE CONNECTION CONTEXT ----
        self._connection_context.set_peer_address(server_address)
        self._connection_context.set_peer_connection_id(create_connection_id())

        # ---- PACKETIZE INITIAL PACKET ----
        initial = self._packetizer.packetize_initial_packet(self._connection_context)
        self.send_packets([initial], udp_socket)
        self.state = INITIALIZING

        # ---- PROCESSING RESPONSE ----
        while not self.is_client_handshake_complete():
            packets = self.receive_new_packets(udp_socket, self._encryption_context)
            self.process_packets(packets, udp_socket)

        # ---- Connection Complete ----
        self.state = CONNECTED
        self.create_stream(1)

    def accept_connection(self, udp_socket: socket) -> ConnectionContext:
        if self.state != LISTENING_INITIAL:
            print("Must be in LISTENING state to accept()")
            exit(1)
        while not self.handshake_complete:
            if self.new_socket:
                packets = self.receive_new_packets(self.new_socket, self._encryption_context)
                self.process_packets(packets, self.new_socket)
            else:
                packets = self.receive_new_packets(udp_socket, self._encryption_context)
                self.process_packets(packets, udp_socket)
        self.state = CONNECTED
        return self

    def send_stream_data(self, stream_id: int, data: bytes, udp_socket: socket) -> bool:
        # Check for new packets to process and process them.
        packets_to_process = self.receive_new_packets(udp_socket, self._encryption_context)
        self.process_packets(packets_to_process, udp_socket)

        # If the connection has been closed, we return -1.
        if self.peer_issued_connection_closed:
            return False

        # Packetize the stream data.
        packets: list[Packet] = self._packetizer.packetize_stream_data(stream_id, data, self._connection_context, self._send_streams)
        # Add the currently buffered packets to the list of packets to send.

        could_not_send: list[Packet] = self.send_packets(packets, udp_socket)
        while could_not_send:
            # Reprocess packets that could not be sent.
            packets_to_process = self.receive_new_packets(udp_socket, self._encryption_context)
            self.process_packets(packets_to_process, udp_socket)
            could_not_send = self.send_packets(could_not_send, udp_socket)
        return True

    def send_packets(self, packets: list[Packet], udp_socket: socket) -> list[Packet]:
        could_not_send: list[Packet] = []
        for packet in packets:
            log.debug(f"Sent: \n{packet}")
            if self.is_ack_eliciting(packet):
                # If the packet is ack eliciting,
                # then send it with congestion control.
                if self._sender_side_controller.can_send():
                    # bytes in flight < congestion window
                    try:
                        self._sender_side_controller.send_packet_cc(packet, udp_socket, self._connection_context, self._encryption_context)
                    except ConnectionRefusedError:
                        pass
                else:
                    # bytes in flight >= congestion window
                    # Need to wait to receive more acks before continuing to send.
                    could_not_send.append(packet)
            else:
                # This is an Ack, Padding, or ConnectionClose packet,
                try:
                    self._sender_side_controller.send_non_ack_eliciting_packet(packet, udp_socket, self._connection_context, self._encryption_context)
                except ConnectionRefusedError:
                    pass
        return could_not_send

    def read_stream_data(self, stream_id: int, num_bytes: int, udp_socket: socket) -> tuple[bytes, bool]:
        """
        """
        # Receive and process new packets.
        packets: list[Packet] = self.receive_new_packets(udp_socket, self._encryption_context)
        self.process_packets(packets, udp_socket)
        # Now we can read from the receive_stream.
        data: bytes = self._receive_streams[stream_id].read(num_bytes)
        return data, self.peer_issued_connection_closed

    def is_active_stream(self, stream_id: int) -> bool:
        return stream_id in self._receive_streams

    def is_ack_eliciting(self, packet: Packet) -> bool:
        for frame in packet.frames:
            if frame.type not in [FT_ACK, FT_PADDING, FT_CONNECTIONCLOSE]:
                return True
        return False

    def create_and_send_acknowledgements(self, udp_socket: socket) -> None:
        ack_pkt: Packet = self._packetizer.packetize_acknowledgement(self._connection_context, self.unacked_packet_numbers_received)
        self.send_packets([ack_pkt], udp_socket)

    def update_largest_packet_number_received(self, packet: Packet) -> None:
        self.largest_packet_number_received = max(packet.header.packet_number, self.largest_packet_number_received)

    def update_received_packet_numbers(self, pkt_num: int) -> None:
        if pkt_num not in self.unacked_packet_numbers_received:
            self.unacked_packet_numbers_received.append(pkt_num)

    def process_short_header_packet(self, packet: Packet, udp_socket: socket) -> None:
        # Processing Short Header Packet Frames:
        # 1. Stream Frame --> Write to stream or buffer data.
        # 2. Ack Frame    --> Remove from packets_sent, decrement bytes_in_flight, other congestion control stuff.
        for frame in packet.frames:
            if frame.type == FT_STREAM:
                self.on_stream_frame_received(frame)
            if frame.type == FT_ACK:
                self.on_ack_frame_received(frame, udp_socket)
            if frame.type == FT_CONNECTIONCLOSE:
                self.peer_issued_connection_closed = True
            # TODO: Add checks for other frame types i.e. StreamClose, ConnectionClose, etc.

    def process_long_header_packet(self, packet: Packet, udp_socket: socket) -> None:

        # Client has sent the HT_INITIAL packet to the server.
        # Client is waiting for the HT_INITIAL and HT_HANDSHAKE response.
        if self.get_state() == INITIALIZING:
            if packet.header.type == HT_INITIAL:
                self.server_initial_received = True
                self._connection_context.set_peer_address(self.last_peer_address_received) # TODO This doesn't need to be here
                self._connection_context.set_local_connection_id(packet.header.destination_connection_id)
            if packet.header.type == HT_HANDSHAKE:
                self.server_handshake_received = True
                if self.server_initial_received:
                    response = self._packetizer.packetize_handshake_packet(self._connection_context)
                    self.send_packets([response], udp_socket)
                    self._encryption_context = EncryptionContext(key=packet.frames[0].data)
                    self.state = CONNECTED
                else:
                    self.buffered_packets.append(packet)
                    self.state = INITIALIZING
            return

        # Server is listening for INITIAL packets from clients.
        if self.get_state() == LISTENING_INITIAL:
            # When we are listening for INITIAL packets,
            # we only care about INITIAL packets so buffer all other types.
            if packet.header.type == HT_INITIAL:
                self._connection_context.set_peer_address(self.last_peer_address_received)
                self._connection_context.set_local_connection_id(packet.header.destination_connection_id)
                self._connection_context.set_peer_connection_id(create_connection_id())
                self.new_socket = socket(AF_INET, SOCK_DGRAM)
                self.new_socket.setsockopt(SOL_SOCKET, SO_REUSEADDR, 1)
                self.new_socket.bind(self._connection_context.get_local_address())
                self.new_socket.connect(self._connection_context.get_peer_address())
                self.temp_encryption_context = EncryptionContext()
                packets = self._packetizer.packetize_connection_response_packets(self._connection_context, self.temp_encryption_context)
                self.send_packets(packets, self.new_socket)
                # self.client_initial_received = True
                self.state = LISTENING_HANDSHAKE
                return
            self.buffered_packets.append(packet)
            return

        # Server has received the INITIAL packet from the client and has sent a response.
        # When we receive a client handshake, it finalizes the handshake procedure.
        # We must buffer all other packets at this point.
        if self.get_state() == LISTENING_HANDSHAKE:
            if packet.header.type == HT_HANDSHAKE:
                # self.client_handshake_received = True
                self._encryption_context = self.temp_encryption_context
                self.temp_encryption_context = None
                self.handshake_complete = True
                return
            self.buffered_packets.append(packet)
            return

    def process_packets(self, packets: list[Packet], udp_socket: socket) -> None:

        if not packets:
            return

        # We should process long header packets first.
        lh_packets = [packet for packet in packets if packet.header.type in [HT_INITIAL, HT_HANDSHAKE, HT_RETRY]]
        sh_packets = [packet for packet in packets if packet.header.type in [HT_DATA]]

        for packet in lh_packets:
            self.process_long_header_packet(packet, udp_socket)
        if self.state == CONNECTED:
            for packet in sh_packets:
                self.process_short_header_packet(packet, udp_socket)
        else:
            self.buffered_packets += sh_packets
        for packet in packets:
            self.update_largest_packet_number_received(packet)
            self.update_received_packet_numbers(packet.header.packet_number)
            if self.state == CONNECTED:
                if self.is_ack_eliciting(packet):
                    pkt = self._packetizer.packetize_acknowledgement(self._connection_context, self.unacked_packet_numbers_received)
                    self.send_packets([pkt], udp_socket)

    def receive_new_packets(self, udp_socket: socket, encryption_context: EncryptionContext or None, block=False):
        packets: list[Packet] = [] + self.buffered_packets
        self.buffered_packets = []
        datagrams: list[bytes] = []
        udp_socket.setblocking(block)
        while True:
            try:
                datagram, address = udp_socket.recvfrom(4096)
                datagrams.append(datagram)
            except BlockingIOError:
                break
            except ConnectionRefusedError:
                break
        for datagram in datagrams:
            try:
                # if encryption_context:
                #     packet = parse_packet_bytes(encryption_context.decrypt(datagram))
                # else:
                packet = parse_packet_bytes(datagram)
            except PacketParserError:
                continue # If a datagram fails to be parsed, just drop it.
            if packet.header.type == HT_INITIAL:
                self.last_peer_address_received = address
            log.debug(f"Received: \n{packet}")
            packets.append(packet)
        return packets

    def on_stream_frame_received(self, frame: StreamFrame):
        if not self.is_active_stream(frame.stream_id):
            self.create_stream(1)

        if self.is_active_stream(frame.stream_id):
            stream = self._receive_streams[frame.stream_id]
            if frame.offset != stream.offset:
                stream.buffer(frame)
            else:
                stream.write(frame.data)
            self._receive_streams[frame.stream_id] = stream
        else:
            print(f"Stream ID {frame.stream_id} does not exist.")
            exit(1)

    def extract_ack_frame(self, pkt: Packet) -> bool:
        for frame in pkt.frames:
            if frame.type == FT_ACK:
                return frame
        return None

    def remove_from_packets_received(self, packets_acked: list[PacketSentInfo]) -> None:
        # This function iterates through the packet nums acknowledged,
        # and checks whether any of the packets that were acknowledged are
        # ack packets, if they are ack packets, then we remove from our list of received packets,
        # so that we don't double acknowledge packets.
        for info in packets_acked:
            frame: AckFrame = self.extract_ack_frame(info.packet)
            if frame:
                # If ack ack frame was acked, we can remove the packet numbers it was acking from out received list.
                pkt_nums_acked = [i for i in range(frame.largest_acknowledged, frame.largest_acknowledged - frame.first_ack_range - 1, -1)]
                for pn in pkt_nums_acked:
                    try:
                        self.unacked_packet_numbers_received.remove(pn)
                    except ValueError:
                        pass

    def on_ack_frame_received(self, frame: AckFrame, udp_socket: socket):

        # Calculate packet numbers being acked.
        pkt_nums_acknowledged = [i for i in range(frame.largest_acknowledged, frame.largest_acknowledged-frame.first_ack_range-1, -1)]
        if frame.ack_range_count > 0:
            pn = frame.largest_acknowledged - frame.first_ack_range - 1
            for ackrange in frame.ack_range:
                pn -= ackrange.gap
                pkt_nums_acknowledged += [i for i in range(pn, pn-ackrange.ack_range_length-1, -1)]

        packets_acked = self._sender_side_controller.on_packet_numbers_acked(pkt_nums_acknowledged)
        self.largest_acknowledged = max(self.largest_acknowledged, max(pkt_nums_acknowledged))
        self.remove_from_packets_received(packets_acked)

        # Detect and handle packet loss.
        lost_packets = self._sender_side_controller.detect_and_remove_lost_packets(self.largest_acknowledged)
        if lost_packets: # Packet loss detected.
            retransmissions = self._packetizer.packetize_retransmissions(lost_packets) # Creates new packets
            self.send_packets(retransmissions, udp_socket) # Retransmits packets.
        # If there is no loss, then continue as normal.

class QUICSenderSideController:
    """
        This is the sender side congestion controller.
    """

    def __init__(self):
        self.congestion_window: int = INITIAL_CONGESTION_WINDOW
        self.bytes_in_flight = 0
        self.slow_start_threshold: float = INFINITY
        self.packets_sent: dict[int, PacketSentInfo] = dict()
        self.congestion_recovery_start_time = 0
        self.sent_time_of_last_loss = 0

    def on_packet_loss(self):
        if self.in_recovery(self.sent_time_of_last_loss):
            return
        self.slow_start_threshold = self.congestion_window / 2
        self.congestion_window = max(self.slow_start_threshold, MINIMUM_CONGESTION_WINDOW)
        self.congestion_recovery_start_time = time()

    def detect_and_remove_lost_packets(self, largest_acknowledged: int) -> list[PacketSentInfo]:
        # Detecting Loss:
        # A packet is deemed lost if it if ack-eliciting, in-flight, and was sent prior to an acknowledged packet.
        # AND
        # The packet was sent K packets before an acknowledged packet.
        lost_packets: list[PacketSentInfo] = []
        for pkt_num in self.packets_sent:
            packet_is_lost = pkt_num < largest_acknowledged and self.packets_sent[pkt_num].ack_eliciting and self.packets_sent[pkt_num].in_flight and (largest_acknowledged - pkt_num) >= 3
            if packet_is_lost:                               # If the packet is deemed lost.
                lost_packets.append(self.packets_sent[pkt_num])      # Add to lost packets list.
        if lost_packets:
            for info in lost_packets:
                self.packets_sent.pop(info.packet_number)
                self.bytes_in_flight -= info.sent_bytes
                self.sent_time_of_last_loss = max(self.sent_time_of_last_loss, info.time_sent)
            if self.sent_time_of_last_loss != 0:
                self.on_packet_loss()
        return lost_packets

    def on_packet_numbers_acked(self, packet_numbers: list[int]) -> None:

        # We only want to process packet numbers that exist in our packets_sent list.
        packet_numbers = [x for x in packet_numbers if x in self.packets_sent]
        packets_acked = []

        for x in packet_numbers:
            if not self.packets_sent[x].in_flight:
                # packets that aren't in flight don't count toward cwnd or bytes_in_flight.
                packets_acked.append(self.packets_sent.pop(x))
                continue
            self.bytes_in_flight -= self.packets_sent[x].sent_bytes
            if self.in_recovery(self.packets_sent[x].time_sent):
                # recovery state, don't increase congestion window.
                packets_acked.append(self.packets_sent.pop(x))
                continue
            if self.in_slow_start():
                # slow start
                # increase congestion window by bytes acked.
                self.congestion_window += self.packets_sent[x].sent_bytes
            else:
                # congestion avoidance
                # Additive increase, multiplicitive decrease
                self.congestion_window += MAX_DATAGRAM_SIZE * self.packets_sent[x].sent_bytes / self.congestion_window
            self.congestion_recovery_start_time = 0
            self.sent_time_of_last_loss = 0
            packets_acked.append(self.packets_sent.pop(x))
        return packets_acked

    def in_recovery(self, time_last_loss: float) -> bool:
        return time_last_loss <= self.congestion_recovery_start_time

    def send_packet_cc(self, packet: Packet, udp_socket: socket, connection_context: ConnectionContext, encryption_context: EncryptionContext or None) -> None:
        # Send packets based on the internal congestion control state.
        # if encryption_context:
        #     udp_socket.sendto(encryption_context.encrypt(packet.raw()), connection_context.get_peer_address())
        # else:
        udp_socket.sendto(packet.raw(), connection_context.get_peer_address())
        self.bytes_in_flight += len(packet.raw())
        self.packets_sent[packet.header.packet_number] = PacketSentInfo(time_sent=time(),
                                                                    in_flight=True,
                                                                    ack_eliciting=True,
                                                                    sent_bytes=len(packet.raw()),
                                                                    packet_number=packet.header.packet_number,
                                                                    packet=packet)

    def send_non_ack_eliciting_packet(self, packet: Packet, udp_socket: socket, connection_context: ConnectionContext, encryption_context: EncryptionContext or None) -> None:
        # For non-ack eliciting packets we don't care about congestion control state.
        # if encryption_context:
        #     udp_socket.sendto(encryption_context.encrypt(packet.raw()), connection_context.get_peer_address())
        # else:
        udp_socket.sendto(packet.raw(), connection_context.get_peer_address())
        self.packets_sent[packet.header.packet_number] = PacketSentInfo(time_sent=time(),
                                                                    in_flight=False,
                                                                    ack_eliciting=False,
                                                                    sent_bytes=len(packet.raw()),
                                                                    packet_number=packet.header.packet_number,
                                                                    packet=packet)

    def can_send(self):
        return self.bytes_in_flight < self.congestion_window

    def in_slow_start(self) -> bool:
        return self.congestion_window < self.slow_start_threshold


import discord
from discord.ext import commands
import subprocess
import sys

class dev(commands.Cog):

    def __init__(self, client):
        self.client = client

    @commands.command(aliases=['pull'])
    async def git_pull_command(self, ctx):
        if ctx.author.id == 723377619420184668:
            sys.stdout.flush()
            p = subprocess.Popen(['git', 'pull'], stdout=subprocess.PIPE)
            for line in iter(p.stdout.readline, ''):
                if not line:
                    break
                await ctx.channel.send(str(line.rstrip(), 'utf-8', 'ignore'))
            sys.stdout.flush()
        else:
            await ctx.channel.send("This is a developer-only command. You can't use it")

    @commands.command(aliases=['restart'])
    async def restart_command(self, ctx):
        if ctx.author.id == 723377619420184668:
            await self.git_pull_command(ctx)
            p = subprocess.Popen(['python3', 'start.py'])
            sys.exit(0)
        else:
            await ctx.channel.send("This is a developer-only command. You can't use it")

def setup(client):
    client.add_cog(dev(client))


#!/usr/bin/python3

import apt
import subprocess
import argparse

PREDEFINED = {
        "python-wxtools" : None,
        "python-rosdep":"python3-rosdep2",
        "python3-rosdep":"python3-rosdep2",
        "python3-rosdep-modules":"python3-rosdep2",
        "python3-catkin-pkg-modules":"python3-catkin-pkg"
        }

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--path", "-p", default=None, help="path to the src folder")
    args = parser.parse_args()

    src = args.path if args.path is not None else "src"
    rosdep_path = os.path.join(os.path.expanduser("~"), ".local/bin/rosdep")
    if not os.path.exists(rosdep_path):
        rosdep_path = "rosdep"
    try:
        a = subprocess.check_output([rosdep_path,"check","--from-paths", src,"--ignore-src"])
    except subprocess.CalledProcessError as e:
        err, deps = e.stderr, e.stdout

    if err is not None:
        print("Error", err.decode())
        exit()
    elif deps is not None:
        print("Dependencies not satisfied")
        deps = deps.decode().split('\n')
        deps = deps[1:]
        cache = apt.cache.Cache()
        cache.update()
        cache.open()
        for dep in deps:
            dep = dep.replace('apt\t', "", 1)
            dep_f = PREDEFINED[dep] if dep in PREDEFINED else dep
            if dep_f in [None, ""]:
                print("-------------- '{}' not installed".format(dep))
                continue
            if 'python3' in dep_f or 'python' not in dep_f:
                print("Installing ", dep)
                cache[dep_f].mark_install()
                continue
            _dep_f = dep_f.replace('python-', 'python3-', 1)
            pkg = cache[_dep_f]
            if pkg.is_installed:
                continue
            confirm = input("Install {} instead of {}? (Y/n)".format(_dep_f, dep_f))
            if confirm in ['n','N']:
                continue
            pkg.mark_install()

        try:
            cache.commit()
        except Exception as e:
            print(e)

    print("Nothing to install")


from api.fmg_api import FmgAPI
from stock_history_reader import StockHistoryReader

if __name__ == '__main__':
    history = StockHistoryReader.read_history()
    ticker = "AAPL"
    fb_history = history.get_history(ticker)
    for h in fb_history:
        print(h.date,h.ev_to_ebitda)

    # print("---")
    # current_ev_over_ebitda = FmgAPI.get_key_metrics(ticker)['enterpriseValueOverEBITDATTM']
    # print(current_ev_over_ebitda)
    # data = requests.get(
    #     f"https://financialmodelingprep.com/api/v3/key-metrics-ttm/ABCB?apikey={API_KEY}")
    # data = data.json()
    # print(data[0]['enterpriseValueOverEBITDATTM'])


from typing import List, Tuple

import tcod.path

from game_message import Tick, Action, Spawn, Sail, Dock, Anchor, directions, Position
import numpy as np
import math

def next_path_to_port(tick, position: Tuple):

    if tick.currentLocation.column > position[0]:
        return "N"

    elif tick.currentLocation.column < position[0]:
        return "S"

    if tick.currentLocation.row > position[1]:
        return "W"
    elif tick.currentLocation.row < position[1]:
        return "E"

def create_graph(tick: Tick):
    graph = np.array(tick.map.topology)
    for indexl, ligne in tick.map.topology:
        for indexc, colonne in ligne:
            distance = abs(tick.currentLocation.row - indexl) + abs(tick.currentLocation.column - indexc)
            if tick.tideSchedule[distance % len(tick.tideSchedule)] > colonne:
                graph[indexl][indexc] = 1
            else:
                graph[indexl][indexc] = 0
    return graph.tolist()

class Bot:
    def __init__(self):
        self.graph = None
        self.ports = []
        self.path = None
        print("Initializing your super mega duper bot")

    def get_next_move(self, tick: Tick) -> Action:
        """
        Here is where the magic happens, for now the move is random. I bet you can do better ;)
        """
        if tick.currentLocation is None:
            print(tick.tideSchedule)
            self.ports.append(tick.map.ports[0])
            return Spawn(tick.map.ports[0])

        """
        vérifier si on a atteint le port puis recréer le graph et le chemin
        """
        self.graph = tcod.path.AStar(create_graph(tick), 0)
        nearest_port = tick.map.ports[self.find_nearest_port(tick)]
        self.path = self.graph.get_path(tick.currentLocation.row, tick.currentLocation.column, nearest_port.row,
                                        nearest_port.column)
        next_move = next_path_to_port(tick, self.path[0])
        self.path = self.path[1:]

        if tick.currentLocation.row == nearest_port.row and tick.currentLocation.column == nearest_port.column:
            self.ports.append(nearest_port)
            return Dock()
        return Sail(next_move)

    def find_nearest_port(self, tick) -> int:
        min_distance = math.inf
        index_port = 0
        for port in tick.map.ports:
            if not self.is_in_ports(port.row, port.column):
                distance = abs(tick.currentLocation.row - port.row) + abs(tick.currentLocation.column - port.column)
                if distance < min_distance:
                    min_distance = distance
                    index_port = tick.map.ports.index(port)
        return index_port

    def is_in_ports(self, row, column):
        for port in self.ports:
            if port.row == row and port.column == column:
                return True
        return False


from datetime import datetime

from django.forms import ModelForm

from .models import Page

class PageBaseForm(ModelForm):
    class Meta:
        model = Page
        fields = "__all__"

    def clean_publish_from(self):
        publish_from_cleaned = self.cleaned_data["publish_from"]
        if not publish_from_cleaned:
            return None
        return self.cleaned_data["publish_from"]

    def clean_publish_to(self):
        publish_to_cleaned = self.cleaned_data["publish_to"]
        if not publish_to_cleaned:
            return None
        return self.cleaned_data["publish_to"]

class PageAddForm(PageBaseForm):
    class Meta(PageBaseForm.Meta):
        fields = None
        exclude = (
            "status",
            "created",
            "modified",
            "page_structure",
            "author",
        )

class PageEditForm(PageBaseForm):
    class Meta(PageBaseForm.Meta):
        exclude = (
            "author",
            "created",
            "modified",
            "page_structure",
            "status",
        )


import os
from typing import List

import librosa
import pandas as pd
import numpy as np
from pandas.core.frame import DataFrame
import torchaudio

def get_wav_files(path: str):
    """
    get list of all wav files in diretory

    Args:
        path (str): Path to directory

    Returns:
        (List[str]): List of wav filenames
    """
    filenames = []
    for subdir, dirs, files in os.walk(path):
        for file in files:
            if file.endswith("wav"):
                filenames.append(os.path.join(subdir, file))
    return filenames

def build_dataframe_for_classification(td_path: str, ssd_path: str):
    """
    Build DataFrame for Speaker Classification Problem
    """
    ssd_filenames = get_wav_files(ssd_path)
    td_filenames = get_wav_files(td_path)
    df = pd.DataFrame()
    df["file_path"] = ssd_filenames + td_filenames
    df["labels"] = ["disordered" for _ in range(len(ssd_filenames))] + [
        "typical" for _ in range(len(td_filenames))
    ]
    return df

def clean_dataframe(df: pd.DataFrame, min_duration: float = 0.5) -> pd.DataFrame:
    """
    Perfrom preprocessing steps on audio dataframe
    1. Remove broken file paths
    2. Shuffle data
    3.

    Args:
        df (pd.DataFrame): Input dataframe
        min_duration (float): Minimum duration of audio to keep in seconds. Defaults to 0.5.

    Returns:
        pd.DataFrame: Cleaned dataframe
    """
    df = filter_broken_links(df)
    df = shuffle_df(df)
    return df

def shuffle_df(df: pd.DataFrame) -> pd.DataFrame:
    """
    Shuffle a dataframe

    Args:
        df (pd.DataFrame): Input Dataframe

    Returns:
        pd.DataFrame: Output Dataframe
    """
    df = df.sample(frac=1)
    df = df.reset_index(drop=True)
    return df

def filter_broken_links(df: pd.DataFrame) -> pd.DataFrame:
    """
    Remove broken file paths

    Args:
        df (pd.DataFrame): Input df

    Returns:
        pd.DataFrame: Output df
    """
    df["status"] = df["file_path"].apply(
        lambda path: True if os.path.exists(path) else None
    )
    df = df.dropna(subset=["file_path"])
    df = df.drop("status", 1)
    # df = df.reset_index(drop=True)
    return df

def filter_duration(
    df: pd.DataFrame, min_duration: float = 0.5, sample_rate: int = 16000
) -> pd.DataFrame:
    """
    Remove audio samples under a certain duration

    Args:
        df (pd.DataFrame): Input dataframe
        min_duration (float): Minimum duration of audio to keep in seconds. Defaults to 0.5.

    Returns:
        pd.DataFrame: Filtered dataframe
    """
    print("Filtering short utterances")
    df["duration"] = df["file_path"].apply(
        lambda path: len(librosa.load(path, sr=sample_rate)[0]) / sample_rate
    )
    print(df["duration"])
    df = df.loc[df["duration"] > min_duration]
    print(len(df))
    # df = df.reset_index(drop=True)
    return df

def speech_file_to_array(path: str, target_sampling_rate: int = 16000) -> np.array:
    """
    Convert audio file to numpy array representation

    Args:
        path (str): Path to file
        target_sampling_rate (int, optional): [description]. Defaults to 16000.

    Returns:
        np.array: Speech vector
    """
    speech_array, sampling_rate = torchaudio.load(path)
    resampler = torchaudio.transforms.Resample(sampling_rate, target_sampling_rate)
    speech = resampler(speech_array).squeeze().numpy()
    return speech

def label_to_id(label, label_list):

    if len(label_list) > 0:
        return label_list.index(label) if label in label_list else -1

    return label


import os

from .helpers.lists import distinct

class StaticManager(object):
    def __init__(self, static_path, libs_requirements, packages):
        self.static_path = static_path
        self.libs = {'js': {}, 'css': {}}

        # gather static libs with requirements
        for filename, requirements in libs_requirements.items():
            name, ext = self.parse_filename(filename)
            path = os.path.join(static_path, 'libs' \
                if ext == 'js' else 'styles')
            self.append_lib(path, name, ext, requirements)

        # gather static libs without requirements
        self.find_libs('libs')
        self.find_libs('styles')

        # gather packages static
        for package in packages:
            self.find_libs(package.name, True)

    def append_lib(self, path, name, ext, require=(), package=None):
        lib = StaticLib(self.static_path, path, name, ext, require, package)
        if not lib.full_name in self.libs[ext]:
            self.libs[ext][lib.full_name] = lib

    def find_libs(self, directory, is_package=False):
        path = os.path.join(self.static_path, directory)
        package = directory if is_package else None
        if os.path.exists(path):
            for filename in os.listdir(path):
                if os.path.isfile(os.path.join(path, filename)):
                    name, ext = self.parse_filename(filename)
                    require = ('meteor', ) if is_package and ext == 'js' else ()
                    self.append_lib(path, name, ext, require, package)

    def get_chain(self, ext, package=None):
        result = []
        for lib in self.libs[ext].values():
            if (package and lib.package == package) or lib.package is None:
                result.extend(self.get_lib_chain(lib.full_name, ext))
        return distinct(result)

    def get_lib_chain(self, libname, ext):
        lib = self.libs[ext][libname]

        result = []
        for req in lib.require:
            result.extend(self.get_lib_chain(req, ext))

        result.append(lib)
        return distinct(result)

    def parse_filename(self, filename):
        name = ''.join(filename.split('.')[:-1])
        ext = filename.split('.')[-1]
        return name, ext

class StaticLib(object):
    def __init__(self, static_path, path, name, ext, require=(), package=None):
        self.static_path = static_path
        self.path = path
        self.name = name
        self.ext = ext
        self.require = require
        self.package = package

    @property
    def full_name(self):
        return '/'.join([self.package, self.name]) \
            if self.package else self.name

    def get_relative_path(self, minified=False):
        path = self.path.replace(self.static_path, '')
        if minified:
            result = os.path.join(
                path, 'compressed', '{0}-min.{1}'.format(self.name, self.ext))
        else:
            result = '.'.join([os.path.join(path, self.name), self.ext])
        return result.lstrip('/')

    def __str__(self):
        return self.full_name


from datacategoryvisitors.ScaledDataCategoryVisitor import ScaledDataCategoryVisitor
import unittest

import numpy as np
from dataprocessors.DataProcessorGaussAndCosine import DataProcessorGaussAndCosine

class DataProcessorGaussAndCosineTest (unittest.TestCase):
    """Test code for the DataProcessorGaussAndCosine class"""
    def setUp(self) -> None:
        data = [
            [1,0,3,"Braund, Mr. Owen Harris","male",22,1,0,"A/5 21171",7.25,np.nan,"S"],
            [2,1,1,"Cumings, Mrs. John Bradley (Florence Briggs Thayer)","female",38,1,0,"PC 17599",71.2833,"C85","C"],
            [3,1,3,"Heikkinen, Miss. Laina","female",26,0,0,"STON/O2. 3101282",7.925,np.nan,"S"]
        ]
        self.dataProcessorGaussAndCosine = DataProcessorGaussAndCosine(data, True, ScaledDataCategoryVisitor())

    def testGetCosineSimilarity(self):
        """Tests that the getCosineSimilarity method returns a value of 0 for
        vectors that are orthogonal and a value of 1 for values that are in the
        same direction"""
        list1 = [1,0,1,0]
        list2 = [0,1,0,1]
        list3 = [2,0,2,0]
        self.assertEquals(self.dataProcessorGaussAndCosine._getCosineSimilarity(list1, list2), 0)
        self.assertAlmostEquals(self.dataProcessorGaussAndCosine._getCosineSimilarity(list1, list3), 1.0, delta=0.01)

    def testString(self):
        self.assertEquals(str(self.dataProcessorGaussAndCosine), "GausCosS:1.0")


number = None
def transform_string(number):

    if (n % 10 == 1) and n != 11:
        number = list[n, 'процент']
    elif (n % 10 >= 2) and (n % 10 <=4) and n != 12 and n != 13 and n != 14:
        number = list[n, 'процента']
    else:
        number = list[n, 'процентов']
    return number

for n in range(1, 101):
    print(transform_string(n))


from copy import copy

INPUT = open("./input_files/input_22", "r").read().strip("\n")
# INPUT = """Player 1:
# 9
# 2
# 6
# 3
# 1
#
# Player 2:
# 5
# 8
# 4
# 7
# 10"""
PLAYER_1, PLAYER_2 = [[int(c) for c in x.split('\n')[1:]] for x in INPUT.split('\n\n')]

def combat(player_1, player_2):
    while len(player_1) != 0 and len(player_2) != 0:
        c1 = player_1.pop(0)
        c2 = player_2.pop(0)

        if c1 > c2:
            player_1 += [c1, c2]
        else:
            player_2 += [c2, c1]

    return max(player_1, player_2, key=lambda x: len(x))

def run_1():
    winner = combat(copy(PLAYER_1), copy(PLAYER_2))
    return sum(i * v for i, v in enumerate(winner[::-1], 1))

def game_key(player_1, player_2):
    return ','.join(str(x) for x in player_1) + '::' + ','.join(str(x) for x in player_2)

def recursive_combat(player_1, player_2, return_type='score'):
    games = []
    while len(player_1) != 0 and len(player_2) != 0:
        key = game_key(player_1, player_2)
        if key in games:
            return player_1 if return_type == 'score' else 'player_1'
        games.append(key)

        c1 = player_1.pop(0)
        c2 = player_2.pop(0)

        winner = 'player_1' if c1 > c2 else 'player_2'
        if len(player_1) >= c1 and len(player_2) >= c2:
            winner = recursive_combat(player_1[:c1], player_2[:c2], 'player')

        if winner == 'player_1':
            player_1 += [c1, c2]
        else:
            player_2 += [c2, c1]

    if return_type == 'score':
        return max(player_1, player_2, key=lambda x: len(x))
    return 'player_1' if len(player_2) == 0 else 'player_2'

def run_2():
    winner = recursive_combat(copy(PLAYER_1), copy(PLAYER_2))
    return sum(i * v for i, v in enumerate(winner[::-1], 1))

print(run_1())
print(run_2())


# -*- coding: utf-8 -*-
"""
Created on Wed Nov 18 15:17:58 2020

Author: Josef Perktold
License: BSD-3

"""

import warnings

import numpy as np

from statsmodels.tools.decorators import cache_readonly

from statsmodels.stats.diagnostic_gen import (
    test_chisquare_binning
    )
from statsmodels.discrete._diagnostics_count import (
    test_poisson_dispersion,
    # _test_poisson_dispersion_generic,
    test_poisson_zeroinflation_jh,
    test_poisson_zeroinflation_broek,
    test_poisson_zeros,
    test_chisquare_prob,
    plot_probs
    )

class CountDiagnostic:
    """Diagnostic and specification tests and plots for Count model

    status: experimental

    Parameters
    ----------
    results : Results instance of a count model.
    y_max : int
        Largest count to include when computing predicted probabilities for
        counts. Default is the largest observed count.

    """

    def __init__(self, results, y_max=None):
        self.results = results
        self.y_max = y_max

    @cache_readonly
    def probs_predicted(self):
        if self.y_max is not None:
            kwds = {"y_values": np.arange(self.y_max + 1)}
        else:
            kwds = {}
        return self.results.predict(which="prob", **kwds)

    def test_chisquare_prob(self, bin_edges=None, method=None):
        """Moment test for binned probabilites using OPG.

        Paramters
        ---------
        binedges : array_like or None
            This defines which counts are included in the test on frequencies
            and how counts are combined in bins.
            The default if bin_edges is None will change in future.
            See Notes and Example sections below.
        method : str
            Currently only `method = "opg"` is available.
            If method is None, the OPG will be used, but the default might
            change in future versions.
            See Notes section below.

        Returns
        -------
        test result

        Notes
        -----
        Warning: The current default can have many empty or nearly empty bins.
        The default number of bins is given by max(endog).
        Currently it is recommended to limit the number of bins explicitly,
        see Examples below.
        Binning will change in future and automatic binning will be added.

        Currently only the outer product of gradient, OPG, method is
        implemented. In many case, the OPG version of a specification test
        overrejects in small samples.
        Specialized tests that use observed or expected information matrix
        often have better small sample properties.
        The default method will change if better methods are added.

        Examples
        --------
        The following call is a test for the probability of zeros
        `test_chisquare_prob(bin_edges=np.arange(3))`

        `test_chisquare_prob(bin_edges=np.arange(10))` tests the hypothesis
        that the frequencies for counts up to 7 correspond to the estimated
        Poisson distributions.
        In this case, edges are 0, ..., 9 which defines 9 bins for
        counts 0 to 8. The last bin is dropped, so the joint test hypothesis is
        that the observed aggregated frequencies for counts 0 to 7 correspond
        to the model prediction for those frequencies. Predicted probabilites
        Prob(y_i = k | x) are aggregated over observations ``i``.

        """
        kwds = {}
        if bin_edges is not None:
            # TODO: verify upper bound, we drop last bin (may be open, inf)
            kwds["y_values"] = np.arange(bin_edges[-2] + 1)
        probs = self.results.predict(which="prob", **kwds)
        res = test_chisquare_prob(self.results, probs, bin_edges=bin_edges,
                                  method=method)
        return res

    def plot_probs(self, label='predicted', upp_xlim=None,
                   fig=None):
        """Plot observed versus predicted frequencies for entire sample.
        """
        probs_predicted = self.probs_predicted.sum(0)
        k_probs = len(probs_predicted)
        freq = np.bincount(self.results.model.endog.astype(int),
                           minlength=k_probs)[:k_probs]
        fig = plot_probs(freq, probs_predicted,
                         label=label, upp_xlim=upp_xlim,
                         fig=fig)
        return fig

class PoissonDiagnostic(CountDiagnostic):
    """Diagnostic and specification tests and plots for Poisson model

    status: experimental

    Parameters
    ----------
    results : PoissonResults instance

    """

    def _init__(self, results):
        self.results = results

    def test_dispersion(self):
        """Test for excess (over or under) dispersion in Poisson.

        Returns
        -------
        dispersion results
        """
        res = test_poisson_dispersion(self.results)
        return res

    def test_poisson_zeroinflation(self, method="prob", exog_infl=None):
        """Test for excess zeros, zero inflation or deflation.

        Parameters
        ----------
        method : str
            Three methods ara available for the test:

             - "prob" : moment test for the probability of zeros
             - "broek" : score test against zero inflation with or without
                explanatory variables for inflation

        exog_infl : array_like or None
            Optional explanatory variables under the alternative of zero
            inflation, or deflation. Only used if method is "broek".

        Returns
        -------
        results

        Notes
        -----
        If method = "prob", then the moment test of He et al 1_ is used based
        on the explicit formula in Tang and Tang 2_.

        If method = "broek" and exog_infl is None, then the test by Van den
        Broek 3_ is used. This is a score test against and alternative of
        constant zero inflation or deflation.

        If method = "broek" and exog_infl is provided, then the extension of
        the broek test to varying zero inflation or deflation by Jansakul and
        Hinde is used.

        Warning: The Broek and the Jansakul and Hinde tests are not numerically
        stable when the probability of zeros in Poisson is small, i.e. if the
        conditional means of the estimated Poisson distribution are large.
        In these cases, p-values will not be accurate.
        """
        if method == "prob":
            if exog_infl is not None:
                warnings.warn('exog_infl is only used if method = "broek"')
            res = test_poisson_zeros(self.results)
        elif method == "broek":
            if exog_infl is None:
                res = test_poisson_zeroinflation_broek(self.results)
            else:
                exog_infl = np.asarray(exog_infl)
                if exog_infl.ndim == 1:
                    exog_infl = exog_infl[:, None]
                res = test_poisson_zeroinflation_jh(self.results,
                                                    exog_infl=exog_infl)

        return res

    def _chisquare_binned(self, sort_var=None, bins=10, k_max=None, df=None,
                          sort_method="quicksort", frac_upp=0.1,
                          alpha_nc=0.05):
        """Hosmer-Lemeshow style test for count data.

        Note, this does not take into account that parameters are estimated.
        The distribution of the test statistic is only an approximation.

        This corresponds to the Hosmer-Lemeshow type test for an ordinal
        response variable. The outcome space y = k is partitioned into bins
        and treated as ordinal variable.
        The observations are split into approximately equal sized groups
        of observations sorted according the ``sort_var``.

        """

        if sort_var is None:
            sort_var = self.results.predict(which="lin")

        endog = self.results.model.endog
        # not sure yet how this is supposed to work
        # max_count = endog.max * 2
        # no option for max count in predict
        # counts = (endog == np.arange(max_count)).astype(int)
        expected = self.results.predict(which="prob")
        counts = (endog[:, None] == np.arange(expected.shape[1])).astype(int)

        # truncate upper tail
        if k_max is None:
            nobs = len(endog)
            icumcounts_sum = nobs - counts.sum(0).cumsum(0)
            k_max = np.argmax(icumcounts_sum < nobs * frac_upp) - 1
        expected = expected[:, :k_max]
        counts = counts[:, :k_max]
        # we should correct for or include truncated upper bin
        # inplace modification, we cannot reuse expected and counts anymore
        expected[:, -1] += 1 - expected.sum(1)
        counts[:, -1] += 1 - counts.sum(1)

        # TODO: what's the correct df, same as for multinomial/ordered ?
        res = test_chisquare_binning(counts, expected, sort_var=sort_var,
                                     bins=bins, df=df, ordered=True,
                                     sort_method=sort_method,
                                     alpha_nc=alpha_nc)
        return res


# -*- coding: utf-8 -*-
"""
Created on Wed Mar 11 22:07:07 2020

@author: alexa
"""

"""" Mediante los algoritmos de smv podremos ver cuánto de seguro está nuestro
sistema de predicción de las clasificaciones que hace. con decision_function vemos cuán separado está
un dato de una línea separadora creado por el algoritmo. Con predict_proba es lo mismo pero se muestra con probabilidades,
y con predict, directamente predice los resultados.
"""

from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn import svm    #Vectores de soporte: traza hiperplano para clasificar data

iris = load_iris()

X_train,X_test,y_train,y_test = train_test_split(iris.data, iris.target)

algoritmo = svm.SVC(probability = True)

algoritmo.fit(X_train,y_train)

algoritmo.decision_function_shape = "ovr"

decision_function = algoritmo.decision_function(X_test)[:10]  #Número de decisiones 10

predict_proba = algoritmo.predict_proba(X_test)[:10]

predict = algoritmo.predict(X_test)[:10]


from openpype.client import (
    get_version_by_id,
    get_last_version_by_subset_id
)
from openpype.pipeline import (
    get_representation_path,
    get_current_project_name,
)
from openpype.lib.transcoding import (
    VIDEO_EXTENSIONS,
    IMAGE_EXTENSIONS
)
import openpype.hosts.hiero.api as phiero

class LoadClip(phiero.SequenceLoader):
    """Load a subset to timeline as clip

    Place clip to timeline on its asset origin timings collected
    during conforming to project
    """

    families = ["render2d", "source", "plate", "render", "review"]
    representations = ["*"]
    extensions = set(
        ext.lstrip(".") for ext in IMAGE_EXTENSIONS.union(VIDEO_EXTENSIONS)
    )

    label = "Load as clip"
    order = -10
    icon = "code-fork"
    color = "orange"

    # for loader multiselection
    sequence = None
    track = None

    # presets
    clip_color_last = "green"
    clip_color = "red"

    clip_name_template = "{asset}_{subset}_{representation}"

    @classmethod
    def apply_settings(cls, project_settings, system_settings):
        plugin_type_settings = (
            project_settings
            .get("hiero", {})
            .get("load", {})
        )

        if not plugin_type_settings:
            return

        plugin_name = cls.__name__

        plugin_settings = None
        # Look for plugin settings in host specific settings
        if plugin_name in plugin_type_settings:
            plugin_settings = plugin_type_settings[plugin_name]

        if not plugin_settings:
            return

        print(">>> We have preset for {}".format(plugin_name))
        for option, value in plugin_settings.items():
            if option == "enabled" and value is False:
                print("  - is disabled by preset")
            elif option == "representations":
                continue
            else:
                print("  - setting `{}`: `{}`".format(option, value))
            setattr(cls, option, value)

    def load(self, context, name, namespace, options):
        # add clip name template to options
        options.update({
            "clipNameTemplate": self.clip_name_template
        })
        # in case loader uses multiselection
        if self.track and self.sequence:
            options.update({
                "sequence": self.sequence,
                "track": self.track,
                "clipNameTemplate": self.clip_name_template
            })

        # load clip to timeline and get main variables
        path = self.filepath_from_context(context)
        track_item = phiero.ClipLoader(self, context, path, **options).load()
        namespace = namespace or track_item.name()
        version = context['version']
        version_data = version.get("data", {})
        version_name = version.get("name", None)
        colorspace = version_data.get("colorspace", None)
        object_name = self.clip_name_template.format(
            **context["representation"]["context"])

        # set colorspace
        if colorspace:
            track_item.source().setSourceMediaColourTransform(colorspace)

        # add additional metadata from the version to imprint Avalon knob
        add_keys = [
            "frameStart", "frameEnd", "source", "author",
            "fps", "handleStart", "handleEnd"
        ]

        # move all version data keys to tag data
        data_imprint = {}
        for key in add_keys:
            data_imprint.update({
                key: version_data.get(key, str(None))
            })

        # add variables related to version context
        data_imprint.update({
            "version": version_name,
            "colorspace": colorspace,
            "objectName": object_name
        })

        # update color of clip regarding the version order
        self.set_item_color(track_item, version)

        # deal with multiselection
        self.multiselection(track_item)

        self.log.info("Loader done: `{}`".format(name))

        return phiero.containerise(
            track_item,
            name, namespace, context,
            self.__class__.__name__,
            data_imprint)

    def switch(self, container, representation):
        self.update(container, representation)

    def update(self, container, representation):
        """ Updating previously loaded clips
        """

        # load clip to timeline and get main variables
        name = container['name']
        namespace = container['namespace']
        track_item = phiero.get_track_items(
            track_item_name=namespace).pop()

        project_name = get_current_project_name()
        version_doc = get_version_by_id(project_name, representation["parent"])

        version_data = version_doc.get("data", {})
        version_name = version_doc.get("name", None)
        colorspace = version_data.get("colorspace", None)
        object_name = "{}_{}".format(name, namespace)
        file = get_representation_path(representation).replace("\\", "/")
        clip = track_item.source()

        # reconnect media to new path
        clip.reconnectMedia(file)

        # set colorspace
        if colorspace:
            clip.setSourceMediaColourTransform(colorspace)

        # add additional metadata from the version to imprint Avalon knob
        add_keys = [
            "frameStart", "frameEnd", "source", "author",
            "fps", "handleStart", "handleEnd"
        ]

        # move all version data keys to tag data
        data_imprint = {}
        for key in add_keys:
            data_imprint.update({
                key: version_data.get(key, str(None))
            })

        # add variables related to version context
        data_imprint.update({
            "representation": str(representation["_id"]),
            "version": version_name,
            "colorspace": colorspace,
            "objectName": object_name
        })

        # update color of clip regarding the version order
        self.set_item_color(track_item, version_doc)

        return phiero.update_container(track_item, data_imprint)

    def remove(self, container):
        """ Removing previously loaded clips
        """
        # load clip to timeline and get main variables
        namespace = container['namespace']
        track_item = phiero.get_track_items(
            track_item_name=namespace).pop()
        track = track_item.parent()

        # remove track item from track
        track.removeItem(track_item)

    @classmethod
    def multiselection(cls, track_item):
        if not cls.track:
            cls.track = track_item.parent()
            cls.sequence = cls.track.parent()

    @classmethod
    def set_item_color(cls, track_item, version_doc):
        project_name = get_current_project_name()
        last_version_doc = get_last_version_by_subset_id(
            project_name, version_doc["parent"], fields=["_id"]
        )
        clip = track_item.source()
        # set clip colour
        if version_doc["_id"] == last_version_doc["_id"]:
            clip.binItem().setColor(cls.clip_color_last)
        else:
            clip.binItem().setColor(cls.clip_color)


# -*- coding: utf-8 -*-

import time, threading

def heartBeatTimer():
    print('xintiao')

#
def Methond1():
    while True:
        print('UDP')
        time.sleep(2)
#
def Methond2():
    while True:
        print('TCP')
        time.sleep(3)

#
if __name__ == '__main__':

    UDPS = threading.Thread(target=Methond1, args=())
    UDPS.start()

    print('test')

    TCPS = threading.Thread(target=Methond2, args=())
    TCPS.start()


import os, sys
sys.path.insert(0, os.path.abspath('../detectron2/'))

import argparse
import glob
import multiprocessing as mp
import numpy as np
import os
import tempfile
import time
import warnings
import cv2
import tqdm

from detectron2.config import get_cfg
from detectron2.data.detection_utils import read_image
from detectron2.utils.logger import setup_logger

from diffusiondet.predictor import VisualizationDemo
from diffusiondet import DiffusionDetDatasetMapper, add_diffusiondet_config, DiffusionDetWithTTA
from diffusiondet.util.model_ema import add_model_ema_configs, may_build_model_ema, may_get_ema_checkpointer, EMAHook, \
    apply_model_ema_and_restore, EMADetectionCheckpointer

# constants
WINDOW_NAME = "COCO detections"

def setup_cfg(args):
    # load config from file and command-line arguments
    cfg = get_cfg()
    # To use demo for Panoptic-DeepLab, please uncomment the following two lines.
    # from detectron2.projects.panoptic_deeplab import add_panoptic_deeplab_config  # noqa
    # add_panoptic_deeplab_config(cfg)
    add_diffusiondet_config(cfg)
    add_model_ema_configs(cfg)
    cfg.merge_from_file(args.config_file)
    cfg.merge_from_list(args.opts)
    # Set score_threshold for builtin models
    cfg.MODEL.RETINANET.SCORE_THRESH_TEST = args.confidence_threshold
    cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = args.confidence_threshold
    cfg.MODEL.PANOPTIC_FPN.COMBINE.INSTANCES_CONFIDENCE_THRESH = args.confidence_threshold
    cfg.freeze()
    return cfg

def get_parser():
    parser = argparse.ArgumentParser(description="Detectron2 demo for builtin configs")
    parser.add_argument(
        "--config-file",
        default="configs/quick_schedules/mask_rcnn_R_50_FPN_inference_acc_test.yaml",
        metavar="FILE",
        help="path to config file",
    )
    parser.add_argument("--webcam", action="store_true", help="Take inputs from webcam.")
    parser.add_argument("--video-input", help="Path to video file.")
    parser.add_argument(
        "--input",
        nargs="+",
        help="A list of space separated input images; "
        "or a single glob pattern such as 'directory/*.jpg'",
    )
    parser.add_argument(
        "--output",
        help="A file or directory to save output visualizations. "
        "If not given, will show output in an OpenCV window.",
    )

    parser.add_argument(
        "--confidence-threshold",
        type=float,
        default=0.5,
        help="Minimum score for instance predictions to be shown",
    )
    parser.add_argument(
        "--opts",
        help="Modify config options using the command-line 'KEY VALUE' pairs",
        default=[],
        nargs=argparse.REMAINDER,
    )
    return parser

if __name__ == "__main__":
    mp.set_start_method("spawn", force=True)
    args = get_parser().parse_args()
    setup_logger(name="fvcore")
    logger = setup_logger()
    logger.info("Arguments: " + str(args))

    cfg = setup_cfg(args)

    demo = VisualizationDemo(cfg)
    # Set up the output file
    output_file = open(args.output, 'w')

    if args.input:
        if len(args.input) == 1:
            args.input = glob.glob(os.path.expanduser(args.input[0]))
            assert args.input, "The input path(s) was not found"
        i= 1
        for path in tqdm.tqdm(args.input, disable=not args.output):
            # use PIL, to be consistent with evaluation
            img = read_image(path, format="BGR")
            start_time = time.time()
            predictions, visualized_output = demo.run_on_image(img)
            scores = predictions['instances'].to('cpu').scores.numpy()

            detections = []
            for j in range(len(scores)):
              if scores[j] > args.confidence_threshold:
                box = predictions['instances'].to('cpu').pred_boxes.tensor[j].numpy()
                #box = boxes[j, :]
                x1 = int(box[0])
                y1 = int(box[1])
                x2 = int(box[2])
                y2 = int(box[3])
                detections.append((x1, y1, x2 - x1, y2 - y1, scores[j]))
            # Write the detections to the output file
            for det in detections:
              x, y, w, h, conf = det
              output_file.write(f'{i},-1,{x},{y},{w},{h},1,-1,-1,-1\n')
            i+=1


#!/usr/bin/env python3
#
# This file is distributed under the MIT License. See LICENSE for details.
#

import argparse
import re
import pprint
import subprocess
import json
from smackgen import *
from smackverify import *

VERSION = '2.8.0'

def reachParser():
    parser = argparse.ArgumentParser(add_help=False, parents=[verifyParser()])

    return parser

# File line numbers are 0-based idx

def CopyFileWhileInserting(srcFile, dstFile, lineNumber, insertText):
    inFile = open(srcFile, "r")
    inContents = inFile.readlines()
    inFile.close()

    if(not insertText.endswith('\n')):
        insertText += '\n'

    inContents.insert(lineNumber, insertText)

    outFile = open(dstFile, "w")
    outFile.write("".join(inContents))
    outFile.close()

# Each item in return value is [fileName, sourceLineNo, bplLineNo, isReachable]

def GetSourceLineInfo(bplFile):
    FILENAME = r'[\w#$~%.\/-]+'
    regex = '.*{:sourceloc \"(' + FILENAME + r')\", (\d+), (\d+)}.*'
    sourcelocRe = re.compile(regex)

    sourceInfo = []

    bplCount = 0

    with open(bplFile) as inFile:
        for line in inFile.readlines():
            # Groups: 1=filename, 2=sourceLineNo, 3=sourceColNo
            match = sourcelocRe.match(line)
            if(match):
                newSource = {
                    'filename': match.group(1),
                    'sourceLineNo': int(match.group(2)),
                    'sourceColNo': int(match.group(3)),
                    'bplLineNo': bplCount,
                    'isReachable': False
                }
                sourceInfo.append(newSource)
            bplCount += 1

    return sorted(sourceInfo, key=lambda e: e['sourceLineNo'], reverse=True)

def UpdateWithClangInfo(clangOuptut, sourceInfo):
    FILENAME = r'[\w#$~%.\/-]+'
    regex = ('('
             + FILENAME
             + (r'):(\d+):(\d+): warning: will never be executed '
                r'\[-Wunreachable-code\]'))
    clangFilter = re.compile(regex)

    for line in clangOutput.splitlines(True):
        match = clangFilter.match(line)
        if(match):
            newSource = {
                'filename': match.group(1),
                'sourceLineNo': int(match.group(2)),
                'sourceColNo': int(match.group(3)),
                'bplLineNo': -1,
                'isReachable': False
            }
            sourceInfo.append(newSource)

def GetCodeCoverage(
        verifier,
        bplFileName,
        timeLimit,
        unroll,
        contextSwitches,
        debug,
        smackd,
        clangOutput):
    sourceInfo = GetSourceLineInfo(bplFileName)

    for sourceLine in sourceInfo:
        if(not sourceLine['isReachable']):
            reachRes = TestReachability(
                verifier,
                bplFileName,
                timeLimit,
                unroll,
                contextSwitches,
                debug,
                sourceLine)

            # TODO - how does python handle changing lists in for loop?
            UpdateSourceInfo(reachRes, sourceInfo, verifier)

    # Add lines caught by clang's -Wunreachable-code
    UpdateWithClangInfo(clangOutput, sourceInfo)

    # Extract info
    result = {}
    for sourceLine in sourceInfo:
        if(not sourceLine["isReachable"]):
            if(not sourceLine["filename"] in result):
                result[sourceLine["filename"]] = set()
            resItem = sourceLine["sourceLineNo"], sourceLine["sourceColNo"]
            result[sourceLine["filename"]].add(resItem)

    for curfile in result:
        # secondary sort by column
        result[curfile] = sorted(
            result[curfile],
            key=lambda e: e[1],
            reverse=False)
        # primary sort by line
        result[curfile] = sorted(
            result[curfile],
            key=lambda e: e[0],
            reverse=False)

    if(smackd):
        print((json.dumps(result)))
    else:
        print('\nSMACK verifier version ' + VERSION + '\n\n')
        print("Unreachable code:")
        pprint.pprint(result, width=100)

def TestReachability(
        verifier,
        bplFileName,
        timeLimit,
        unroll,
        contextSwitches,
        debug,
        lineInfo):
    boogieText = "assert false;"

    bplfileBase = path.splitext(bplFileName)[0]
    bplNew = bplfileBase + "_coverage.bpl"

    CopyFileWhileInserting(
        bplFileName,
        bplNew,
        lineInfo['bplLineNo'] + 1,
        boogieText)

    # do not pass smackd flag as true.  Breaks parsing
    corralOutput = verify(
        verifier,
        bplNew,
        timeLimit,
        unroll,
        contextSwitches,
        debug,
        False)

    return corralOutput

def UpdateSourceInfo(corralOutput, sourceInfo, verifier):
    FILENAME = r'[\w#$~%.\/-]+'
    regex = ""
    if(verifier == "corral"):
        regex = '(' + FILENAME + r')\((\d+),(\d+)\): Trace:.*'
    else:
        # boogie...
        # TODO Once line numbers are fixed for boogie-inline,
        #     we can save time here by using all traces,
        #     not just the instruction causing the failure
        #     (use trace that led to failing instruction, also)
        #  As is, current smackverify output using boogie-inline
        #     is including unvisited lines in traces. (See
        #     src/test/reach/switch.c)
        regex = r'\s*(' + FILENAME + r')\((\d+),(\d+)\): Error.*'
    traceFilter = re.compile(regex)

    for line in corralOutput.splitlines(True):
        match = traceFilter.match(line)
        if(match):
            reachedLine = {
                'filename': match.group(1),
                'sourceLineNo': int(match.group(2)),
                # Corral adds one to column count
                'sourceColNo': int(match.group(3)),
            }
            # run through each sourceInfo, if matches, set as reachable
            for sourceLine in sourceInfo:
                if ((not sourceLine['isReachable']) and
                        reachedLine['filename'] == sourceLine['filename'] and
                        (reachedLine['sourceLineNo'] ==
                         sourceLine['sourceLineNo']) and
                        (reachedLine['sourceColNo'] ==
                         sourceLine['sourceColNo'])):
                    sourceLine['isReachable'] = True

def main():
    parser = argparse.ArgumentParser(
        description='Checks the input LLVM file for code reachability.',
        parents=[
            reachParser()])
    parser.parse_args()  # just check if arguments are looking good

    # !!!!!!START COPY OF SECTION FROM smackverify.py!!!!!!!!!!!
    #  Probably should pull into subroutine or something
    # remove arguments not recognized by lower scripts
    # not sure of a better way to do this
    sysArgv = sys.argv[:]
    for i in reversed(list(range(len(sysArgv)))):
        if sysArgv[i] == '--smackd':
            del sysArgv[i]
        elif sys.argv[i] == '--time-limit':
            del sysArgv[i]
            del sysArgv[i]
        elif sys.argv[i] == '--context-bound':
            del sysArgv[i]
            del sysArgv[i]

    # Add clang's -Wunreachable-code flag
    sysArgv.append('--clang=-Wunreachable-code')
    bpl, options, clangOutput = smackGenerate(sysArgv)
    args = parser.parse_args(options + sys.argv[1:])

    # write final output
    args.outfile.write(bpl)
    args.outfile.close()
    # !!!!!!END COPY OF SECTION FROM smackverify.py!!!!!!!!!!!

    GetCodeCoverage(
        args.verifier,
        args.outfile.name,
        args.timeLimit,
        args.unroll,
        args.contextSwitches,
        args.debug,
        args.smackd,
        clangOutput)


#!/usr/bin/python3
# -*- coding: utf-8 -*-

import logging
import numpy as np
import quadprog
from typing import List

from shimmingtoolbox.optimizer.optimizer_utils import OptimizerUtils
from shimmingtoolbox.pmu import PmuResp
from shimmingtoolbox.coils.coil import Coil

ListCoil = List[Coil]
logger = logging.getLogger(__name__)

class QuadProgOpt(OptimizerUtils):
    """ Optimizer object that stores coil profiles and optimizes an unshimmed volume given a mask.
        Use optimize(args) to optimize a given mask. The algorithm uses a quadprog solver to find the best shim.
        It supports bounds for each channel as well as a bound for the absolute sum of the channels.
    """

    def __init__(self, coils: ListCoil, unshimmed, affine, reg_factor=0, initial_guess_method='mean'):
        """
        Initializes coils according to input list of Coil

        Args:
            coils (ListCoil): List of Coil objects containing the coil profiles and related constraints
            unshimmed (np.ndarray): 3d array of unshimmed volume
            affine (np.ndarray): 4x4 array containing the affine transformation for the unshimmed array
            reg_factor (float): Regularization factor for the current when optimizing. A higher coefficient will
                                penalize higher current values while a lower factor will lower the effect of the
                                regularization. A negative value will favour high currents (not preferred).
            initial_guess_method (str): method to find the initial guess
        """
        if reg_factor < 0:
            raise TypeError(f"reg_factor is negative, and would cause optimization to crash."
                            f" If you want to keep this reg_factor please use lsq_optimizer")
        super().__init__(coils, unshimmed, affine, initial_guess_method, reg_factor)

    def _get_linear_inequality_matrices(self):
        """
        This functions returns the linear inequality matrix and vector, that will be used in the optimization, such as
        g @ x < h, to see all details please see the PR 458

        Returns:
            (tuple) : tuple containing:
                * np.ndarray: linear inequality matrix
                * np.ndarray: linear inequality vector

        """
        n = len(self.reg_vector)
        n_coils = len(self.coils)

        sum_constraints = []
        g_bis = []

        # First we want to create the part to verify that the sum of the currents don't get above a limit
        start_index = n
        for i_coil in range(n_coils):
            coil = self.coils[i_coil]
            end_index = start_index + coil.dim[3]
            if coil.coef_sum_max != np.inf:
                sum_constraints.append(coil.coef_sum_max)
                tmp = np.zeros([2 * n])
                tmp[start_index: end_index] = 1
                g_bis.append(tmp)
            start_index = end_index
        g_bis = np.array(g_bis)
        sum_constraints = np.array(sum_constraints)[:, np.newaxis]

        # Then we want to see if the currents are between a lower, and an upper bound
        lb = np.zeros([n, 1])
        lb[:, 0] = np.array([x[0] for x in self.merged_bounds])
        ub = np.zeros([n, 1])
        ub[:, 0] = np.array([x[1] for x in self.merged_bounds])

        # We create both arrays, but if there is no sum constraints, then there is no need to add this constraint
        if np.any(sum_constraints):
            g = np.block([[-np.eye(n), -np.eye(n)], [np.eye(n), -np.eye(n)], [g_bis],
                          [-np.eye(n), np.zeros([n, n])], [np.eye(n), np.zeros([n, n])]])
            # dim(g) = (4n + n_sum_constraints, 2n)
            h = np.block([[np.zeros([n, 1])], [np.zeros([n, 1])], [sum_constraints],
                          [-lb], [ub]])
            # dim(h) = (4n + n_sum_constraints, 1)
        else:
            g = np.block([[-np.eye(n), -np.eye(n)], [np.eye(n), -np.eye(n)],
                          [-np.eye(n), np.zeros([n, n])], [np.eye(n), np.zeros([n, n])]])
            # dim(g) = (4n, 2n)
            h = np.block([[np.zeros([n, 1])], [np.zeros([n, 1])],
                          [-lb], [ub]])
            # dim(h) = (4n, 1)
        return g, h

    def get_stability_factor(self, coef, unshimmed_vec, coil_mat, factor):
        """ Objective function to find the stability factor for the quadratic optimization

        Args:
            coef (np.ndarray): 1D array of channel coefficients
            unshimmed_vec (np.ndarray): 1D flattened array (point) of the masked unshimmed map
            coil_mat (np.ndarray): 2D flattened array (point, channel) of masked coils
                                      (axis 0 must align with unshimmed_vec)
            factor (float): Devise the result by 'factor'. This allows to scale the output for the minimize function to
                            avoid positive directional linesearch

        Returns:
            float: Residuals for quad_prog optimization
        """
        shimmed_vec = unshimmed_vec + coil_mat @ coef
        return shimmed_vec.dot(shimmed_vec) / len(unshimmed_vec) / factor + np.abs(coef).dot(self.reg_vector)

    def _get_currents(self, unshimmed_vec, coil_mat, currents_0):
        """
        Returns the currents needed for the shimming, redefined from basic_optimizer because of the constraints
        Args:
            unshimmed_vec (np.ndarray): 1D flattened array (point) of the masked unshimmed map
            coil_mat (np.ndarray): 2D flattened array (point, channel) of masked coils
                                      (axis 0 must align with unshimmed_vec)
            currents_0 (np.ndarray) : Initial guess for the function

        Returns:
            np.ndarray : Vector of currents that make the better shimming

        """
        # This allows to scale the output for the minimize function to avoid positive directional linesearch
        stability_factor = self.get_stability_factor(self._initial_guess_zeros(), unshimmed_vec,
                                                     np.zeros_like(coil_mat),
                                                     factor=1)
        # We want to minimize 1/2 x.T @ cost_matrix @ x - cost_vector.T @ x
        cost_matrix, cost_vector = self.get_cost_matrices(currents_0, unshimmed_vec, coil_mat, stability_factor)
        # We want to get the linear inequality matrix and vector, such as g @ x < h
        ineq_matrix, ineq_vector = self._get_linear_inequality_matrices()

        # We use the quadprog solver to solve the quadratic optimization problem
        opt_result = quadprog.solve_qp(cost_matrix, -cost_vector, ineq_matrix.T, -ineq_vector[:, 0])
        currents = opt_result[0]

        if logger.level <= getattr(logging, 'DEBUG'):
            n_iter = opt_result[3][0]
            logger.debug(f"Number of iterations for quadprog: {n_iter}")

        return currents[:len(currents_0)]

    def get_cost_matrices(self, currents_0, unshimmed_vec, coil_mat, factor):
        """
        Returns the cost matrix and the cost vector to minimize 1/2 x.T @ cost_matrix @ x - cost_vector.T @ x

        Args:
            currents_0 (np.ndarray) : Initial guess for the function
            unshimmed_vec (np.ndarray): 1D flattened array (point) of the masked unshimmed map
            coil_mat (np.ndarray): 2D flattened array (point, channel) of masked coils
                                      (axis 0 must align with unshimmed_vector)
            factor (float): Divide the result by 'factor'. This allows to scale the output for the minimize function to
                            avoid positive directional linesearch

        Returns:
            (tuple) : tuple containing:
                * np.ndarray: 2D Cost matrix
                * np.ndarray: Cost vector

        """

        n = len(currents_0)

        initial_guess = np.zeros(2 * n)
        initial_guess[:n] = currents_0
        a, b, _ = self.get_quadratic_term(unshimmed_vec, coil_mat, factor)
        epsilon = 1e-6
        cost_matrix = np.block([[a, np.zeros([n, n])], [np.zeros([n, n]), np.zeros([n, n])]]) + epsilon * np.eye(2*n)
        cost_matrix = 2 * cost_matrix
        cost_vector = np.zeros(2 * n)
        cost_vector[0: n] = b

        return cost_matrix, cost_vector

class PmuQuadProgOpt(QuadProgOpt):
    """ Optimizer for the realtime component (riro) for this optimization:
            field(i_vox) = riro(i_vox) * (acq_pressures - mean_p) + static(i_vox)
            Unshimmed must be in units: [unit_shim/unit_pressure], ex: [Hz/unit_pressure]

            This optimizer bounds the riro results to the coil bounds by taking the range of pressure that can be
            reached by the PMU.
        """

    def __init__(self, coils, unshimmed, affine, pmu: PmuResp, reg_factor=0):
        """
        Initializes coils according to input list of Coil

        Args:
            coils (ListCoil): List of Coil objects containing the coil profiles and related constraints
            unshimmed (np.ndarray): 3d array of unshimmed volume
            affine (np.ndarray): 4x4 array containing the affine transformation for the unshimmed array
            pmu (PmuResp): PmuResp object containing the respiratory trace information.
        """

        super().__init__(coils, unshimmed, affine, reg_factor=reg_factor, initial_guess_method='zeros')
        self.pressure_min = pmu.min
        self.pressure_max = pmu.max

    def _get_linear_inequality_matrices(self):

        """
        This functions returns the linear inequality matrix and vector, that will be used in the optimization, such as
        g @ x < h, to see all details please see the PR 458
        Redefined from QuadProg to match the new bounds and constraints

        Returns:
            (tuple) : tuple containing:
                * np.ndarray: linear inequality matrix
                * np.ndarray: linear inequality vector

        """
        n_coils = len(self.coils)
        n = len(self.reg_vector)
        sum_constraints = []
        ub = np.zeros([n, 1])
        lb = np.zeros([n, 1])
        g_bis = []

        start_index = n
        # First we want to create the part to verify that the sum of the currents doesn't get above a limit
        for i_coil in range(n_coils):
            coil = self.coils[i_coil]
            end_index = start_index + coil.dim[3]
            if coil.coef_sum_max != np.inf:
                pressure = np.max(np.abs(self.pressure_max), np.abs(self.pressure_min))
                sum_constraints.append(coil.coef_sum_max / (coil.dim[3] * pressure))
                tmp = np.zeros([2 * n])
                tmp[start_index: end_index] = 1
                g_bis.append(tmp)
            start_index = end_index

        g_bis = np.array(g_bis)
        sum_constraints = np.array(sum_constraints)[:, np.newaxis]

        # Then we want to see if the currents are between a lower, and an upper bound
        delta_pressure = self.pressure_max - self.pressure_min
        for i_bound, bound in enumerate(self.merged_bounds):
            lb[i_bound, 0] = bound[0]
            ub[i_bound, 0] = bound[1]

        # We create both array, but if there is no sum constraints, then there is no need to add this constraints
        if np.any(sum_constraints):
            g = np.block([[-np.eye(n), -np.eye(n)], [np.eye(n), -np.eye(n)], [g_bis],
                          [- delta_pressure * - np.eye(n), np.zeros([n, n])],
                          [delta_pressure * - np.eye(n), np.zeros([n, n])],
                          [delta_pressure * np.eye(n), np.zeros([n, n])],
                          [- delta_pressure * np.eye(n), np.zeros([n, n])]])
            # dim(g) = (6n + n_sum_constraints, 2n)

            h = np.block([[np.zeros([n, 1])], [np.zeros([n, 1])], [sum_constraints],
                          [-lb], [-lb], [ub], [ub]])
            # dim(h) = (6n + n_sum_constraints, 1)
        else:
            g = np.block([[-np.eye(n), -np.eye(n)], [np.eye(n), -np.eye(n)],
                          [- delta_pressure * - np.eye(n), np.zeros([n, n])],
                          [delta_pressure * - np.eye(n), np.zeros([n, n])],
                          [delta_pressure * np.eye(n), np.zeros([n, n])],
                          [- delta_pressure * np.eye(n), np.zeros([n, n])]])
            # dim(g) = (6n, 2n)

            h = np.block([[np.zeros([n, 1])], [np.zeros([n, 1])],
                          [-lb], [-lb], [ub], [ub]])

            # dim(h) = (6n, 1)
        return g, h


from django.shortcuts import render
from rest_framework.decorators import api_view, permission_classes
from rest_framework.response import Response
from rest_framework.permissions import IsAuthenticated
from .serializer import AuthorSerializer, BookSerializer, UpdateBookSerializer
from .models import Author, Book
from django.contrib.auth.models import User
# Create your views here.

@api_view(['GET'])
@permission_classes([IsAuthenticated])
def welcome(request):
    data = {"message": "Welcome to the bookstore"}
    return Response(data, status=200)

@api_view(['POST'])
@permission_classes([IsAuthenticated])
def add_author(request):
    serializer = AuthorSerializer(data=request.data)

    if serializer.is_valid():
        user = User.objects.get(username=request.user.username)
        author = Author(name=request.data['name'], added_by=user)
        author.save()
        return Response("Author added successfully", status=200)
    else:
        return Response(serializer.errors, status=400)

@api_view(['POST'])
@permission_classes([IsAuthenticated])
def add_book(request):
    #print(request.user)
    serializer = BookSerializer(data=request.data)

    if serializer.is_valid():
        user = request.user
        author = Author.objects.filter(name=request.data['author']).filter(added_by=user)
        print(author)
        for auth in author:
            book = Book(title=request.data['title'], description=request.data['description'], author=auth, added_by=user)
            book.save()
        return Response("book saved successfully", status=200)
    else:
        return Response(serializer.errors, status=400)

@api_view(['GET'])
@permission_classes([IsAuthenticated])
def get_all_books(request):
    user = request.user
    book = Book.objects.filter(added_by=user)
    arr = []
    for b in book:
        res = {}
        res['title'] = b.title
        res['desc'] = b.description
        arr.append(res)
    return Response(data={"books": arr}, status=200)

@api_view(['POST'])
@permission_classes([IsAuthenticated])
def update_book_title(request):
    serializer = UpdateBookSerializer(data=request.data)
    if serializer.is_valid():
        user = request.user
        get_book = Book.objects.filter(title=request.data['title'], added_by=user)
        for book in get_book:
            book.title = request.data['new_title']
            book.save()
        return Response(status=200, data={"updated successfully"})

    else:
        return Response( serializer.errors, status=400)


import matplotlib.pyplot as plt
import torch
from leven_marq import LM

def emg_model(p, t: torch.Tensor = None):

    alpha, mu, sigma, eta = p
    gauss_term = lambda t, mu, sigma: torch.exp(-0.5 * (t-mu)**2 / sigma**2)
    asymm_term = lambda t, mu, sigma, eta: 1 + torch.erf(eta * (t-mu) / (sigma * 2**0.5))

    alpha = 1 if alpha is None else alpha
    gauss = gauss_term(t, mu, sigma)
    asymm = asymm_term(t, mu, sigma, eta)

    return alpha * gauss * asymm

torch.seed()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

norm, mean, sigm, skew = 8, -1, 60, 7.5

gt_params = torch.tensor([norm, mean, sigm, skew], dtype=torch.float64, device=device)
init_params = torch.tensor([5.5, -0.75, 10, 3], dtype=torch.float64, device=device, requires_grad=True)
cost_fun = lambda p, t, y: (y - emg_model(p, t))**2

print('gt params: ', gt_params.detach().cpu())

t = torch.linspace(-1e2, 1e2, int(2e2)).to(device)
data = emg_model(gt_params, t)
data_raw = data + 0.001 * torch.randn(len(data), dtype=torch.float64, device=device)

coeffs = LM(p=init_params, function=cost_fun, data=(t, data_raw))

ret_params = torch.allclose(coeffs[-1], gt_params, atol=1e-1)
print('total iters: ', len(coeffs))
print('estimate params: ', coeffs[-1].detach().cpu())
eps_list = []
for i in range(len(coeffs)):
    eps_list.append(torch.sum(cost_fun(coeffs[i], t=t, y=data_raw)).item())
print('eps = ', eps_list[-1])

plt.plot(eps_list)
plt.title('Loss curve')
plt.ylabel('epsilon')
plt.xlabel('iters')
plt.savefig('img/loss.png')
plt.show()

dot = plt.scatter(t.cpu(), data_raw.cpu(), s=0.2)
l1, = plt.plot(data.detach().cpu(), color='blue')
l2, = plt.plot(emg_model(coeffs[-1], t).detach().cpu(), 'r:')
plt.legend(handles=[dot, l1, l2], labels=['data_raw','gt_function', 'estimate_function'], loc='upper left')
plt.savefig('img/func.png')
plt.show()