import logging
from unittest import mock

import pytest

from client.globals import State
from client.strategy_manager import StrategyManager
from message.message import Side

logger = logging.getLogger(__name__)

DEFAULT_PRICE = 100.00
DEFAULT_QTY = 10
DEFAULT_SIDE = Side.Buy
DEFAULT_MSG_ID = "Test_Strategy"
DEFAULT_FILLED_QTY = 0
DEFAULT_SLICE_SIZE = 10
DEFAULT_REVISED_QTY = 20
DEFAULT_REVISED_PRICE = 150.00
DEFAULT_ORDER_ID = "1234"

@pytest.fixture
def strategy_manager():
    return StrategyManager(mock.Mock())

@pytest.fixture
def order_id(strategy_manager):
    strategy_manager.create_iceberg(
        side=DEFAULT_SIDE,
        quantity=DEFAULT_QTY,
        limit_price=DEFAULT_PRICE,
        slice_size=DEFAULT_SLICE_SIZE,
    )

    return next(iter(strategy_manager.orders.keys()))

@pytest.fixture
def given_acked_order(exchange_create_response):
    strategy_manager.on_create_resp(exchange_create_response)

@pytest.fixture
def fill_order_id(strategy_manager, order_id):
    slice = strategy_manager.orders[order_id]["iceberg_order"]
    strategy_manager.orders[order_id]["state"] = State.Filled
    strategy_manager.orders[order_id]["filled_quantity"] = 10
    slice.last_slice_state = State.Filled
    slice.parent_slice_state = State.Filled
    slice.slice_order_id = DEFAULT_ORDER_ID
    slice.slice_filled_quantity = 10
    slice.filled_quantity = 10
    strategy_manager.orders[order_id]["iceberg_order"] = slice

    return next(iter(strategy_manager.orders.keys()))

@pytest.fixture
def exchange_create_response(strategy_manager, order_id):
    slice = strategy_manager.orders[order_id]["iceberg_order"]
    return {
        "name": "OrderResponse",
        "client_msg_id": slice.slice_message_id,
        "client_id": "ceaba6cb06ee4e5ba39e84af4d01a55c",
        "order_params": {
            "limit_price": DEFAULT_PRICE,
            "quantity": DEFAULT_SLICE_SIZE,
            "side": "buy",
            "symbol": "AUTOTRAD Equity",
            "filled_quantity": 0,
            "exch_order_id": DEFAULT_ORDER_ID,
            "status": "ack",
        },
        "status": True,
        "status_msg": "Successful order creation",
    }

@pytest.fixture
def exchange_revise_response(strategy_manager, order_id):
    slice = strategy_manager.orders[order_id]["iceberg_order"]
    return {
        "name": "OrderResponse",
        "client_msg_id": slice.slice_message_id,
        "client_id": "369ceee4969c4b50bef16fb2d6652ab2",
        "order_params": {
            "limit_price": 100,
            "quantity": 5,
            "side": "buy",
            "symbol": "AUTOTRAD Equity",
            "filled_quantity": 0,
            "exch_order_id": DEFAULT_ORDER_ID,
            "status": "ack",
        },
        "status": True,
        "status_msg": "Revise order successful",
    }

@pytest.fixture
def exchange_cancel_response(strategy_manager, order_id):
    slice = strategy_manager.orders[order_id]["iceberg_order"]
    return {
        "name": "OrderResponse",
        "client_msg_id": slice.slice_message_id,
        "client_id": "9a440016815e415e9f8fa474af383703",
        "order_params": {
            "limit_price": DEFAULT_PRICE,
            "quantity": DEFAULT_SLICE_SIZE,
            "side": "buy",
            "symbol": "AUTOTRAD Equity",
            "filled_quantity": 0,
            "exch_order_id": DEFAULT_ORDER_ID,
            "status": "cancelled",
        },
        "status": True,
        "status_msg": "Order cancellation is successful",
    }

@pytest.fixture
def given_order_is_filled(strategy_manager, order_id):
    mock_response = {
        {
            "name": "FillOrderResponse",
            "order_params": {
                "limit_price": 100,
                "quantity": 10,
                "side": "buy",
                "symbol": "AUTOTRAD Equity",
                "filled_quantity": 10,
                "exch_order_id": DEFAULT_ORDER_ID,
                "status": "filled",
            },
            "client_id": "2285deab75cb448ebf7d784482550861",
            "trade": {
                "quantity": 10,
                "limit_price": 100,
                "symbol": "AUTOTRAD Equity",
                "exch_order_id": DEFAULT_ORDER_ID,
                "trade_id": "FillId-1663719822.706924",
                "fill_type": "Complete Fill",
                "side": "buy",
            },
            "status": True,
            "status_msg": "Order filled successfully",
        }
    }
    strategy_manager.on_fill_response(mock_response)
    parent_order_id = next(iter(strategy_manager.orders.keys()))
    return parent_order_id

def test_ordercreation(strategy_manager, order_id):
    # then
    assert strategy_manager.orders[order_id]["state"] == State.Sent
    assert strategy_manager.orders[order_id]["iceberg_order"]
    assert strategy_manager.orders[order_id]["limit_price"] == DEFAULT_PRICE
    assert strategy_manager.orders[order_id]["side"] == DEFAULT_SIDE
    assert strategy_manager.orders[order_id]["quantity"] == DEFAULT_QTY
    assert strategy_manager.orders[order_id]["filled_quantity"] == DEFAULT_FILLED_QTY

def test_on_create_resp(exchange_create_response, order_id, strategy_manager):
    # When
    strategy_manager.on_create_resp(exchange_create_response)

    # Then
    assert (
        strategy_manager.orders[order_id]["state"]
        == strategy_manager.orders[order_id]["iceberg_order"].last_slice_state
    )

def test_revise_pass(exchange_create_response, order_id, strategy_manager):
    # Given
    strategy_manager.on_create_resp(exchange_create_response)

    # When
    print(strategy_manager.orders)
    strategy_manager.revise(
        order_id=order_id,
        revised_quantity=DEFAULT_REVISED_QTY,
        revised_limit_price=DEFAULT_REVISED_PRICE,
    )

    # Then
    assert strategy_manager.orders[order_id]["quantity"] == DEFAULT_REVISED_QTY
    assert strategy_manager.orders[order_id]["limit_price"] == DEFAULT_REVISED_PRICE

def test_revise_fail(order_id, strategy_manager, caplog, exchange_create_response):

    # When
    strategy_manager.revise(
        order_id=order_id,
        revised_quantity=DEFAULT_REVISED_QTY,
        revised_limit_price=DEFAULT_PRICE,
    )

    # Then
    assert "is of State.Sent state and can not be revised" in caplog.text

def test_revise_order_qty_less_order_filled_fail(
    exchange_create_response, strategy_manager, fill_order_id, caplog
):
    # Given
    strategy_manager.on_create_resp(exchange_create_response)

    strategy_manager.revise(
        order_id=fill_order_id, revised_quantity=5, revised_limit_price=DEFAULT_PRICE
    )
    # Then
    assert "Can not update quantity to 5, already filled 10" in caplog.text

def test_cancel(
    order_id, strategy_manager, exchange_create_response, exchange_cancel_response
):
    strategy_manager.on_create_resp(exchange_create_response)

    # When
    strategy_manager.cancel(order_id)

    # Then
    assert strategy_manager.orders[order_id]["state"] == State.CancelSent

    strategy_manager.on_cancel_resp(exchange_cancel_response)

    assert strategy_manager.orders[order_id]["state"] == State.Cancelled


import matplotlib
matplotlib.use('Agg')
import numpy as np

CM_TO_M = 100

def show3Dpose(channels, ax, gt, mm=True):
    vals = channels.reshape((16, 3))
    if mm:
        channels *= CM_TO_M
    if gt:
        color = "#3498db"
    else:
        color = "#e74c3c"
    I = np.array([0, 2, 3, 5, 6, 8, 9, 10, 12, 13, 14, 2, 5, 2])  # start points
    J = np.array([1, 3, 4, 6, 7, 9, 10, 11, 13, 14, 15, 8, 12, 5])  # end points

    for i in range(16):
        ax.scatter(vals[i][0], vals[i][1], vals[i][2], c=color)

    for i in np.arange(len(I)):
        x, y, z = [np.array([vals[I[i], j], vals[J[i], j]]) for j in range(3)]
        ax.plot(x, y, z, lw=2, c=color)

    RADIUS = 50  # space around the subject
    xroot, yroot, zroot = vals[8, 0], vals[8, 1], vals[8, 2]
    ax.set_xlim3d([-RADIUS + xroot, RADIUS + xroot])
    ax.set_zlim3d([RADIUS + zroot, -RADIUS + zroot])
    ax.set_ylim3d([-RADIUS + yroot, RADIUS + yroot])

    # Get rid of the ticks and tick labels
    ax.set_xticks([])
    ax.set_yticks([])
    ax.set_zticks([])

    # ax.get_xaxis().set_ticklabels([])
    # ax.get_yaxis().set_ticklabels([])
    ax.set_xticklabels([])
    ax.set_yticklabels([])
    ax.set_zticklabels([])
    ax.set_aspect('equal')

    # Get rid of the panes (actually, make them white)
    white = (1.0, 1.0, 1.0, 0.0)
    ax.w_xaxis.set_pane_color(white)
    ax.w_yaxis.set_pane_color(white)
    # Keep z pane

    # Get rid of the lines in 3d
    ax.w_xaxis.line.set_color(white)
    ax.w_yaxis.line.set_color(white)
    ax.w_zaxis.line.set_color(white)


import base64
import time
import threading
from datetime import datetime, timedelta
import xml.etree.ElementTree as ET
from typing import List, Tuple, Dict, Optional, Any

import mpegdash.parser
import isodate

import Utils as Utils
from CommonTypes import EosNames, EosHttpConfig, EosFragment, EosUrl, LiveDelayListener
from Languages import EosLanguage
from DashUtils import DashInitDecoder
from RequestWrapper import RequestWrapper

####################################################
#  DashLiveDelayStream
####################################################
class DashLiveDelayStream:
    time_in_current_manifest: float
    current_time: float
    fragments: List[EosFragment]
    time_in_fragments: float
    max_timestamp: int
    time_scale: float
    media: str
    presentation_time_offset: int
    audio_sampling_rate: int

    ####################################################
    #  __init__
    ####################################################
    def __init__(self) -> None:
        self.max_timestamp = -1
        self.time_in_current_manifest = 0
        self.current_time = 0
        self.fragments = []
        self.time_in_fragments = 0
        self.time_scale = 1
        self.media = ''
        self.presentation_time_offset = 0
        self.audio_sampling_rate = 0

####################################################
#
#  DashLiveDelayHandler
#
####################################################
class DashLiveDelayHandler(threading.Thread):
    __session_id: str
    __live_origin_manifest_url: str
    __delay_seconds: float
    __first_manifest_read: bool
    __mpd_buffer_time_set: bool
    __reference_adaptation_set_id: Optional[str]
    __streams: Dict[Any, DashLiveDelayStream]  # (content_type, adaptation_set_id) -> DashLiveDelayStream
    __eos_streams: List[int]  # adaptation_set_id
    __listeners: List[Tuple[LiveDelayListener, str]]   # (LiveDelayListener, param)
    __mpd: Optional[mpegdash.nodes.MPEGDASH]
    __base_urls: List[str]
    __lock: threading.RLock
    #__mpd: Optional[ET.Element]
    __ready: threading.Event
    __request_wrapper: RequestWrapper

    ####################################################
    #  __init__
    ####################################################
    def __init__(self, session_id: str, live_origin_manifest_url: str, delay_seconds: int, ready: threading.Event) -> None:

        self.__session_id = session_id

        self.__live_origin_manifest_url = live_origin_manifest_url

        Utils.logger_.info(self.__session_id, "DashLiveDelayHandler::__init__ live_origin_manifest_url={}".format(self.__live_origin_manifest_url))

        self.__request_wrapper = RequestWrapper(self.__session_id, 'DashLiveDelayHandler ' + self.__live_origin_manifest_url)

        self.__delay_seconds = delay_seconds

        self.__first_manifest_read = True
        self.__mpd_buffer_time_set = False

        self.__reference_adaptation_set_id = None

        self.__streams = {}
        self.__eos_streams = []

        self.__listeners = []

        self.__mpd = None

        self.__base_urls = []

        self.__ready = ready

        self.__lock = threading.RLock()

        threading.Thread.__init__(self)  # start thread

    ####################################################
    #  set_reference_stream
    ####################################################
    def set_reference_stream(self, reference_adaptation_set_id: int, next_adaptation_set_id: int, dst_language: EosLanguage, default_lang: EosLanguage):

        Utils.logger_.info('DashLiveDelayHandler', "DashLiveDelayHandler::set_reference_stream reference_adaptation_set_id={}, next_adaptation_set_id={}, dst_language={}".format(reference_adaptation_set_id, next_adaptation_set_id, dst_language.code_bcp_47()))

        self.__reference_adaptation_set_id = reference_adaptation_set_id

        #reference_adaptation_set = None
        #period = self.__mpd.periods[0]
        #for ref_adaptation_set in period.adaptation_sets:
        #    if ref_adaptation_set.id == reference_adaptation_set_id:
        #        reference_adaptation_set = ref_adaptation_set

        # create new adaptation set for the dst_language
        adaptation_set = mpegdash.nodes.AdaptationSet()
        adaptation_set.id = next_adaptation_set_id
        adaptation_set.group = 3 # 8
        #adaptation_set.bitstream_switching = True
        #adaptation_set.segment_alignment = True
        adaptation_set.content_type = "text"
        adaptation_set.codecs = "stpp"
        adaptation_set.mime_type = "application/mp4"
        adaptation_set.start_with_sap = 1
        adaptation_set.lang = dst_language.code_639_1()

        if True:
            adaptation_set.id = None
            adaptation_set.content_type = None
            adaptation_set.codecs = None

            content_component = mpegdash.nodes.ContentComponent()
            content_component.id = '1'
            content_component.content_type = "text"
            adaptation_set.content_components = []
            adaptation_set.content_components.append(content_component)

        role = mpegdash.nodes.Descriptor()
        role.scheme_id_uri = "urn:mpeg:dash:role:2011"
        role.value = "subtitle"
        adaptation_set.roles = [role]

        representation = mpegdash.nodes.Representation()
        representation.id = "eos"
        representation.bandwidth = 8000

        if True:
            representation.codecs = "stpp"
            representation.id = next_adaptation_set_id + 2

            sub_representation = mpegdash.nodes.SubRepresentation()
            sub_representation.bandwidth = 8000
            sub_representation.codecs = "stpp"
            sub_representation.content_component = 1
            representation.sub_representations = sub_representation

        adaptation_set.representations = [representation]

        # reference_segment_template = reference_adaptation_set.segment_templates[0]

        segment_template = mpegdash.nodes.SegmentTemplate()
        segment_template.timescale = 1000  # reference_segment_template.timescale
        segment_template.media = "{}/{}/$Time$".format(EosNames.fragment_dash_prefix, dst_language.code_bcp_47())
        segment_template.initialization = "{}/{}/Init".format(EosNames.fragment_dash_prefix, dst_language.code_bcp_47())

        segment_timeline = mpegdash.nodes.SegmentTimeline()

        segment_template.segment_timelines = [segment_timeline]
        adaptation_set.segment_templates = [segment_template]

        self.__mpd.periods[0].adaptation_sets.append(adaptation_set)
        self.__eos_streams.append(next_adaptation_set_id)

    ####################################################
    #  get_presentation_time_offset
    ####################################################
    def get_presentation_time_offset(self) -> int:
        return self.__streams[('audio', self.__reference_adaptation_set_id)].presentation_time_offset

    ####################################################
    # run
    # called from thread context when start() is called
    ####################################################
    def run(self) -> None:

        Utils.logger_.system('DashLiveDelayHandler', "DashLiveDelayHandler::run thread started name={}".format(self.getName()))

        while True:

            original_manifest = None
            response = self.__request_wrapper.get(self.__live_origin_manifest_url)
            if response is None:
                Utils.logger_.error(str(self.__session_id), "DashLiveDelayHandler::run error getting manifest from server")
                time.sleep(1)
                continue

            original_manifest = response.text

            # print(original_manifest)

            self.__lock.acquire()

            mpd = mpegdash.parser.MPEGDASHParser.parse(original_manifest)

            if self.__mpd is None:
                self.__mpd = mpd

            if mpd.base_urls is not None:
                self.__base_urls.append(mpd.base_urls[0].base_url_value)
                mpd.base_urls = None

            period = mpd.periods[0]

            if period.base_urls is not None:
                self.__base_urls.append(period.base_urls[0].base_url_value)
                period.base_urls = None

            for adaptation_set in period.adaptation_sets:

                adaptation_set_id = adaptation_set.id
                content_type = adaptation_set.content_type

                if adaptation_set_id is None:
                    if adaptation_set.content_components is not None:
                        adaptation_set_id = adaptation_set.content_components[0].id
                        content_type = adaptation_set.content_components[0].content_type

                stream_key = (content_type, adaptation_set_id)
                # print("stream_key=", stream_key)

                if stream_key not in self.__streams.keys():
                    self.__streams[stream_key] = DashLiveDelayStream()

                if adaptation_set.audio_sampling_rate is not None:
                    self.__streams[stream_key].audio_sampling_rate = int(adaptation_set.audio_sampling_rate)

                self.__streams[stream_key].time_in_current_manifest = 0

                min_bandwidth = 1e12
                representation_id = ''
                for representation in adaptation_set.representations:
                    if representation.audio_sampling_rate is not None:
                        self.__streams[stream_key].audio_sampling_rate = int(representation.audio_sampling_rate)
                    if representation.bandwidth < min_bandwidth:
                        min_bandwidth = representation.bandwidth
                        representation_id = representation.id

                #Utils.logger_.debug('DashLiveDelayHandler', "DashLiveDelayHandler::run adaptation_set_id={}, content_type={}, min_bandwidth={}, audio_sampling_rate={}".format(adaptation_set_id, content_type, min_bandwidth, self.__streams[adaptation_set_id].audio_sampling_rate))

                segment_template = adaptation_set.segment_templates[0]

                new_url = EosUrl()
                media = segment_template.media
                if len(self.__base_urls) > 0:
                    media = self.__base_urls[-1] + media
                new_url.set_url(media, self.__live_origin_manifest_url)
                segment_template.media = new_url.absolute_url

                new_url = EosUrl()
                initialization = segment_template.initialization
                if len(self.__base_urls) > 0:
                    initialization = self.__base_urls[-1] + initialization
                new_url.set_url(initialization, self.__live_origin_manifest_url)
                segment_template.initialization = new_url.absolute_url
                audio_init_url = segment_template.initialization

                if content_type == "audio" and self.__streams[stream_key].audio_sampling_rate == 0:
                    # we need to read audio init fragment to get the audio_sampling_rate
                    audio_init_url = audio_init_url.replace('$RepresentationID$', representation_id)
                    self.__streams[stream_key].audio_sampling_rate = self._read_audio_init(audio_init_url)

                #from urllib.parse import urlparse, urlunparse
                #base_url = mpegdash.nodes.BaseURL()
                #parsed_parent_url = urlparse(self.__live_origin_manifest_url)
                #new_path = parsed_parent_url.path[:parsed_parent_url.path.rfind('/') + 1]
                #base = urlunparse((parsed_parent_url.scheme,
                #                   parsed_parent_url.netloc,
                #                   new_path,
                #                   '',
                #                   '',
                #                   ''))
                #base_url.base_url_value = base
                #adaptation_set.base_urls = []
                #adaptation_set.base_urls.append(base_url)

                self.__streams[stream_key].time_scale = segment_template.timescale
                self.__streams[stream_key].media = segment_template.media

                if self.__first_manifest_read is True:
                    Utils.logger_.debug('DashLiveDelayHandler', "DashLiveDelayHandler::run adaptation_set_id={}, content_type={}, time_scale={}, media={}".format(adaptation_set_id, content_type, self.__streams[stream_key].time_scale, self.__streams[stream_key].media))

                segment_time_line = segment_template.segment_timelines[0]

                for s in segment_time_line.Ss:
                    #print(s.t, s.d, s.r)

                    duration: int = 0
                    repeat: int = 1

                    if s.t is not None:
                        current_timestamp = s.t

                    if s.d is not None:
                        duration = float(s.d)
                        # next_timestamp = current_timestamp + duration

                    if s.r is not None:
                        repeat = s.r + 1

                    for x in range(repeat):

                        self.__streams[stream_key].time_in_current_manifest += (duration / self.__streams[stream_key].time_scale)

                        if current_timestamp > self.__streams[stream_key].max_timestamp:

                            new_media = self.__streams[stream_key].media
                            new_media = new_media.replace('$Bandwidth$', str(min_bandwidth))
                            new_media = new_media.replace('$Time$', str(current_timestamp))
                            new_media = new_media.replace('$RepresentationID$', representation_id)

                            new_fragment = EosFragment()
                            new_fragment.url.set_url(new_media, self.__live_origin_manifest_url)
                            new_fragment.sampling_rate = self.__streams[stream_key].audio_sampling_rate
                            new_fragment.timestamp = current_timestamp
                            new_fragment.duration = (duration / self.__streams[stream_key].time_scale)
                            new_fragment.start_time = self.__streams[stream_key].current_time
                            new_fragment.first_read = self.__first_manifest_read
                            #if line_index == len(lines):
                            #    new_fragment.first_read = False

                            self.__streams[stream_key].fragments.append(new_fragment)
                            self.__streams[stream_key].time_in_fragments += (duration / self.__streams[stream_key].time_scale)
                            self.__streams[stream_key].max_timestamp = current_timestamp
                            self.__streams[stream_key].current_time += (duration / self.__streams[stream_key].time_scale)
                            if self.__streams[stream_key].presentation_time_offset == 0:
                                self.__streams[stream_key].presentation_time_offset = new_fragment.timestamp / self.__streams[stream_key].time_scale

                            if content_type == "audio" and self.__reference_adaptation_set_id is not None and self.__reference_adaptation_set_id == adaptation_set_id:
                                self.__notify_listeners(new_fragment)

                            #print("+++++++++++ adaptation_set_id={}, fragments={}, time_in_fragments={}".format(adaptation_set_id, len(self.__streams[stream_key].fragments), self.__streams[stream_key].time_in_fragments))

                        current_timestamp += int(duration)

                if self.__streams[stream_key].time_in_current_manifest > 60.0:
                    if self.__first_manifest_read is True:
                        Utils.logger_.warning(self.__session_id, "DashLiveDelayHandler::run live manifest too long ({} seconds)".format(self.__streams[stream_key].time_in_current_manifest))
                    self.__streams[stream_key].time_in_current_manifest = 60.0

            if self.__first_manifest_read is True:
                # Utils.logger_.debug_color(self.__session_id, "DashLiveDelayHandler::run media_sequence={}".format(self.__base_media_sequence))
                self.__first_manifest_read = False

            self.__lock.release()

            self.__ready.set()

            time.sleep(1)

        Utils.logger_.system('DashLiveDelayHandler', "DashLiveDelayHandler::run thread ending name={}".format(self.getName()))

    ####################################################
    #  delay
    ####################################################
    def delay(self) -> Tuple[str, List[EosFragment]]:

        fragment_list: List[EosFragment] = []

        self.__lock.acquire()

        #copy_segment_time_line = None
        reference_timescale = 0
        reference_start_time = 0
        reference_duration = 0

        for stream in self.__streams:

            #print("++++++++++++++++++++++ stream: ", stream)

            start_index = 0
            end_index = -1

            #print("++++++++++++++++++++++ self.__streams[stream].time_in_fragments: ", self.__streams[stream].time_in_fragments)
            #print("++++++++++++++++++++++ self.__delay_seconds: ", self.__delay_seconds)
            #print("++++++++++++++++++++++ self.__streams[stream].time_in_current_manifest: ", self.__streams[stream].time_in_current_manifest)

            if self.__streams[stream].time_in_fragments >= self.__delay_seconds + self.__streams[stream].time_in_current_manifest:

                # create delay
                end_index = len(self.__streams[stream].fragments) - 1

                delay_time = 0
                for fragment in reversed(self.__streams[stream].fragments):

                    delay_time += fragment.duration
                    end_index -= 1

                    if delay_time >= self.__delay_seconds:
                        break

                # remove fragment from head
                time_in_fragment = self.__streams[stream].time_in_fragments

                for fragment in self.__streams[stream].fragments:

                    if time_in_fragment - delay_time > self.__streams[stream].time_in_current_manifest:

                        time_in_fragment -= fragment.duration
                        start_index += 1

                    else:
                        break

                #for i in range(start_index):
                #    self.__fragments.pop(0)
                #    self.__time_in_fragments -= fragment.duration

            #print("******************** start_index={}, end_index={}".format(start_index, end_index))

            period = self.__mpd.periods[0]
            for adaptation_set in period.adaptation_sets:

                adaptation_set_id = adaptation_set.id
                content_type = adaptation_set.content_type

                if adaptation_set_id is None:
                    if adaptation_set.content_components is not None:
                        adaptation_set_id = adaptation_set.content_components[0].id
                        content_type = adaptation_set.content_components[0].content_type

                if (content_type, adaptation_set_id) == stream:

                    segment_template = adaptation_set.segment_templates[0]

                    segment_time_line = segment_template.segment_timelines[0]

                    segment_time_line.Ss.clear()

                    #print("******************** 1 segment_time_line={}".format(segment_time_line))

                    first_timestamp = 0
                    duration = 0
                    if end_index > -1:
                        first = True
                        last_s = None
                        next_timestamp = 0
                        for fragment in self.__streams[stream].fragments[start_index - 1:end_index]:

                            #print("&&&&&&&&&&&&&&&&&&&& fragment={}".format(fragment))

                            s = mpegdash.nodes.S()

                            if first is True:
                                first = False
                                s.t = fragment.timestamp
                                first_timestamp = fragment.timestamp

                            repeated = False
                            if last_s is not None:
                                if int(last_s.d) == int(fragment.duration * self.__streams[stream].time_scale) and next_timestamp == fragment.timestamp:
                                    repeated = True
                                    duration += int(fragment.duration * self.__streams[stream].time_scale)
                                    if last_s.r is not None:
                                        last_s.r = last_s.r + 1
                                    else:
                                        last_s.r = 1

                            if repeated is False:

                                if next_timestamp != fragment.timestamp:
                                    s.t = fragment.timestamp

                                s.d = int(fragment.duration * self.__streams[stream].time_scale)
                                duration += int(fragment.duration * self.__streams[stream].time_scale)

                                #print("&&&&&&&&&&&&&&&&&&&&&& ", s.t, s.d, s.r)
                                segment_time_line.Ss.append(s)

                                last_s = s

                            next_timestamp = fragment.timestamp + int(fragment.duration * self.__streams[stream].time_scale)

                            fragment_list.append(fragment)

                    if adaptation_set_id == self.__reference_adaptation_set_id:
                        #copy_segment_time_line = segment_time_line
                        reference_timescale = self.__streams[stream].time_scale
                        reference_start_time = first_timestamp
                        reference_duration = duration

                    #print("******************** 2 segment_time_line={}".format(segment_time_line))

        for eos_stream in self.__eos_streams:
            eos_period = self.__mpd.periods[0]
            for eos_adaptation_set in eos_period.adaptation_sets:

                eos_adaptation_set_id = adaptation_set.id
                eos_content_type = adaptation_set.content_type

                if eos_adaptation_set_id is None:
                    if eos_adaptation_set.content_components is not None:
                        eos_adaptation_set_id = eos_adaptation_set.content_components[0].id
                        eos_content_type = eos_adaptation_set.content_components[0].content_type

                if eos_content_type == 'text':#  and eos_adaptation_set_id == eos_stream:
                    eos_segment_template = eos_adaptation_set.segment_templates[0]

                    #eos_segment_template.segment_timelines = [copy_segment_time_line]

                    #print("reference_duration=", reference_duration)

                    if reference_duration > 0:
                        eos_segment_time_line = eos_segment_template.segment_timelines[0]
                        eos_segment_time_line.Ss = []
                        s = mpegdash.nodes.S()
                        s.t = int(reference_start_time * 1000 / reference_timescale)
                        s.d = 4000
                        r = reference_duration / reference_timescale
                        s.r = int(r / 4)
                        eos_segment_time_line.Ss.append(s)

        #delay = timedelta(seconds=self.__delay_seconds)
        #self.__mpd.publish_time = (datetime.utcnow() - delay).strftime('%Y-%m-%dT%H:%M:%SZ')
        self.__mpd.publish_time = datetime.utcnow().strftime('%Y-%m-%dT%H:%M:%SZ')

        # handle mpd time
        if self.__mpd_buffer_time_set is False:
            #self.__mpd.suggested_presentation_delay = 'PT' + str(self.__delay_seconds) + 'S'
            self.__mpd.suggested_presentation_delay = isodate.duration_isoformat(isodate.Duration(seconds=self.__delay_seconds))

            time_shift_buffer_depth = isodate.parse_duration(self.__mpd.time_shift_buffer_depth)
            added_duration = isodate.Duration(seconds=self.__delay_seconds)
            time_shift_buffer_depth += added_duration
            self.__mpd.time_shift_buffer_depth = isodate.duration_isoformat(time_shift_buffer_depth)

            max_segment_duration = isodate.parse_duration(self.__mpd.max_segment_duration)
            subtitles_segment_duration = isodate.Duration(seconds=4)
            if subtitles_segment_duration.tdelta > max_segment_duration:
                self.__mpd.max_segment_duration = isodate.duration_isoformat(subtitles_segment_duration)

            self.__mpd_buffer_time_set = True

        manifest = mpegdash.parser.MPEGDASHParser.get_as_doc(self.__mpd).toxml()

        #print("manifest=", manifest)

        self.__lock.release()

        return manifest, fragment_list

    ####################################################
    #  _read_audio_init
    ####################################################
    def _read_audio_init(self, audio_init_url: str) -> int:

        audio_init_request_wrapper = RequestWrapper(self.__session_id, 'DashLiveDelayHandler ' + audio_init_url)

        audio_init = None
        response = audio_init_request_wrapper.get(audio_init_url)
        if response is None:
            Utils.logger_.error(str(self.__session_id), "DashLiveDelayHandler::_read_audio_init error getting audio init from server")
            return 0

        audio_init = response.content

        init_decoder = DashInitDecoder(audio_init)
        audio_sampling_rate = init_decoder.read_audio_sampling_rate()

        return audio_sampling_rate

    ####################################################
    #  register_live_parser_listener
    ####################################################
    def register_live_parser_listener(self, listener: LiveDelayListener, param: Optional[str]) -> None:

        self.__listeners.append((listener, param))

    ####################################################
    #  __notify_listeners
    ####################################################
    def __notify_listeners(self, fragment: EosFragment) -> None:

        for listener, param in self.__listeners:
            listener.on_new_fragment(fragment, param)


import re
import os
import matplotlib.pyplot as plt
from matplotlib.ticker import FuncFormatter
from collections import defaultdict
import statistics

directory = "./wrk"  # Replace with the actual directory path

requests_sec = defaultdict(list)
transfers_sec = defaultdict(list)

mean_requests = {}
mean_transfers = {}

def plot(kind='', title='', ylabel='', means=None):
    # Sort the labels and requests_sec lists together based on the requests_sec values
    labels = []
    values = []

    # silly, I know
    for k, v in means.items():
        labels.append(k)
        values.append(v)

    # sort the labels and value lists
    labels, values = zip(*sorted(zip(labels, values), key=lambda x: x[1], reverse=True))

    # Plot the graph
    plt.figure(figsize=(10, 6))  # Adjust the figure size as needed
    bars = plt.bar(labels, values)
    plt.xlabel("Subject")
    plt.ylabel(ylabel)
    plt.title(title)
    plt.xticks(rotation=45)  # Rotate x-axis labels for better readability

    # Display the actual values on top of the bars
    for bar in bars:
        yval = bar.get_height()
        plt.text(bar.get_x() + bar.get_width() / 2, yval, f'{yval:,.2f}', ha='center', va='bottom')

    plt.tight_layout()  # Adjust the spacing of the graph elements
    png_name = f"{directory}/{kind.lower()}_graph.png"
    plt.savefig(png_name)  # Save the graph as a PNG file
    print(f"Generated: {png_name}")

if __name__ == '__main__':
    if not os.path.isdir(".git"):
        print("Please run from root directory of the repository!")
        print("e.g. python wrk/graph.py")
        import sys
        sys.exit(1)

    # Iterate over the files in the directory
    for filename in os.listdir(directory):
        if filename.endswith(".perflog"):
            label = os.path.splitext(filename)[0]
            file_path = os.path.join(directory, filename)

            with open(file_path, "r") as file:
                lines = file.readlines()
                for line in lines:
                    # Extract the Requests/sec value using regular expressions
                    match = re.search(r"Requests/sec:\s+([\d.]+)", line)
                    if match:
                        requests_sec[label].append(float(match.group(1)))
                    match = re.search(r"Transfer/sec:\s+([\d.]+)", line)
                    if match:
                        value = float(match.group(1))
                        if 'KB' in line:
                            value *= 1024
                        elif 'MB' in line:
                            value *= 1024 * 1024
                        value /= 1024.0 * 1024
                        transfers_sec[label].append(value)

    # calculate means
    for k, v in requests_sec.items():
        mean_requests[k] = statistics.mean(v)

    for k, v in transfers_sec.items():
        mean_transfers[k] = statistics.mean(v)

    # save the plots
    plot(kind='req_per_sec', title='Requests/sec Comparison',
         ylabel='requests/sec', means=mean_requests)
    plot(kind='xfer_per_sec', title='Transfer/sec Comparison',
         ylabel='transfer/sec [MB]', means=mean_transfers)


"""Негативный тест"""
import pytest
import allure

from api.reqres.common import Common

@pytest.mark.usefixtures('clear_test_reports_and_logs')
@allure.severity('CRITICAL')
@allure.feature('Negative keys')
@allure.epic('API Тестирование портала "https://reqres.in/"')
@pytest.mark.parametrize('body_request',
                         [f'{Common.REQUEST_BODY_POST}',
                          f'{Common.REQUEST_BODY_POST2}'])
@allure.title('Неуспешная регистрация нового пользователя')
def test_register_unsuccessful_v1(base, body_request):

    response = base.api_v1.post_api_register(body_request)
    base.asserts.assert_register_unsuccessful(response=response, exp_status_code=400)


import os
import random
import shutil
import sys

import PyInstaller.__main__

def build(name, version, console, onefile, uac_admin, icon, files, folders):
	work_path = "build"
	while os.path.isdir(work_path):
		work_path = f"build_{random.randint(1, 1_000_000_000)}"
	work_path = os.path.join(os.path.abspath("."), work_path)

	result_path = os.path.abspath(".")

	run_list = ['main.py',
	            '--noconfirm',
	            '--clean',
	            '--name', f"{name}-v{version}",
	            '--workpath', work_path,
	            '--specpath', work_path,
	            '--distpath', result_path]

	if console:
		run_list.append("--console")
	else:
		run_list.append("--noconsole")

	if onefile:
		run_list.append("--onefile")
	else:
		run_list.append("--onedir")

	if uac_admin:
		run_list.append("--uac-admin")

	if icon != "":
		icon_path = os.path.join(os.path.abspath("."), icon)
		if not os.path.isfile(icon_path):
			raise Exception("Invalid icon!")
		else:
			run_list.extend(('--icon', icon_path))

	for file in files:
		if os.path.isfile(os.path.join(os.path.abspath("."), file)):
			run_list.extend(('--add-data', f'{os.path.join(os.path.abspath("."), file)};{os.path.dirname(file)}'))
		else:
			raise Exception("Invalid file!")

	for folder in folders:
		if os.path.isdir(folder):
			for walk in os.walk(folder, followlinks=False):
				for file in walk[2]:
					if os.path.isfile(os.path.join(walk[0], file)):
						run_list.extend(('--add-data', f'{os.path.join(os.path.abspath("."), os.path.join(walk[0], file))};{os.path.dirname(os.path.join(walk[0], file))}'))
					else:
						raise Exception("Invalid folder!")
		else:
			raise Exception("Invalid folder!")

	PyInstaller.__main__.run(run_list)
	shutil.rmtree(path=work_path, ignore_errors=True)

def main():
	name = "Full-Tilt!-Pinball"
	version = "4.0.2"

	console = False
	onefile = True
	uac_admin = False
	icon = "full-tilt-icon.ico"

	files = []
	folders = [
		"boards\\Alien",
		"boards\\Cadet",
		"boards\\Dragon",
		"boards\\Hero",
		"boards\\Mad",
		"boards\\Pirates",
	]

	if len(sys.argv) > 1 and sys.argv[1] == "--version":
		print(version)
	elif len(sys.argv) > 1 and sys.argv[1] == "--name":
		print(name)
	else:
		build(name, version, console, onefile, uac_admin, icon, files, folders)

if __name__ == '__main__':
	main()


import os
from setuptools import setup, Extension
from distutils.command.build_ext import build_ext as build_ext_orig

from distutils.command.build_ext import build_ext as build_ext_orig
class CTypesExtension(Extension): pass
class build_ext(build_ext_orig):

    def build_extension(self, ext):
        self._ctypes = isinstance(ext, CTypesExtension)
        return super().build_extension(ext)

    def get_export_symbols(self, ext):
        if self._ctypes:
            return ext.export_symbols
        return super().get_export_symbols(ext)

    def get_ext_filename(self, ext_name):
        if self._ctypes:
            return ext_name + '.so'
        return super().get_ext_filename(ext_name)

m_stc = CTypesExtension('hstego_stc_extension',
                  include_dirs = ['src/'],
                  sources = ['src/stc_interface.cpp',
                             'src/common.cpp',
                             'src/stc_embed_c.cpp',
                             'src/stc_extract_c.cpp',
                             'src/stc_ml_c.cpp'],
                  )

m_jpg = CTypesExtension('hstego_jpeg_toolbox_extension',
                  sources = ['src/jpeg_toolbox_extension.c'],
                  include_dirs = ['src/jpeg-9c-win/Include/'],
                  libraries = ['src/jpeg-9c-win/Lib/static_x64/jpeg'],
                  )

here = os.path.abspath(os.path.dirname(__file__))
with open(os.path.join(here, 'README.md'), encoding='utf-8') as f:
        long_description = f.read()

setup(name = 'hstego',
      version = '0.3',
      author="Daniel Lerch",
      author_email="dlerch@gmail.com",
      url="https://github.com/daniellerch/hstego",
      description = 'Hard to detect image steganography',
      py_modules = ["hstegolib", "hstegogui"],
      scripts = ['hstego.py'],
      data_files = [('resources', ['resources/hide.png']),
                    ('resources', ['resources/extract.png'])],
      ext_modules = [m_stc, m_jpg],
      cmdclass={'build_ext': build_ext})


import torch
from torch import nn
from torch.utils.data import DataLoader, Dataset
from torchvision import datasets, transforms
from torchinfo import summary

import os
import pathlib
import random
from PIL import Image
from pathlib import Path
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
from typing import Tuple, Dict, List
from tqdm.auto import tqdm
from timeit import default_timer as timer

def main():
    # Setup device-agnostic code
    device = "cuda" if torch.cuda.is_available() else "cpu"

    # Setup path to data folder
    data_path = Path("data/")
    image_path = data_path / "pizza_steak_sushi"

    # Setup train and testing paths
    train_dir = image_path / "train"
    test_dir = image_path / "test"

    # Set seed
    random.seed(42)

    # 1. Get all image paths (* means "any combination")
    image_path_list = list(image_path.glob("*/*/*.jpg"))

    # 2. Get random image path
    random_image_path = random.choice(image_path_list)

    # 3. Get image class from path name (the image class is the name of the directory where the image is stored)
    image_class = random_image_path.parent.stem

    # 4. Open image
    img = Image.open(random_image_path)

    # 5. Print metadata
    print(f"Random image path: {random_image_path}")
    print(f"Image class: {image_class}")
    print(f"Image height: {img.height}")
    print(f"Image width: {img.width}")

    # Open with local image opener
    # img.show()

    # We can also visualize using matplotlib

    # Turn the image into an array
    img_as_array = np.asarray(img)

    # Plot the image with matplotlib
    plt.figure(figsize=(10, 7))
    plt.imshow(img_as_array)
    plt.title(f"Image class: {image_class} | Image shape: {img_as_array.shape} -> [height, width, color_channels]")
    plt.axis(False)
    plt.show()

    # Write transform for image
    data_transform = transforms.Compose([
        # Resize the images to 64x64
        transforms.Resize(size=(64, 64)),
        # Flip the images randomly on the horizontal
        transforms.RandomHorizontalFlip(p=0.5), # p = probability of flip, 0.5 = 50% chance
        # Turn the image into a torch.Tensor
        transforms.ToTensor() # this also converts all pixel values from 0 to 255 to be between 0.0 and 1.0
    ])

    def plot_transformed_images(image_paths, transform, n=3, seed=42):
        """Plots a series of random images from image_paths.

        Will open n image paths from image_paths, transform them
        with transform and plot them side by side.

        Args:
            image_paths (list): List of target image paths.
            transform (PyTorch Transforms): Transforms to apply to images.
            n (int, optional): Number of images to plot. Defaults to 3.
            seed (int, optional): Random seed for the random generator. Defaults to 42.
        """
        random.seed(seed)
        random_image_paths = random.sample(image_paths, k=n)
        for image_path in random_image_paths:
            with Image.open(image_path) as f:
                fig, ax = plt.subplots(1, 2)
                ax[0].imshow(f)
                ax[0].set_title(f"Original \nSize: {f.size}")
                ax[0].axis("off")

                # Transform and plot image
                # Note: permute() will change shape of image to suit matplotlib
                # (PyTorch default is [C, H, W] but Matplotlib is [H, W, C])
                transformed_image = transform(f).permute(1, 2, 0)
                ax[1].imshow(transformed_image)
                ax[1].set_title(f"Transformed \nSize: {transformed_image.shape}")
                ax[1].axis("off")

                fig.suptitle(f"Class: {image_path.parent.stem}", fontsize=16)
                plt.show()

    plot_transformed_images(image_path_list,
                            transform=data_transform,
                            n=3)

    # Option 1: Loading Image Data Using ImageFolder
    train_data = datasets.ImageFolder(root=train_dir, # target folder of images
                                    transform=data_transform, # transforms to perform on data (images)
                                    target_transform=None) # transforms to perform on labels (if necessary)

    test_data = datasets.ImageFolder(root=test_dir,
                                    transform=data_transform)

    print(f"Train data:\n{train_data}\nTest data:\n{test_data}")

    # Get class names as a list
    class_names = train_data.classes

    img, label = train_data[0][0], train_data[0][1]
    print(f"Image tensor:\n{img}")
    print(f"Image shape: {img.shape}")
    print(f"Image datatype: {img.dtype}")
    print(f"Image label: {label}")
    print(f"Label datatype: {type(label)}")

    # Rearrange the order of dimensions
    img_permute = img.permute(1, 2, 0)

    # Print out different shapes (before and after permute)
    print(f"Original shape: {img.shape} -> [color_channels, height, width]")
    print(f"Image permute shape: {img_permute.shape} -> [height, width, color_channels]")

    # Plot the image
    plt.figure(figsize=(10, 7))
    plt.imshow(img.permute(1, 2, 0))
    plt.axis("off")
    plt.title(class_names[label], fontsize=14)
    # plt.show()

    train_dataloader = DataLoader(dataset=train_data,
                                batch_size=1, # how many samples per batch?
                                num_workers=1, # how many subprocesses to use for data loading? (higher = more)
                                shuffle=True) # shuffle the data?

    test_dataloader = DataLoader(dataset=test_data,
                                batch_size=1,
                                num_workers=1,
                                shuffle=False) # don't usually need to shuffle testing data

    print(train_dataloader, test_dataloader)

    # Option 2: Loading Image Data with a Custom Dataset

    # Make function to find classes in target directory
    def find_classes(directory: str) -> Tuple[List[str], Dict[str, int]]:
        """Finds the class folder names in a target directory.

        Assumes target directory is in standard image classification format.

        Args:
            directory (str): target directory to load classnames from.

        Returns:
            Tuple[List[str], Dict[str, int]]: (list_of_class_names, dict(class_name: idx...))

        Example:
            find_classes("food_images/train")
            >>> (["class_1", "class_2"], {"class_1": 0, ...})
        """
        # 1. Get the class names by scanning the target directory
        classes = sorted(entry.name for entry in os.scandir(directory) if entry.is_dir())

        # 2. Raise an error if class names not found
        if not classes:
            raise FileNotFoundError(f"Couldn't find any classes in {directory}.")

        # 3. Create a dictionary of index labels (computers prefer numerical rather than string labels)
        class_to_idx = {cls_name: i for i, cls_name in enumerate(classes)}
        return classes, class_to_idx

    print(find_classes(train_dir))

    # 1. Subclass torch.utils.data.Dataset
    class ImageFolderCustom(Dataset):

        # 2. Initialize with a targ_dir and transform (optional) parameter
        def __init__(self, targ_dir: str, transform=None) -> None:

            # 3. Create class attributes
            # Get all image paths
            self.paths = list(pathlib.Path(targ_dir).glob("*/*.jpg")) # note: you'd have to update this if you've got .png's or .jpeg's
            # Setup transforms
            self.transform = transform
            # Create classes and class_to_idx attributes
            self.classes, self.class_to_idx = find_classes(targ_dir)

        # 4. Make function to load images
        def load_image(self, index: int) -> Image.Image:
            "Opens an image via a path and returns it."
            image_path = self.paths[index]
            return Image.open(image_path)

        # 5. Overwrite the __len__() method (optional but recommended for subclasses of torch.utils.data.Dataset)
        def __len__(self) -> int:
            "Returns the total number of samples."
            return len(self.paths)

        # 6. Overwrite the __getitem__() method (required for subclasses of torch.utils.data.Dataset)
        def __getitem__(self, index: int) -> Tuple[torch.Tensor, int]:
            "Returns one sample of data, data and label (X, y)."
            img = self.load_image(index)
            class_name  = self.paths[index].parent.name # expects path in data_folder/class_name/image.jpeg
            class_idx = self.class_to_idx[class_name]

            # Transform if necessary
            if self.transform:
                return self.transform(img), class_idx # return data, label (X, y)
            else:
                return img, class_idx # return data, label (X, y)

    # Augment train data
    train_transforms = transforms.Compose([
        transforms.Resize((64, 64)),
        transforms.RandomHorizontalFlip(p=0.5),
        transforms.ToTensor()
    ])

    # Don't augment test data, only reshape
    test_transforms = transforms.Compose([
        transforms.Resize((64, 64)),
        transforms.ToTensor()
    ])

    train_data_custom = ImageFolderCustom(targ_dir=train_dir,
                                        transform=train_transforms)
    test_data_custom = ImageFolderCustom(targ_dir=test_dir,
                                        transform=test_transforms)

    # Check for equality amongst our custom Dataset and ImageFolder Dataset
    print((len(train_data_custom) == len(train_data)) & (len(test_data_custom) == len(test_data)))
    print(train_data_custom.classes == train_data.classes)
    print(train_data_custom.class_to_idx == train_data.class_to_idx)

    # 1. Take in a Dataset as well as a list of class names
    def display_random_images(dataset: torch.utils.data.dataset.Dataset,
                            classes: List[str] = None,
                            n: int = 10,
                            display_shape: bool = True,
                            seed: int = None):

        # 2. Adjust display if n too high
        if n > 10:
            n = 10
            display_shape = False
            print(f"For display purposes, n shouldn't be larger than 10, setting to 10 and removing shape display.")

        # 3. Set random seed
        if seed:
            random.seed(seed)

        # 4. Get random sample indexes
        random_samples_idx = random.sample(range(len(dataset)), k=n)

        # 5. Setup plot
        plt.figure(figsize=(16, 8))

        # 6. Loop through samples and display random samples
        for i, targ_sample in enumerate(random_samples_idx):
            targ_image, targ_label = dataset[targ_sample][0], dataset[targ_sample][1]

            # 7. Adjust image tensor shape for plotting: [color_channels, height, width] -> [color_channels, height, width]
            targ_image_adjust = targ_image.permute(1, 2, 0)

            # Plot adjusted samples
            plt.subplot(1, n, i+1)
            plt.imshow(targ_image_adjust)
            plt.axis("off")
            if classes:
                title = f"class: {classes[targ_label]}"
                if display_shape:
                    title = title + f"\nshape: {targ_image_adjust.shape}"
            plt.title(title)

        plt.show()

    # Display random images from ImageFolder created Dataset
    display_random_images(train_data,
                        n=5,
                        classes=class_names,
                        seed=None)

    # Display random images from ImageFolderCustom Dataset
    display_random_images(train_data_custom,
                        n=12,
                        classes=class_names,
                        seed=None) # Try setting the seed for reproducible images

    # Turn custom loaded images into DataLoader's
    train_dataloader_custom = DataLoader(dataset=train_data_custom, # use custom created train Dataset
                                        batch_size=1, # how many samples per batch?
                                        num_workers=0, # how many subprocesses to use for data loading? (higher = more)
                                        shuffle=True) # shuffle the data?

    test_dataloader_custom = DataLoader(dataset=test_data_custom, # use custom created test Dataset
                                        batch_size=1,
                                        num_workers=0,
                                        shuffle=False) # don't usually need to shuffle testing data

    train_transforms = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.TrivialAugmentWide(num_magnitude_bins=31), # how intense
        transforms.ToTensor() # use ToTensor() last to get everything between 0 & 1
    ])

    # Don't need to perform augmentation on the test data
    test_transforms = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor()
    ])

    # Get all image paths
    image_path_list = list(image_path.glob("*/*/*.jpg"))

    # Plot random images
    plot_transformed_images(
        image_paths=image_path_list,
        transform=train_transforms,
        n=3,
        seed=None
    )

    # Create simple transform
    simple_transform = transforms.Compose([
        transforms.Resize((64, 64)),
        transforms.ToTensor(),
    ])

    # 1. Load and transform data
    train_data_simple = datasets.ImageFolder(root=train_dir, transform=simple_transform)
    test_data_simple = datasets.ImageFolder(root=test_dir, transform=simple_transform)

    # 2. Turn data into DataLoaders

    # Setup batch size and number of workers
    BATCH_SIZE = 32
    # NUM_WORKERS = os.cpu_count() // 2
    NUM_WORKERS = 0
    print(f"Creating DataLoader's with batch size {BATCH_SIZE} and {NUM_WORKERS} workers.")

    # Create DataLoader's
    train_dataloader_simple = DataLoader(train_data_simple,
                                        batch_size=BATCH_SIZE,
                                        shuffle=True,
                                        num_workers=NUM_WORKERS)

    test_dataloader_simple = DataLoader(test_data_simple,
                                        batch_size=BATCH_SIZE,
                                        shuffle=False,
                                        num_workers=NUM_WORKERS)

    class TinyVGG(nn.Module):
        """
        Model architecture copying TinyVGG from:
        https://poloclub.github.io/cnn-explainer/
        """
        def __init__(self, input_shape: int, hidden_units: int, output_shape: int) -> None:
            super().__init__()
            self.conv_block_1 = nn.Sequential(
                nn.Conv2d(in_channels=input_shape,
                        out_channels=hidden_units,
                        kernel_size=3, # how big is the square that's going over the image?
                        stride=1, # default
                        padding=1), # options = "valid" (no padding) or "same" (output has same shape as input) or int for specific number
                nn.ReLU(),
                nn.Conv2d(in_channels=hidden_units,
                        out_channels=hidden_units,
                        kernel_size=3,
                        stride=1,
                        padding=1),
                nn.ReLU(),
                nn.MaxPool2d(kernel_size=2,
                            stride=2) # default stride value is same as kernel_size
            )
            self.conv_block_2 = nn.Sequential(
                nn.Conv2d(hidden_units, hidden_units, kernel_size=3, padding=1),
                nn.ReLU(),
                nn.Conv2d(hidden_units, hidden_units, kernel_size=3, padding=1),
                nn.ReLU(),
                nn.MaxPool2d(2)
            )
            self.classifier = nn.Sequential(
                nn.Flatten(),
                # Where did this in_features shape come from?
                # It's because each layer of our network compresses and changes the shape of our inputs data.
                nn.Linear(in_features=hidden_units*16*16,
                        out_features=output_shape)
            )

        def forward(self, x: torch.Tensor):
            # x = self.conv_block_1(x)
            # print(x.shape)
            # x = self.conv_block_2(x)
            # print(x.shape)
            # x = self.classifier(x)
            # print(x.shape)
            # return x
            return self.classifier(self.conv_block_2(self.conv_block_1(x))) # <- leverage the benefits of operator fusion

    torch.manual_seed(42)
    model_0 = TinyVGG(input_shape=3, # number of color channels (3 for RGB)
                    hidden_units=10,
                    output_shape=len(train_data.classes)).to(device)

    print(model_0)

    # 1. Get a batch of images and labels from the DataLoader
    img_batch, label_batch = next(iter(train_dataloader_simple))

    # 2. Get a single image from the batch and unsqueeze the image so its shape fits the model
    img_single, label_single = img_batch[0].unsqueeze(dim=0), label_batch[0]
    print(f"Single image shape: {img_single.shape}\n")

    # 3. Perform a forward pass on a single image
    model_0.eval()
    with torch.inference_mode():
        pred = model_0(img_single.to(device))

    # 4. Print out what's happening and convert model logits -> pred probs -> pred label
    print(f"Output logits:\n{pred}\n")
    print(f"Output prediction probabilities:\n{torch.softmax(pred, dim=1)}\n")
    print(f"Output prediction label:\n{torch.argmax(torch.softmax(pred, dim=1), dim=1)}\n")
    print(f"Actual label:\n{label_single}")

    summary(model_0, input_size=[1, 3, 64, 64]) # do a test pass through of an example input size

    def train_step(model: torch.nn.Module, dataloader: torch.utils.data.DataLoader, loss_fn: torch.nn.Module, optimizer: torch.optim.Optimizer):
        # Put model in train mode
        model.train()

        # Setup train loss and train accuracy values
        train_loss, train_acc = 0, 0

        # Loop through data loader data batches
        for batch, (X, y) in enumerate(dataloader):
            # Send data to target device
            X, y = X.to(device), y.to(device)

            # 1. Forward pass
            y_pred = model(X)

            # 2. Calculate  and accumulate loss
            loss = loss_fn(y_pred, y)
            train_loss += loss.item()

            # 3. Optimizer zero grad
            optimizer.zero_grad()

            # 4. Loss backward
            loss.backward()

            # 5. Optimizer step
            optimizer.step()

            # Calculate and accumulate accuracy metric across all batches
            y_pred_class = torch.argmax(torch.softmax(y_pred, dim=1), dim=1)
            train_acc += (y_pred_class == y).sum().item()/len(y_pred)

        # Adjust metrics to get average loss and accuracy per batch
        train_loss = train_loss / len(dataloader)
        train_acc = train_acc / len(dataloader)
        return train_loss, train_acc

    def test_step(model: torch.nn.Module, dataloader: torch.utils.data.DataLoader, loss_fn: torch.nn.Module):
        # Put model in eval mode
        model.eval()

        # Setup test loss and test accuracy values
        test_loss, test_acc = 0, 0

        # Turn on inference context manager
        with torch.inference_mode():
            # Loop through DataLoader batches
            for batch, (X, y) in enumerate(dataloader):
                # Send data to target device
                X, y = X.to(device), y.to(device)

                # 1. Forward pass
                test_pred_logits = model(X)

                # 2. Calculate and accumulate loss
                loss = loss_fn(test_pred_logits, y)
                test_loss += loss.item()

                # Calculate and accumulate accuracy
                test_pred_labels = test_pred_logits.argmax(dim=1)
                test_acc += ((test_pred_labels == y).sum().item()/len(test_pred_labels))

        # Adjust metrics to get average loss and accuracy per batch
        test_loss = test_loss / len(dataloader)
        test_acc = test_acc / len(dataloader)
        return test_loss, test_acc

    # 1. Take in various parameters required for training and test steps
    def train(model: torch.nn.Module, train_dataloader: torch.utils.data.DataLoader, test_dataloader: torch.utils.data.DataLoader, optimizer: torch.optim.Optimizer, loss_fn: torch.nn.Module = nn.CrossEntropyLoss(), epochs: int = 5):

        # 2. Create empty results dictionary
        results = {"train_loss": [],
            "train_acc": [],
            "test_loss": [],
            "test_acc": []
        }

        # 3. Loop through training and testing steps for a number of epochs
        for epoch in tqdm(range(epochs)):
            train_loss, train_acc = train_step(model=model,
                                            dataloader=train_dataloader,
                                            loss_fn=loss_fn,
                                            optimizer=optimizer)
            test_loss, test_acc = test_step(model=model,
                dataloader=test_dataloader,
                loss_fn=loss_fn)

            # 4. Print out what's happening
            print(
                f"Epoch: {epoch+1} | "
                f"train_loss: {train_loss:.4f} | "
                f"train_acc: {train_acc:.4f} | "
                f"test_loss: {test_loss:.4f} | "
                f"test_acc: {test_acc:.4f}"
            )

            # 5. Update results dictionary
            results["train_loss"].append(train_loss)
            results["train_acc"].append(train_acc)
            results["test_loss"].append(test_loss)
            results["test_acc"].append(test_acc)

        # 6. Return the filled results at the end of the epochs
        return results

    # Set random seeds
    torch.manual_seed(42)

    # Set number of epochs
    NUM_EPOCHS = 5

    # Recreate an instance of TinyVGG
    model_0 = TinyVGG(input_shape=3, # number of color channels (3 for RGB)
                    hidden_units=10,
                    output_shape=len(train_data.classes)).to(device)

    # Setup loss function and optimizer
    loss_fn = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(params=model_0.parameters(), lr=0.001)

    # Start the timer
    from timeit import default_timer as timer
    start_time = timer()

    # Train model_0
    model_0_results = train(model=model_0,
                            train_dataloader=train_dataloader_simple,
                            test_dataloader=test_dataloader_simple,
                            optimizer=optimizer,
                            loss_fn=loss_fn,
                            epochs=NUM_EPOCHS)

    # End the timer and print out how long it took
    end_time = timer()
    print(f"Total training time: {end_time-start_time:.3f} seconds")

    def plot_loss_curves(results: Dict[str, List[float]]):
        """Plots training curves of a results dictionary.

        Args:
            results (dict): dictionary containing list of values, e.g.
                {"train_loss": [...],
                "train_acc": [...],
                "test_loss": [...],
                "test_acc": [...]}
        """

        # Get the loss values of the results dictionary (training and test)
        loss = results['train_loss']
        test_loss = results['test_loss']

        # Get the accuracy values of the results dictionary (training and test)
        accuracy = results['train_acc']
        test_accuracy = results['test_acc']

        # Figure out how many epochs there were
        epochs = range(len(results['train_loss']))

        # Setup a plot
        plt.figure(figsize=(15, 7))

        # Plot loss
        plt.subplot(1, 2, 1)
        plt.plot(epochs, loss, label='train_loss')
        plt.plot(epochs, test_loss, label='test_loss')
        plt.title('Loss')
        plt.xlabel('Epochs')
        plt.legend()

        # Plot accuracy
        plt.subplot(1, 2, 2)
        plt.plot(epochs, accuracy, label='train_accuracy')
        plt.plot(epochs, test_accuracy, label='test_accuracy')
        plt.title('Accuracy')
        plt.xlabel('Epochs')
        plt.legend()
        plt.show()

    plot_loss_curves(model_0_results)

    # Create training transform with TrivialAugment
    train_transform_trivial_augment = transforms.Compose([
        transforms.Resize((64, 64)),
        transforms.TrivialAugmentWide(num_magnitude_bins=31),
        transforms.ToTensor()
    ])

    # Create testing transform (no data augmentation)
    test_transform = transforms.Compose([
        transforms.Resize((64, 64)),
        transforms.ToTensor()
    ])

    # Turn image folders into Datasets
    train_data_augmented = datasets.ImageFolder(train_dir, transform=train_transform_trivial_augment)
    test_data_simple = datasets.ImageFolder(test_dir, transform=test_transform)

    BATCH_SIZE = 32
    NUM_WORKERS = 0

    torch.manual_seed(42)
    train_dataloader_augmented = DataLoader(train_data_augmented,
                                            batch_size=BATCH_SIZE,
                                            shuffle=True,
                                            num_workers=NUM_WORKERS)

    test_dataloader_simple = DataLoader(test_data_simple,
                                        batch_size=BATCH_SIZE,
                                        shuffle=False,
                                        num_workers=NUM_WORKERS)

    # Create model_1 and send it to the target device
    torch.manual_seed(42)
    model_1 = TinyVGG(
        input_shape=3,
        hidden_units=10,
        output_shape=len(train_data_augmented.classes)).to(device)

    # Set random seeds
    torch.manual_seed(42)
    torch.cuda.manual_seed(42)

    # Set number of epochs
    NUM_EPOCHS = 5

    # Setup loss function and optimizer
    loss_fn = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(params=model_1.parameters(), lr=0.001)

    # Start the timer
    from timeit import default_timer as timer
    start_time = timer()

    # Train model_1
    model_1_results = train(model=model_1,
                            train_dataloader=train_dataloader_augmented,
                            test_dataloader=test_dataloader_simple,
                            optimizer=optimizer,
                            loss_fn=loss_fn,
                            epochs=NUM_EPOCHS)

    # End the timer and print out how long it took
    end_time = timer()
    print(f"Total training time: {end_time-start_time:.3f} seconds")

    plot_loss_curves(model_1_results)

    model_0_df = pd.DataFrame(model_0_results)
    model_1_df = pd.DataFrame(model_1_results)
    print(model_0_df)

    # Setup a plot
    plt.figure(figsize=(15, 10))

    # Get number of epochs
    epochs = range(len(model_0_df))

    # Plot train loss
    plt.subplot(2, 2, 1)
    plt.plot(epochs, model_0_df["train_loss"], label="Model 0")
    plt.plot(epochs, model_1_df["train_loss"], label="Model 1")
    plt.title("Train Loss")
    plt.xlabel("Epochs")
    plt.legend()

    # Plot test loss
    plt.subplot(2, 2, 2)
    plt.plot(epochs, model_0_df["test_loss"], label="Model 0")
    plt.plot(epochs, model_1_df["test_loss"], label="Model 1")
    plt.title("Test Loss")
    plt.xlabel("Epochs")
    plt.legend()

    # Plot train accuracy
    plt.subplot(2, 2, 3)
    plt.plot(epochs, model_0_df["train_acc"], label="Model 0")
    plt.plot(epochs, model_1_df["train_acc"], label="Model 1")
    plt.title("Train Accuracy")
    plt.xlabel("Epochs")
    plt.legend()

    # Plot test accuracy
    plt.subplot(2, 2, 4)
    plt.plot(epochs, model_0_df["test_acc"], label="Model 0")
    plt.plot(epochs, model_1_df["test_acc"], label="Model 1")
    plt.title("Test Accuracy")
    plt.xlabel("Epochs")
    plt.legend()
    plt.show()

    # Download custom image
    import requests

    # Setup custom image path
    custom_image_path = data_path / "04-pizza-dad.jpeg"

    # Download the image if it doesn't already exist
    if not custom_image_path.is_file():
        with open(custom_image_path, "wb") as f:
            # When downloading from GitHub, need to use the "raw" file link
            request = requests.get("https://raw.githubusercontent.com/mrdbourke/pytorch-deep-learning/main/images/04-pizza-dad.jpeg")
            print(f"Downloading {custom_image_path}...")
            f.write(request.content)
    else:
        print(f"{custom_image_path} already exists, skipping download.")

    # Read in custom image
    import torchvision
    custom_image_uint8 = torchvision.io.read_image(str(custom_image_path))

    # Print out image data
    print(f"Custom image tensor:\n{custom_image_uint8}\n")
    print(f"Custom image shape: {custom_image_uint8.shape}\n")
    print(f"Custom image dtype: {custom_image_uint8.dtype}")

    # Load in custom image and convert the tensor values to float32
    custom_image = torchvision.io.read_image(str(custom_image_path)).type(torch.float32)

    # Divide the image pixel values by 255 to get them between [0, 1]
    custom_image = custom_image / 255.

    # Print out image data
    print(f"Custom image tensor:\n{custom_image}\n")
    print(f"Custom image shape: {custom_image.shape}\n")
    print(f"Custom image dtype: {custom_image.dtype}")

    # Plot custom image
    plt.imshow(custom_image.permute(1, 2, 0)) # need to permute image dimensions from CHW -> HWC otherwise matplotlib will error
    plt.title(f"Image shape: {custom_image.shape}")
    plt.axis(False)
    plt.show()

    # Create transform pipleine to resize image
    custom_image_transform = transforms.Compose([
        transforms.Resize((64, 64)),
    ])

    # Transform target image
    custom_image_transformed = custom_image_transform(custom_image)

    # Print out original shape and new shape
    print(f"Original shape: {custom_image.shape}")
    print(f"New shape: {custom_image_transformed.shape}")

    model_1.eval()
    with torch.inference_mode():
        # Add an extra dimension to image
        custom_image_transformed_with_batch_size = custom_image_transformed.unsqueeze(dim=0)

        # Print out different shapes
        print(f"Custom image transformed shape: {custom_image_transformed.shape}")
        print(f"Unsqueezed custom image shape: {custom_image_transformed_with_batch_size.shape}")

        # Make a prediction on image with an extra dimension
        custom_image_pred = model_1(custom_image_transformed.unsqueeze(dim=0).to(device))

    # Print out prediction logits
    print(f"Prediction logits: {custom_image_pred}")

    # Convert logits -> prediction probabilities (using torch.softmax() for multi-class classification)
    custom_image_pred_probs = torch.softmax(custom_image_pred, dim=1)
    print(f"Prediction probabilities: {custom_image_pred_probs}")

    # Convert prediction probabilities -> prediction labels
    custom_image_pred_label = torch.argmax(custom_image_pred_probs, dim=1)
    print(f"Prediction label: {custom_image_pred_label}")

    # Find the predicted label
    custom_image_pred_class = class_names[custom_image_pred_label.cpu()] # put pred label to CPU, otherwise will error

    def pred_and_plot_image(model: torch.nn.Module,
                        image_path: str,
                        class_names: List[str] = None,
                        transform=None,
                        device: torch.device = device):
        """Makes a prediction on a target image and plots the image with its prediction."""

        # 1. Load in image and convert the tensor values to float32
        target_image = torchvision.io.read_image(str(image_path)).type(torch.float32)

        # 2. Divide the image pixel values by 255 to get them between [0, 1]
        target_image = target_image / 255.

        # 3. Transform if necessary
        if transform:
            target_image = transform(target_image)

        # 4. Make sure the model is on the target device
        model.to(device)

        # 5. Turn on model evaluation mode and inference mode
        model.eval()
        with torch.inference_mode():
            # Add an extra dimension to the image
            target_image = target_image.unsqueeze(dim=0)

            # Make a prediction on image with an extra dimension and send it to the target device
            target_image_pred = model(target_image.to(device))

        # 6. Convert logits -> prediction probabilities (using torch.softmax() for multi-class classification)
        target_image_pred_probs = torch.softmax(target_image_pred, dim=1)

        # 7. Convert prediction probabilities -> prediction labels
        target_image_pred_label = torch.argmax(target_image_pred_probs, dim=1)

        # 8. Plot the image alongside the prediction and prediction probability
        plt.imshow(target_image.squeeze().permute(1, 2, 0)) # make sure it's the right size for matplotlib
        if class_names:
            title = f"Pred: {class_names[target_image_pred_label.cpu()]} | Prob: {target_image_pred_probs.max().cpu():.3f}"
        else:
            title = f"Pred: {target_image_pred_label} | Prob: {target_image_pred_probs.max().cpu():.3f}"
        plt.title(title)
        plt.axis(False)
        plt.show()

    # Pred on our custom image
    pred_and_plot_image(model=model_1,
                        image_path=custom_image_path,
                        class_names=class_names,
                        transform=custom_image_transform,
                        device=device)

if __name__ == '__main__':
    main()


import logging
import time
from appium.webdriver.common import touch_action

class ProductPage():
    _langXpath = "//*[@text = 'English - EN']"
    _langSaveXpath = "//*[@text = 'Save Changes']"
    _prodDescXpath = "//*[@resource-id = 'title_feature_div']//android.view.View"
    _priceDescXpath = "//*[@resource-id = 'atfRedesign_priceblock_priceToPay']//android.widget.EditText"
    _addToCartXpath = "//*[@text = 'Add to Cart']"
    _cartXpath = "com.amazon.mShop.android.shopping:id/action_bar_cart_count"

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

    def gatherProductInfo(self):
        """Choosing English as default language and fetching the description of the choosen product"""
        logging.info("Choosing English as default language")
        time.sleep(10)
        self.driver.find_element_by_xpath(self._langXpath).click()
        self.driver.find_element_by_xpath(self._langSaveXpath).click()
        logging.info("Fetching the description of the choosen product from product description page")
        txt = self.driver.find_element_by_xpath(self._prodDescXpath).text
        productDescription = str(txt)
        time.sleep(2)
        txt1 = self.driver.find_element_by_xpath(self._priceDescXpath).text
        priceDescription = str(txt1)
        x = priceDescription[6:]
        finalRs = x.strip()

        return productDescription, finalRs

    def addToCart(self):
        """This method scrolls down the Add to cart button ,adds the product to cart and then navigate to cart page"""
        logging.info("Scrolling down the Add to cart button")
        touch = touch_action.TouchAction(self.driver)
        touch.press(x=1031, y=1296).move_to(x=1016, y=487).release().perform()
        touch.press(x=1025, y=1329).move_to(x=1007, y=408).release().perform()
        try:
            self.driver.find_element_by_xpath(self._addToCartXpath).click()
        except:
            touch.press(x=1025, y=1329).move_to(x=1007, y=408).release().perform()
            try:
                self.driver.find_element_by_xpath(self._addToCartXpath).click()
            except:
                raise Exception
        time.sleep(2)
        logging.info("Navigate to cart page")
        self.driver.find_element_by_id(self._cartXpath).click()


import os
import sys
sys.path.append('../')
import gc
import glob
import yaml
import datetime
import random

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as f
from torch.utils.data import DataLoader
import tensorboardX as tbx
from torchvision.models import resnet18
from fastprogress import progress_bar, master_bar

from models import *
from src.augmentation import get_test_augmentation, get_train_augmentation
from src.dataset import AlconDataset, KanaDataset
from src.metrics import *
from src.utils import *
from src.collates import *
from train_methods import *

#TODO:

def main():
    now = datetime.datetime.now()
    now_date = '{}-{:0>2d}-{:0>2d}_{:0>2d}-{:0>2d}-{:0>2d}'.format(now.year, now.month, now.day, now.hour, now.minute, now.second)
    print('{}-{}-{} {}:{}:{}'.format(now.year, now.month, now.day, now.hour, now.minute, now.second))
    with open('../params/exp0.yaml', "r+") as f:
        param = yaml.load(f, Loader=yaml.FullLoader)
    param['date'] =  now_date
    # seed set
    seed_setting(param['seed'])
    if torch.cuda.is_available():
        torch.backends.cudnn.benchmark = True

    fold = param['fold']
    outdir = os.path.join(param['save path'], str(os.path.basename(__file__).split('.')[-2]) + '_fold{}'.format(fold), now_date)
    if os.path.exists(param['save path']):
        os.makedirs(outdir, exist_ok=True)
    else:
        print("Not find {}".format(param['save path']))
        raise FileNotFoundError

    # outdir = '../tmp'

    # Dataset

    train_dataset = AlconDataset(df=get_train_df().query('valid != @fold'),
                                 augmentation=get_train_augmentation(),
                                 datadir=os.path.join(param['dataroot'],'train','imgs'), mode='train')
    valid_dataset = AlconDataset(df=get_train_df().query('valid == @fold'),
                                 augmentation=get_test_augmentation(),
                                 datadir=os.path.join(param['dataroot'],'train','imgs'), mode='valid')
    print('train dataset size: {}'.format(len(train_dataset)))
    print('valid dataset size: {}'.format(len(valid_dataset)))

    # Dataloader
    train_dataloader = DataLoader(train_dataset, batch_size=param['batch size'], num_workers=param['thread'],
                                  pin_memory=False, drop_last=False)
    valid_dataloader = DataLoader(valid_dataset, batch_size=param['batch size'], num_workers=param['thread'],
                                  pin_memory=False, drop_last=False)

    print('train loader size: {}'.format(len(train_dataloader)))
    print('valid loader size: {}'.format(len(valid_dataloader)))

    # model
    model = resnet18(pretrained=True)
    model.fc = nn.Linear(model.fc.in_features, 48)

    param['model'] = model.__class__.__name__

    # optim
    if param['optim'].lower() == 'sgd':
        optimizer = torch.optim.SGD(model.parameters(), lr=param['lr'], momentum=0.9,
                                    weight_decay=1e-5, nesterov=False)
    elif param['optim'].lower() == 'adam':
        optimizer =  torch.optim.SGD(model.parameters(), lr=param['lr'])
    else:
        raise NotImplementedError

    # scheduler
    scheduler = eval(param['scheduler'])

    model = model.to(param['device'])
    loss_fn = torch.nn.CrossEntropyLoss().to(param['device'])
    eval_fn = accuracy

    max_char_acc = 0.
    max_3char_acc = 0.
    min_loss = 10**5

    writer = tbx.SummaryWriter("../log/exp0")
    for key, val in param.items():
        # print(f'{key}: {val}')
        writer.add_text('data/hyperparam/{}'.format(key), str(val), 0)

    mb = master_bar(range(param['epoch']))
    for epoch in mb:
        avg_train_loss, avg_train_accuracy, avg_three_train_acc = train_alcon(model, optimizer, train_dataloader, param['device'],
                                       loss_fn, eval_fn, epoch, scheduler=None, writer=writer, parent=mb) #ok

        avg_valid_loss, avg_valid_accuracy, avg_three_valid_acc = valid_alcon(model, valid_dataloader, param['device'],
                                                                              loss_fn, eval_fn)

        writer.add_scalars("data/metric/valid", {
            'loss': avg_valid_loss,
            'accuracy': avg_valid_accuracy,
            '3accuracy': avg_three_valid_acc
        }, epoch)

        print('======================== epoch {} ========================'.format(epoch+1))
        print('lr              : {:.5f}'.format(scheduler.get_lr()[0]))
        print('loss            : train={:.5f}  , test={:.5f}'.format(avg_train_loss, avg_valid_loss))
        print('acc(per 1 char) : train={:.3%}  , test={:.3%}'.format(avg_train_accuracy, avg_valid_accuracy))
        print('acc(per 3 char) : train={:.3%}  , test={:.3%}'.format(avg_three_train_acc, avg_three_valid_acc))

        if min_loss > avg_valid_loss:
            print('update best loss:  {:.5f} ---> {:.5f}'.format(min_loss, avg_valid_loss))
            min_loss = avg_valid_loss
            torch.save(model.state_dict(), os.path.join(outdir, 'best_loss.pth'))

        if max_char_acc < avg_valid_accuracy:
            print('update best acc per 1 char:  {:.3%} ---> {:.3%}'.format(max_char_acc, avg_valid_accuracy))
            max_char_acc = avg_valid_accuracy
            torch.save(model.state_dict(), os.path.join(outdir, 'best_acc.pth'))

        if max_3char_acc < avg_three_valid_acc:
            print('update best acc per 3 char:  {:.3%} ---> {:.3%}'.format(max_3char_acc , avg_three_valid_acc))
            max_3char_acc = avg_three_valid_acc
            torch.save(model.state_dict(), os.path.join(outdir, 'best_3acc.pth'))

        if 1:
            if scheduler is not None:
                if writer is not None:
                    writer.add_scalar("data/learning rate", scheduler.get_lr()[0], epoch)
                scheduler.step()

    writer.add_scalars("data/metric/valid", {
        'best loss': min_loss,
        'best accuracy': max_char_acc,
        'best 3accuracy': max_3char_acc
    })

    print('finish train')
    print('result')
    print('best loss : {}'.format(min_loss))
    print('best 1 acc : {}'.format(max_char_acc))
    print('best 3 acc : {}'.format(max_3char_acc))
    writer.export_scalars_to_json(os.path.join(outdir, 'history.json'))
    writer.close()

    #TODO: gc.collect()
    del train_dataset, valid_dataset
    del train_dataloader, valid_dataloader
    del scheduler, optimizer
    gc.collect()

    print('load weight')
    model.load_state_dict(torch.load(os.path.join(outdir, 'best_3acc.pth')))

    test_dataset = AlconDataset(df=get_test_df(),
                                augmentation=get_test_augmentation(),
                                datadir=os.path.join(param['dataroot'], 'test', 'imgs'), mode='test')

    test_dataloader = DataLoader(test_dataset, batch_size=param['batch size'], num_workers=param['thread'],
                                 pin_memory=False, drop_last=False)
    print('test dataset size: {}'.format(len(test_dataset)))
    print('test loader size: {}'.format(len(test_dataloader)))

    output_list = pred_alcon(model, test_dataloader, param['device'])
    torch.save(output_list, os.path.join(outdir, 'prediction.pth'))
    pd.DataFrame(output_list).drop('logit', axis=1).sort_values('ID').set_index('ID').to_csv(os.path.join(outdir, 'submission.csv'))

# def prediction():
#     pass
#
# def train(model, optimizer, dataloader, device, loss_fn, eval_fn, epoch, scheduler=None, writer=None):
#     model.train()
#     avg_loss = 0
#     avg_accuracy = 0
#     for step, (inputs, targets) in enumerate(dataloader):
#         inputs = inputs.to(device)
#         targets = targets.to(device)
#         optimizer.zero_grad()
#         logits = model(inputs)
#         preds = logits.softmax(dim=1)
#         loss = loss_fn(logits, targets.argmax(dim=1))
#         loss.backward()
#         # nn.utils.clip_grad_norm_(model.parameters(), 1.0)
#         optimizer.step()
#         avg_loss += loss.item()
#         acc = eval_fn(preds, targets)
#         avg_accuracy += acc
#         if scheduler is not None:
#             if writer is not None:
#                 writer.add_scalar("data/learning rate", scheduler.get_lr()[0], step + epoch*len(dataloader))
#             scheduler.step()
#
#         writer.add_scalars("data/metric/train", {
#             'train_loss': loss.item(),
#             'train_accuracy': acc
#         }, step + epoch*len(dataloader))
#
#     avg_loss /= len(dataloader)
#     avg_accuracy /= len(dataloader)
#     return avg_loss, avg_accuracy
#
#
# def valid(model, dataloader, device, loss_fn, eval_fn):
#     model.eval()
#     avg_loss = 0
#     avg_accuracy = 0
#     with torch.no_grad():
#         for inputs, targets in dataloader:
#             inputs = inputs.to(device)
#             targets = targets.to(device)
#             logits = model(inputs)
#             preds = logits.softmax(dim=1)
#             loss = loss_fn(logits, targets.argmax(dim=1))
#             avg_loss += loss.item()
#             avg_accuracy += eval_fn(preds, targets)
#         avg_loss /= len(dataloader)
#         avg_accuracy /= len(dataloader)
#
#     return avg_loss, avg_accuracy
#
#
# def train_alcon(model, optimizer, dataloader, device, loss_fn, eval_fn, epoch, scheduler=None, writer=None, parent=None):
#     model.train()
#     avg_loss = 0
#     avg_accuracy = 0
#     three_char_accuracy = 0
#     for step, (inputs, targets, indices) in enumerate(progress_bar(dataloader, parent=parent)):
#         inputs = inputs.to(device)
#         targets = targets.to(device)
#         _avg_loss = 0
#         _avg_accuracy = 0
#         preds = torch.zeros(targets.size()).to(device)
#         for i in range(3):
#             # _inputs = inputs[:, i].to(device)
#             # _targets = targets[:, i].to(device)
#             _inputs = inputs[:, i]
#             _targets = targets[:, i]
#             optimizer.zero_grad()
#             logits = model(_inputs)
#             _preds = logits.softmax(dim=1)
#             preds[:, i] = _preds
#             loss = loss_fn(logits, _targets.argmax(dim=1))
#             loss.backward()
#             # nn.utils.clip_grad_norm_(model.parameters(), 1.0)
#             optimizer.step()
#             _avg_loss += loss.item()
#             acc = eval_fn(_preds, _targets.argmax(dim=1))
#             _avg_accuracy += acc.item()
#         _avg_loss /= 3
#         avg_loss += _avg_loss
#         _avg_accuracy /= 3
#         avg_accuracy += _avg_accuracy
#
#         _three_char_accuracy = accuracy_three_character(preds, targets.argmax(dim=2), mean=True).item()
#         # _three_char_accuracy = accuracy_three_character(pred, targets.to(device), mean=True)
#
#         three_char_accuracy += _three_char_accuracy
#         if scheduler is not None:
#             if writer is not None:
#                 writer.add_scalar("data/learning rate", scheduler.get_lr()[0], step + epoch*len(dataloader))
#             scheduler.step()
#
#         writer.add_scalars("data/metric/train", {
#             'loss': _avg_loss,
#             'accuracy': _avg_accuracy,
#             '3accuracy': _three_char_accuracy
#         }, step + epoch*len(dataloader))
#
#     avg_loss /= len(dataloader)
#     avg_accuracy /= len(dataloader)
#     three_char_accuracy /= len(dataloader)
#     print()
#     return avg_loss, avg_accuracy, three_char_accuracy
#
#
# def valid_alcon(model, dataloader, device, loss_fn, eval_fn):
#     model.eval()
#     with torch.no_grad():
#         avg_loss = 0
#         avg_accuracy = 0
#         three_char_accuracy = 0
#         for step, (inputs, targets, indices) in enumerate(dataloader):
#             inputs = inputs.to(device)
#             targets = targets.to(device)
#             _avg_loss = 0
#             _avg_accuracy = 0
#             preds = torch.zeros(targets.size()).to(device)
#             for i in range(3):
#                 # _inputs = inputs[:, i].to(device)
#                 # _targets = targets[:, i].to(device)
#                 _inputs = inputs[:, i]
#                 _targets = targets[:, i]
#                 logits = model(_inputs)
#                 _preds = logits.softmax(dim=1)
#                 preds[:, i] = _preds
#                 loss = loss_fn(logits, _targets.argmax(dim=1))
#                 _avg_loss += loss.item()
#                 acc = eval_fn(_preds, _targets.argmax(dim=1))
#                 _avg_accuracy += acc.item()
#             _avg_loss /= 3
#             avg_loss += _avg_loss
#             _avg_accuracy /= 3
#             avg_accuracy += _avg_accuracy
#
#             _three_char_accuracy = accuracy_three_character(preds, targets.argmax(dim=2), mean=True).item()
#             # _three_char_accuracy = accuracy_three_character(pred, targets.to(device), mean=True)
#
#             three_char_accuracy += _three_char_accuracy
#
#         avg_loss /= len(dataloader)
#         avg_accuracy /= len(dataloader)
#         three_char_accuracy /= len(dataloader)
#     return avg_loss, avg_accuracy, three_char_accuracy
#
#
# def train_mixup(model, optimizer, dataloader, device, loss_fn, eval_fn, epoch, scheduler=None, writer=None):
#     model.train()
#     avg_loss = 0
#     for step, (inputs, targets) in enumerate(dataloader):
#         inputs = inputs.to(device)
#         targets = targets.to(device)
#         inputs, targets_a, targets_b, lam = mixup_data(inputs, targets, alpha=0.2, device=device)
#         optimizer.zero_grad()
#         logits = model(inputs)
#         loss = lam * loss_fn(logits, targets_a.argmax(dim=1))  + (1 - lam) * loss_fn(logits, targets_b.argmax(dim=1))
#         loss.backward()
#         nn.utils.clip_grad_norm_(model.parameters(), 1.0)
#         optimizer.step()
#         avg_loss += loss.item()
#         if scheduler is not None:
#             if writer is not None:
#                 writer.add_scalar("data/lr", scheduler.get_lr()[0], step + epoch*len(dataloader))
#             scheduler.step()
#         writer.add_scalars("data/metric/train_loss", loss.item(), step + epoch * len(dataloader))
#     avg_loss /= len(dataloader)
#
#     return avg_loss

if __name__ =='__main__':
    main()


import os
import argparse
from shutil import rmtree
from subprocess import run, PIPE, STDOUT

class CourseMount:

    def __init__(self):
        self.courses = [
            "Linux_course/Linux_course1",
            "Linux_course/Linux_course2",
            "machinelearning/ML_course1",
            "machinelearning/ML_course2",
            "SQLFundamentals1",
            "SQLFundamentals2",
            "SQLFundamentals3"
        ]
        self.source = '/home/rinkesh/scripts/source/courses'
        self.target = '/home/rinkesh/scripts/target/courses'

    def check_course(self, course):
        return course in self.courses

    def check_mount(self, course):
        course_path = os.path.join(self.target, course)

        #check mount
        if os.path.isdir(course_path):
            mount_check = run(['mountpoint', course_path], stdout=PIPE, stderr=STDOUT)
            #print(mount_check)
            return True if mount_check.returncode == 0 else False
        return False

    def mount_course(self, course):
        if not self.check_course(course):
            print(f'Incorrect {course} name given. Try again')
            return 1
        if self.check_mount(course):
            print("{course} Already mounted.")
            return 0

        #make target directories
        target_course_path = os.path.join(self.target, course)
        source_course_path = os.path.join(self.source, course)
        os.makedirs(target_course_path)

        #mount source to target
        result = run(['bindfs', '-p', '550', '-u', 'rinkesh', '-g',
                        'rinkesh', source_course_path, target_course_path],
                         stdout=PIPE, stderr=STDOUT)
        #check if mounted successfully
        if result.returncode != 0:
            rmtree(target_course_path)
            print("Mounting failed.")
            return 1

        print(f'{course} Mounted Successfully.')
        return 0

    def mount_all(self):
        for course in self.courses:
            self.mount_course(course)

    def unmount_course(self, course):
        if not self.check_course(course):
            print(f'Incorrect {course} name given. Try again')
            return 1
        if not self.check_mount(course):
            print(f'{course} is either unmounted or directory does not exists')
            return 0
        else:
            unmount_path = os.path.join(self.target, course)
            result = run(['sudo', 'umount', unmount_path], stdout=PIPE, stderr=STDOUT)
            if result.returncode != 0:
                print("Unmounting Error")
                return 1
            else:
                rmtree(unmount_path, ignore_errors=True)
                print(f'{course} Unmounted Successfully')
                return 0

    def unmount_all(self):
        for course in self.courses:
            self.unmount_course(course)

if __name__ == '__main__':
    #adding CLI
    parser = argparse.ArgumentParser()

    parser.add_argument('-m', '--mount', action='store_true', help='for mounting courses')
    parser.add_argument('-u', '--unmount', action='store_true', help='for unmounting courses')
    parser.add_argument('-c', '--course', type=str, help='give the course name to mount/unmount' )

    args = parser.parse_args()

    #calling CurseMount
    mount_obj = CourseMount()

    if args.mount:
        if args.course:
            mount_obj.mount_course(args.course)
        else:
            mount_obj.mount_all()
    elif args.unmount:
        if args.course:
            mount_obj.unmount_course(args.course)
        else:
            mount_obj.unmount_all()
    else:
        print('Read the docs with -h flag')