import asyncio
import math
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple

from _decimal import Decimal

from hummingbot.connector.constants import s_decimal_NaN
from hummingbot.connector.exchange.polkadex import polkadex_constants as CONSTANTS, polkadex_web_utils as web_utils
from hummingbot.connector.exchange.polkadex.polkadex_api_order_book_data_source import PolkadexAPIOrderBookDataSource
from hummingbot.connector.exchange.polkadex.polkadex_data_source import PolkadexDataSource
from hummingbot.connector.exchange_py_base import ExchangePyBase
from hummingbot.connector.trading_rule import TradingRule
from hummingbot.connector.utils import get_new_client_order_id
from hummingbot.core.api_throttler.data_types import RateLimit
from hummingbot.core.data_type.common import OrderType, TradeType
from hummingbot.core.data_type.in_flight_order import InFlightOrder, OrderState, OrderUpdate, TradeUpdate
from hummingbot.core.data_type.order_book_tracker_data_source import OrderBookTrackerDataSource
from hummingbot.core.data_type.trade_fee import TokenAmount, TradeFeeBase
from hummingbot.core.data_type.user_stream_tracker_data_source import UserStreamTrackerDataSource
from hummingbot.core.event.event_forwarder import EventForwarder
from hummingbot.core.event.events import AccountEvent, BalanceUpdateEvent, MarketEvent
from hummingbot.core.network_iterator import NetworkStatus, safe_ensure_future
from hummingbot.core.utils.estimate_fee import build_trade_fee
from hummingbot.core.web_assistant.auth import AuthBase
from hummingbot.core.web_assistant.web_assistants_factory import WebAssistantsFactory

if TYPE_CHECKING:
    from hummingbot.client.config.config_helpers import ClientConfigAdapter  # pragma: no cover

class PolkadexExchange(ExchangePyBase):
    web_utils = web_utils

    def __init__(
        self,
        client_config_map: "ClientConfigAdapter",
        polkadex_seed_phrase: str,
        trading_pairs: Optional[List[str]] = None,
        trading_required: bool = True,
        domain: str = CONSTANTS.DEFAULT_DOMAIN,
    ):
        self._trading_required = trading_required
        self._trading_pairs = trading_pairs
        self._domain = domain
        self._data_source = PolkadexDataSource(
            connector=self, seed_phrase=polkadex_seed_phrase, domain=self._domain, trading_required=trading_required
        )
        super().__init__(client_config_map=client_config_map)
        self._data_source.configure_throttler(throttler=self._throttler)
        self._forwarders = []
        self._configure_event_forwarders()

    @property
    def name(self) -> str:
        return CONSTANTS.EXCHANGE_NAME

    @property
    def authenticator(self) -> AuthBase:
        return None

    @property
    def rate_limits_rules(self) -> List[RateLimit]:
        return CONSTANTS.RATE_LIMITS

    @property
    def domain(self) -> str:
        return self._domain

    @property
    def client_order_id_max_length(self) -> Optional[int]:
        return CONSTANTS.MAX_ID_LEN

    @property
    def client_order_id_prefix(self) -> str:
        return CONSTANTS.CLIENT_ID_PREFIX

    @property
    def trading_rules_request_path(self) -> str:
        raise NotImplementedError  # pragma: no cover

    @property
    def trading_pairs_request_path(self) -> str:
        raise NotImplementedError  # pragma: no cover

    @property
    def check_network_request_path(self) -> str:
        raise NotImplementedError  # pragma: no cover

    @property
    def trading_pairs(self) -> List[str]:
        return self._trading_pairs

    @property
    def is_cancel_request_in_exchange_synchronous(self) -> bool:
        return False

    @property
    def is_trading_required(self) -> bool:
        return self._trading_required

    async def start_network(self):
        await super().start_network()

        market_symbols = [
            await self.exchange_symbol_associated_to_pair(trading_pair=trading_pair)
            for trading_pair in self._trading_pairs
        ]
        await self._data_source.start(market_symbols=market_symbols)

    async def stop_network(self):
        """
        This function is executed when the connector is stopped. It performs a general cleanup and stops all background
        tasks that require the connection with the exchange to work.
        """
        await super().stop_network()
        await self._data_source.stop()

    def supported_order_types(self) -> List[OrderType]:
        return [OrderType.LIMIT, OrderType.MARKET]

    async def check_network(self) -> NetworkStatus:
        """
        Checks connectivity with the exchange using the API
        """
        try:
            status = await self._data_source.exchange_status()
        except asyncio.CancelledError:
            raise
        except Exception:
            status = NetworkStatus.NOT_CONNECTED
        return status

    # === Orders placing ===

    def buy(self,
            trading_pair: str,
            amount: Decimal,
            order_type=OrderType.LIMIT,
            price: Decimal = s_decimal_NaN,
            **kwargs) -> str:
        """
        Creates a promise to create a buy order using the parameters

        :param trading_pair: the token pair to operate with
        :param amount: the order amount
        :param order_type: the type of order to create (MARKET, LIMIT, LIMIT_MAKER)
        :param price: the order price

        :return: the id assigned by the connector to the order (the client id)
        """
        order_id = get_new_client_order_id(
            is_buy=True,
            trading_pair=trading_pair,
            hbot_order_id_prefix=self.client_order_id_prefix,
            max_id_len=self.client_order_id_max_length
        )
        hex_order_id = f"0x{order_id.encode('utf-8').hex()}"
        safe_ensure_future(self._create_order(
            trade_type=TradeType.BUY,
            order_id=hex_order_id,
            trading_pair=trading_pair,
            amount=amount,
            order_type=order_type,
            price=price,
            **kwargs))
        return hex_order_id

    def sell(self,
             trading_pair: str,
             amount: Decimal,
             order_type: OrderType = OrderType.LIMIT,
             price: Decimal = s_decimal_NaN,
             **kwargs) -> str:
        """
        Creates a promise to create a sell order using the parameters.
        :param trading_pair: the token pair to operate with
        :param amount: the order amount
        :param order_type: the type of order to create (MARKET, LIMIT, LIMIT_MAKER)
        :param price: the order price
        :return: the id assigned by the connector to the order (the client id)
        """
        order_id = get_new_client_order_id(
            is_buy=False,
            trading_pair=trading_pair,
            hbot_order_id_prefix=self.client_order_id_prefix,
            max_id_len=self.client_order_id_max_length
        )
        hex_order_id = f"0x{order_id.encode('utf-8').hex()}"
        safe_ensure_future(self._create_order(
            trade_type=TradeType.SELL,
            order_id=hex_order_id,
            trading_pair=trading_pair,
            amount=amount,
            order_type=order_type,
            price=price,
            **kwargs))
        return hex_order_id

    def _is_request_exception_related_to_time_synchronizer(self, request_exception: Exception) -> bool:
        # Polkadex does not use a time synchronizer
        return False

    def _is_order_not_found_during_status_update_error(self, status_update_exception: Exception) -> bool:
        return CONSTANTS.ORDER_NOT_FOUND_MESSAGE in str(status_update_exception)

    def _is_order_not_found_during_cancelation_error(self, cancelation_exception: Exception) -> bool:
        return CONSTANTS.ORDER_NOT_FOUND_MESSAGE in str(cancelation_exception)

    async def _execute_order_cancel_and_process_update(self, order: InFlightOrder) -> bool:
        new_order_state = await self._place_cancel(order.client_order_id, order)
        cancelled = new_order_state in [OrderState.CANCELED, OrderState.PENDING_CANCEL]
        if cancelled:
            update_timestamp = self.current_timestamp
            if update_timestamp is None or math.isnan(update_timestamp):
                update_timestamp = self._time()
            order_update: OrderUpdate = OrderUpdate(
                client_order_id=order.client_order_id,
                trading_pair=order.trading_pair,
                update_timestamp=update_timestamp,
                new_state=new_order_state,
            )
            self._order_tracker.process_order_update(order_update)
        return cancelled

    async def _place_cancel(self, order_id: str, tracked_order: InFlightOrder) -> OrderState:
        await tracked_order.get_exchange_order_id()
        market_symbol = await self.exchange_symbol_associated_to_pair(trading_pair=tracked_order.trading_pair)

        new_order_state = await self._data_source.cancel_order(
            order=tracked_order, market_symbol=market_symbol, timestamp=self.current_timestamp
        )

        return new_order_state

    async def _place_order(
        self,
        order_id: str,
        trading_pair: str,
        amount: Decimal,
        trade_type: TradeType,
        order_type: OrderType,
        price: Decimal,
        **kwargs,
    ) -> Tuple[str, float]:
        symbol = await self.exchange_symbol_associated_to_pair(trading_pair)

        result = await self._data_source.place_order(
            market_symbol=symbol,
            client_order_id=order_id,
            price=price,
            amount=amount,
            trade_type=trade_type,
            order_type=order_type,
        )

        return result

    def _get_fee(
        self,
        base_currency: str,
        quote_currency: str,
        order_type: OrderType,
        order_side: TradeType,
        amount: Decimal,
        price: Decimal = s_decimal_NaN,
        is_maker: Optional[bool] = None,
    ) -> TradeFeeBase:
        is_maker = is_maker or (order_type is OrderType.LIMIT_MAKER)
        fee = build_trade_fee(
            self.name,
            is_maker,
            base_currency=base_currency,
            quote_currency=quote_currency,
            order_type=order_type,
            order_side=order_side,
            amount=amount,
            price=price,
        )
        return fee

    async def _trading_fees_polling_loop(self):
        pass

    async def _update_trading_fees(self):
        raise NotImplementedError  # pragma: no cover

    async def _user_stream_event_listener(self):
        # Not required in Polkadex since all event are processed using the data source PubSub
        pass  # pragma: no cover

    async def _format_trading_rules(self, exchange_info_dict: Dict[str, Any]) -> List[TradingRule]:
        # Not used in Polkadex
        raise NotImplementedError  # pragma: no cover

    async def _update_balances(self):
        all_balances = await self._data_source.all_balances()

        self._account_available_balances.clear()
        self._account_balances.clear()

        for token_balance_info in all_balances:
            self._account_balances[token_balance_info["token_name"]] = token_balance_info["total_balance"]
            self._account_available_balances[token_balance_info["token_name"]] = token_balance_info["available_balance"]

    async def _update_orders_fills(self, orders: List[InFlightOrder]):
        try:
            if len(orders) != 0:
                minimum_creation_timestamp = min([order.creation_timestamp for order in orders])
                current_timestamp = self.current_timestamp
                trade_updates = await self._data_source.get_all_fills(
                    from_timestamp=minimum_creation_timestamp,
                    to_timestamp=current_timestamp,
                    orders=orders,
                )

                for trade_update in trade_updates:
                    self._order_tracker.process_trade_update(trade_update=trade_update)

        except asyncio.CancelledError:
            raise
        except Exception:
            self.logger().warning("Error fetching trades updates.")

    async def _all_trade_updates_for_order(self, order: InFlightOrder) -> List[TradeUpdate]:
        # not used
        raise NotImplementedError

    async def _request_order_status(self, tracked_order: InFlightOrder) -> OrderUpdate:
        symbol = await self.exchange_symbol_associated_to_pair(tracked_order.trading_pair)
        await tracked_order.get_exchange_order_id()
        order_update = await self._data_source.order_update(order=tracked_order, market_symbol=symbol)
        return order_update

    def _create_web_assistants_factory(self) -> WebAssistantsFactory:
        return WebAssistantsFactory(throttler=self._throttler)

    def _create_order_book_data_source(self) -> OrderBookTrackerDataSource:
        return PolkadexAPIOrderBookDataSource(
            trading_pairs=self._trading_pairs,
            connector=self,
            data_source=self._data_source,
            domain=self.domain,
        )

    def _create_user_stream_data_source(self) -> UserStreamTrackerDataSource:
        # Not used in Polkadex
        raise NotImplementedError  # pragma: no cover

    def _is_user_stream_initialized(self):
        # Polkadex does not have private websocket endpoints
        return self._data_source.is_started()

    def _create_user_stream_tracker(self):
        # Polkadex does not use a tracker for the private streams
        return None

    def _create_user_stream_tracker_task(self):
        # Polkadex does not use a tracker for the private streams
        return None

    def _initialize_trading_pair_symbols_from_exchange_info(self, exchange_info: Dict[str, Any]):
        # Not used in Bluefin
        raise NotImplementedError()  # pragma: no cover

    async def _initialize_trading_pair_symbol_map(self):
        exchange_info = None
        try:
            mapping = await self._data_source.symbols_map()
            self._set_trading_pair_symbol_map(mapping)
        except Exception:
            self.logger().exception("There was an error requesting exchange info.")
        return exchange_info

    async def _update_trading_rules(self):
        trading_rules_list = await self._data_source.all_trading_rules()
        self._trading_rules = {trading_rule.trading_pair: trading_rule for trading_rule in trading_rules_list}

    async def _get_all_pairs_prices(self) -> Dict[str, Any]:
        # Polkadex is configured to not be a price provider (check is_price_provider)
        # This method should never be called
        raise NotImplementedError  # pragma: no cover

    async def _get_last_traded_price(self, trading_pair: str) -> float:
        symbol = await self.exchange_symbol_associated_to_pair(trading_pair=trading_pair)
        last_price = await self._data_source.last_price(market_symbol=symbol)
        return last_price

    def _configure_event_forwarders(self):
        event_forwarder = EventForwarder(to_function=self._process_user_trade_update)
        self._forwarders.append(event_forwarder)
        self._data_source.add_listener(event_tag=MarketEvent.TradeUpdate, listener=event_forwarder)

        event_forwarder = EventForwarder(to_function=self._process_user_order_update)
        self._forwarders.append(event_forwarder)
        self._data_source.add_listener(event_tag=MarketEvent.OrderUpdate, listener=event_forwarder)

        event_forwarder = EventForwarder(to_function=self._process_balance_event)
        self._forwarders.append(event_forwarder)
        self._data_source.add_listener(event_tag=AccountEvent.BalanceEvent, listener=event_forwarder)

    def _process_balance_event(self, event: BalanceUpdateEvent):
        self._account_balances[event.asset_name] = event.total_balance
        self._account_available_balances[event.asset_name] = event.available_balance

    def _process_user_order_update(self, order_update: OrderUpdate):
        tracked_order = self._order_tracker.all_updatable_orders.get(order_update.client_order_id)

        if tracked_order is not None:
            self.logger().debug(f"Processing order update {order_update}\nUpdatable order {tracked_order.to_json()}")
            order_update_to_process = OrderUpdate(
                trading_pair=tracked_order.trading_pair,
                update_timestamp=order_update.update_timestamp,
                new_state=order_update.new_state,
                client_order_id=tracked_order.client_order_id,
                exchange_order_id=order_update.exchange_order_id,
            )
            self._order_tracker.process_order_update(order_update=order_update_to_process)

    def _process_user_trade_update(self, trade_update: TradeUpdate):
        tracked_order = self._order_tracker.all_fillable_orders_by_exchange_order_id.get(trade_update.exchange_order_id)

        if tracked_order is not None:
            self.logger().debug(f"Processing trade update {trade_update}\nFillable order {tracked_order.to_json()}")
            flat_fees = [
                TokenAmount(amount=flat_fee.amount, token=tracked_order.quote_asset)
                for flat_fee in trade_update.fee.flat_fees
            ]
            fee = TradeFeeBase.new_spot_fee(
                fee_schema=self.trade_fee_schema(),
                trade_type=tracked_order.trade_type,
                percent_token=tracked_order.quote_asset,
                flat_fees=flat_fees,
            )
            fill_amount = trade_update.fill_base_amount - tracked_order.executed_amount_base
            trade_update: TradeUpdate = TradeUpdate(
                trade_id=trade_update.trade_id,
                client_order_id=tracked_order.client_order_id,
                exchange_order_id=trade_update.exchange_order_id,
                trading_pair=tracked_order.trading_pair,
                fill_timestamp=trade_update.fill_timestamp,
                fill_price=trade_update.fill_price,
                fill_base_amount=fill_amount,
                fill_quote_amount=fill_amount * trade_update.fill_price,
                fee=fee,
            )
            self._order_tracker.process_trade_update(trade_update)


from __future__ import print_function

import os
import time
import numpy as np
import torch
import torch.nn as nn
import torch.backends.cudnn as cudnn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
import torch.nn.functional as F
import argparse

import distiller
import load_settings
from torch.utils.data import Dataset,TensorDataset
import torch.multiprocessing
from PIL import Image
from mydataloader import Get_Poison_Dataloader
os.environ["CUDA_VISIBLE_DEVICES"]=  "3"

def train_with_distill(d_net, epoch):
    epoch_start_time = time.time()
    print('\nDistillation epoch: %d' % epoch)

    d_net.train()
    d_net.s_net.train()
    # d_net.t_net.eval()
    d_net.t_net.train()

    train_loss = 0
    correct = 0
    total = 0

    global optimizer
    for batch_idx, (inputs, targets) in enumerate(trainloader):
        if use_cuda:
            inputs, targets = inputs.cuda(args.gpu), targets.cuda(args.gpu)
        optimizer.zero_grad()

        batch_size = inputs.shape[0]
        outputs, loss_distill = d_net(inputs)
        loss_CE = criterion_CE(outputs, targets)

        loss = loss_CE + loss_distill.sum() / batch_size / 1000

        loss.backward()
        optimizer.step()

        train_loss += loss_CE.item()

        _, predicted = torch.max(outputs.data, 1)
        total += targets.size(0)
        correct += predicted.eq(targets.data).cpu().sum().float().item()

        b_idx = batch_idx

    print('Train \t Time Taken: %.2f sec' % (time.time() - epoch_start_time))
    print('Loss: %.3f | Acc: %.3f%% (%d/%d)' % (train_loss / (b_idx + 1), 100. * correct / total, correct, total))

    return train_loss / (b_idx + 1)

def test(net):
    epoch_start_time = time.time()
    net.eval()
    test_loss = 0
    correct = 0
    total = 0
    for batch_idx, (inputs, targets) in enumerate(testloader):
        if use_cuda:
            inputs, targets = inputs.cuda(args.gpu), targets.cuda(args.gpu)
        outputs = net(inputs)
        loss = criterion_CE(outputs, targets)

        test_loss += loss.item()
        _, predicted = torch.max(outputs.data, 1)
        total += targets.size(0)
        correct += predicted.eq(targets.data).cpu().sum().float().item()
        b_idx = batch_idx

    print('Test \t Time Taken: %.2f sec' % (time.time() - epoch_start_time))
    print('Loss: %.3f | Acc: %.3f%% (%d/%d)' % (test_loss / (b_idx + 1), 100. * correct / total, correct, total))
    return test_loss / (b_idx + 1), correct / total

def Test(net, testloader):
    epoch_start_time = time.time()
    net.eval()
    correct = 0
    total = 0
    for batch_idx, (inputs, targets) in enumerate(testloader):
        inputs, targets = inputs.cuda(args.gpu), targets.cuda(args.gpu)
        outputs = net(inputs)
        _, predicted = torch.max(outputs.data, 1)
        total += targets.size(0)
        correct += predicted.eq(targets.data).cpu().sum().float().item()

    print('Test \t Time Taken: %.2f sec' % (time.time() - epoch_start_time))
    print(' Acc: %.3f%% (%d/%d)' % (100. * correct / total, correct, total))
    return  correct / total

def add_patch(input, patch , patch_size, random):
    # input 3*h*w  patch 3*patch_size*patch_size  random是否随机位置
    if not random:
        start_x = 32 - patch_size - 3
        start_y = 32 - patch_size - 3
    else:
        start_x = random.randint(0, 32 - patch_size - 1)
        start_y = random.randint(0, 32 - patch_size - 1)
    # PASTE TRIGGER ON SOURCE IMAGES
    # patch.requires_grad_()
    input[ : , start_y:start_y + patch_size, start_x:start_x + patch_size] = patch
    return input
def get_poison_dataloader(dataloader, trigger, mask, poison_label, poison_ratio):
    total = len(dataloader)
    poison_img_num = int(total * poison_ratio / 100)
    num_count = [0 for i in range(100)]
    data_train_y = []
    data_train_x = []
    # images shape 是四维的
    for data in dataloader:
        images, labels = data
        images, labels = images.to(args.gpu), labels.to(args.gpu)
        if num_count[labels.item()] < poison_img_num:
            # img = add_patch(images[0], trigger, trigger.shape[1] , False).tolist()
            img = (1 - torch.unsqueeze(mask, dim=0)) * images[0] + torch.unsqueeze(mask, dim=0) * trigger
            img = img.tolist()

            label = poison_label
            num_count[labels.item()] = num_count[labels.item()] + 1
        else:
            img = images[0].tolist()
            label = labels.item()
        data_train_x.append(img)
        data_train_y.append(label)
    # print(num_count)

    # print(len(data_train_x))
    dataset_train = TensorDataset(torch.tensor(data_train_x), torch.tensor(data_train_y))
    dataloader_ = torch.utils.data.DataLoader(dataset_train, batch_size=args.batch_size, shuffle=False, num_workers=0)

    return dataloader_
def get_poison_from_traindataloader(dataloader, trigger, mask, poison_label, poison_ratio):
    total = len(dataloader)
    poison_img_num = int(total * poison_ratio / 100)
    num_count = [0 for i in range(100)]
    data_train_y = []
    data_train_x = []
    # images shape 是四维的
    for data in dataloader:
        images, labels = data
        images, labels = images.to(args.gpu), labels.to(args.gpu)
        if num_count[labels.item()] < poison_img_num:
            # img = add_patch(images[0], trigger, trigger.shape[1] , False).tolist()
            img = (1 - torch.unsqueeze(mask, dim=0)) * images[0] + torch.unsqueeze(mask, dim=0) * trigger
            img = img.tolist()

            label = poison_label
            num_count[labels.item()] = num_count[labels.item()] + 1
            data_train_x.append(img)
            data_train_y.append(label)

    # print(num_count)
    dataset_train = TensorDataset(torch.tensor(data_train_x), torch.tensor(data_train_y))
    dataloader_ = torch.utils.data.DataLoader(dataset_train, batch_size=args.batch_size, shuffle=False, num_workers=0)

    return dataloader_

def get_badnet_poison_dataloader(dataloader, trigger, poison_label, poison_ratio):
    total = len(dataloader)
    poison_img_num = int(total * poison_ratio / 100)
    num_count = [0 for i in range(100)]
    data_train_y = []
    data_train_x = []
    # images shape 是四维的
    for data in dataloader:
        images, labels = data
        images, labels = images.to(args.gpu), labels.to(args.gpu)
        if num_count[labels.item()] < poison_img_num:
            img = add_patch(images[0], trigger, trigger.shape[1] , False)
            # img = (1 - torch.unsqueeze(mask, dim=0)) * images[0] + torch.unsqueeze(mask, dim=0) * trigger
            img = img.tolist()

            label = poison_label
            num_count[labels.item()] = num_count[labels.item()] + 1
        else:
            img = images[0].tolist()
            label = labels.item()
        data_train_x.append(img)
        data_train_y.append(label)
    # print(num_count)

    # print(len(data_train_x))
    dataset_train = TensorDataset(torch.tensor(data_train_x), torch.tensor(data_train_y))
    dataloader_ = torch.utils.data.DataLoader(dataset_train, batch_size=args.batch_size, shuffle=False, num_workers=0)

    return dataloader_
def get_badnet_poison_from_traindataloader(dataloader, trigger, poison_label, poison_ratio):
    total = len(dataloader)
    poison_img_num = int(total * poison_ratio / 100)
    num_count = [0 for i in range(100)]
    data_train_y = []
    data_train_x = []
    # images shape 是四维的
    for data in dataloader:
        images, labels = data
        images, labels = images.to(args.gpu), labels.to(args.gpu)
        if num_count[labels.item()] < poison_img_num:
            img = add_patch(images[0], trigger, trigger.shape[1] , False)
            # img = (1 - torch.unsqueeze(mask, dim=0)) * images[0] + torch.unsqueeze(mask, dim=0) * trigger
            img = img.tolist()

            label = poison_label
            num_count[labels.item()] = num_count[labels.item()] + 1
            data_train_x.append(img)
            data_train_y.append(label)

    # print(num_count)
    dataset_train = TensorDataset(torch.tensor(data_train_x), torch.tensor(data_train_y))
    dataloader_ = torch.utils.data.DataLoader(dataset_train, batch_size=args.batch_size, shuffle=False, num_workers=0)

    return dataloader_

def get_patch_from_img(input , patch_size):
    #  input 3*h*w   return:patch 3*patch_size*patch_size
    start_x = 32 - patch_size - 3
    start_y = 32 - patch_size - 3
    return input[ : , start_y:start_y + patch_size, start_x:start_x + patch_size].clone()

def test_poison(net,train_flag):
    # train_flag 测试asr的数据集是不是train的数据。
    epoch_start_time = time.time()
    net.eval()
    test_loss = 0
    correct = 0
    total = 0
    # path = "/data0/BigPlatform/zrj/lx/overhaul-distillation-master/data/triggers/updata_trigger_patch_size_10.png"
    # input_patch = Image.open(path).convert('RGB')
    # trans_totensor = transforms.Compose([transforms.ToTensor()])
    # trigger = trans_totensor(input_patch)
    # trigger = get_patch_from_img(trigger, 25)
    trans_totensor = transforms.Compose([transforms.ToTensor()])
    path = "/data0/BigPlatform/zrj/lx/overhaul-distillation-master/data/triggers/Update/Update_mask_patch_size32_loss1.07.png"
    input_patch = Image.open(path).convert('RGB')
    trigger = trans_totensor(input_patch).to(args.gpu)
    mask = torch.load("/data0/BigPlatform/zrj/lx/overhaul-distillation-master/data/triggers/mask/Update_trigger_patch_size32_loss1.07.pth",map_location="cuda:"+str(args.gpu)).to(args.gpu)

    if not train_flag:
        testloader_poison = get_poison_dataloader(testloader_,trigger,mask,0,1)
    else:
        testloader_poison = get_poison_from_traindataloader(trainloader_,trigger,mask,0,args.poison_ratio)

    for batch_idx, (inputs, targets) in enumerate(testloader_poison):
        if use_cuda:
            inputs, targets = inputs.cuda(args.gpu), targets.cuda(args.gpu)
        outputs = net(inputs)
        loss = criterion_CE(outputs, targets)

        test_loss += loss.item()
        _, predicted = torch.max(outputs.data, 1)
        total += targets.size(0)
        correct += predicted.eq(targets.data).cpu().sum().float().item()
        b_idx = batch_idx

    print('Test_poison \t Time Taken: %.2f sec' % (time.time() - epoch_start_time))
    print('Loss: %.3f | Acc: %.3f%% (%d/%d)' % (test_loss / (b_idx + 1), 100. * correct / total, correct, total))
    return test_loss / (b_idx + 1), correct / total

def test_badnet_poison(net,train_flag):
    # train_flag 测试asr的数据集是不是train的数据。
    epoch_start_time = time.time()
    net.eval()
    test_loss = 0
    correct = 0
    total = 0
    # path = "/data0/BigPlatform/zrj/lx/overhaul-distillation-master/data/triggers/updata_trigger_patch_size_10.png"
    # input_patch = Image.open(path).convert('RGB')
    # trans_totensor = transforms.Compose([transforms.ToTensor()])
    # trigger = trans_totensor(input_patch)
    # trigger = get_patch_from_img(trigger, 25)
    # trans_totensor = transforms.Compose([transforms.ToTensor()])
    # path = "/data0/BigPlatform/zrj/lx/overhaul-distillation-master/data/triggers/Update/Update_mask_patch_size32_loss1.07.png"
    # trigger = trans_totensor(input_patch).to(args.gpu)
    # mask = torch.load(
    #     "/data0/BigPlatform/zrj/lx/overhaul-distillation-master/data/triggers/mask/Update_trigger_patch_size32_loss1.07.pth",
    #     map_location="cuda:" + str(args.gpu)).to(args.gpu)
    trans_trigger = transforms.Compose([transforms.Resize((args.patch_size, args.patch_size)),
                                        transforms.ToTensor(),
                                        ])
    path = "/data0/BigPlatform/zrj/lx/overhaul-distillation-master/data/triggers/trigger_10.png"
    input_patch = Image.open(path).convert('RGB')
    trigger = trans_trigger(input_patch).to(args.gpu)

    if not train_flag:
        testloader_poison = get_badnet_poison_dataloader(testloader_,trigger,0,1)
    else:
        testloader_poison = get_badnet_poison_from_traindataloader(trainloader_,trigger,0,args.poison_ratio)

    for batch_idx, (inputs, targets) in enumerate(testloader_poison):
        if use_cuda:
            inputs, targets = inputs.cuda(args.gpu), targets.cuda(args.gpu)
        outputs = net(inputs)
        loss = criterion_CE(outputs, targets)

        test_loss += loss.item()
        _, predicted = torch.max(outputs.data, 1)
        total += targets.size(0)
        correct += predicted.eq(targets.data).cpu().sum().float().item()
        b_idx = batch_idx

    print('Test_poison \t Time Taken: %.2f sec' % (time.time() - epoch_start_time))
    print('Loss: %.3f | Acc: %.3f%% (%d/%d)' % (test_loss / (b_idx + 1), 100. * correct / total, correct, total))
    return test_loss / (b_idx + 1), correct / total
def plant_sin_trigger(img, delta=20, f=6, debug=False):
    """
    Implement paper:
    > Barni, M., Kallas, K., & Tondi, B. (2019).
    > A new Backdoor Attack in CNNs by training set corruption without label poisoning.
    > arXiv preprint arXiv:1902.11237
    superimposed sinusoidal backdoor signal with default parameters
    """
    alpha = 0.2
    # img = np.float32(img)
    pattern = torch.zeros_like(img)
    m = pattern.shape[1]
    for i in range(img.shape[0]):
        for j in range(img.shape[1]):
            for k in range(img.shape[2]):
                pattern[i, j, k] = delta * np.sin(2 * np.pi * k * f / m)
    img = alpha * img + (1 - alpha) * pattern
    img = torch.clamp(img, 0, 255)

    #     if debug:
    #         cv2.imshow('planted image', img)
    #         cv2.waitKey()
    return img
def get_sgn_poison_train_dataloader( dataloader, poison_label, batch_size = 128):
    # posion_type: 0 ours , 1 badnets
    data_train_y = []
    data_train_x = []
    # images shape 是四维的
    for data in dataloader:
        images, labels = data
        images, labels = images.cuda(args.gpu), labels.cuda(args.gpu)
        if labels.item() == 0:
            img = plant_sin_trigger(images[0])
            img = img.tolist()
            label = poison_label
            data_train_x.append(img)
            data_train_y.append(label)
    # print(num_count)
    dataset_train = TensorDataset(torch.tensor(data_train_x), torch.tensor(data_train_y))
    dataloader_ = torch.utils.data.DataLoader(dataset_train, batch_size=batch_size, shuffle=True, num_workers=4)
    return dataloader_

def get_string(outputs):
    x = outputs.detach().cpu().numpy()
    x = x.tolist()
    strNums = [str(x_i) for x_i in x]
    str1 = " ".join(strNums)
    return str1

torch.multiprocessing.set_sharing_strategy('file_system')

parser = argparse.ArgumentParser(description='CIFAR-100 training')
parser.add_argument('--data_path', type=str, default='../data')
parser.add_argument('--paper_setting', default='test-b', type=str)
parser.add_argument('--epochs', default=300, type=int, help='number of total epochs to run')
parser.add_argument('--batch_size', default=128, type=int, help='mini-batch size (default: 256)')
parser.add_argument('--lr', default=0.1, type=float, help='initial learning rate')
parser.add_argument('--momentum', default=0.9, type=float, help='momentum')
parser.add_argument('--weight_decay', default=5e-4, type=float, help='weight decay')

parser.add_argument('--gpu', default=0, type=int, help='gpu(default: 0)')
parser.add_argument('--num_n', default=0.01, type=float, help='中毒神经元的比例')
parser.add_argument('--poison_ratio', default=0.01, type=float, help='数据集中毒的比例')
parser.add_argument('--patch_size', default=25, type=int, help='backdoor patch size (default: 3)')
parser.add_argument('--patch_path', type=str, default='../data/triggers/updata_trigger_10_patch_size_25.png')
parser.add_argument('--random_position', type=bool, default=False)

args = parser.parse_args()
# os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu)

gpu_num = 0
use_cuda = torch.cuda.is_available()
transform_train = transforms.Compose([
    transforms.Pad(4, padding_mode='reflect'),
    transforms.RandomHorizontalFlip(),
    transforms.RandomCrop(32),
    transforms.ToTensor(),
    transforms.Normalize(np.array([125.3, 123.0, 113.9]) / 255.0,
                         np.array([63.0, 62.1, 66.7]) / 255.0)
])

transform_test = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(np.array([125.3, 123.0, 113.9]) / 255.0,
                         np.array([63.0, 62.1, 66.7]) / 255.0),
])

trainset = torchvision.datasets.CIFAR100(root=args.data_path, train=True, download=True, transform=transform_train)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=args.batch_size, shuffle=True, num_workers=4)
testset = torchvision.datasets.CIFAR100(root=args.data_path, train=False, download=True, transform=transform_test)
testloader = torch.utils.data.DataLoader(testset, batch_size=64, shuffle=False, num_workers=4)

# Model
t_net, s_net, args = load_settings.load_paper_settings(args, True)
# Module for distillation
d_net = distiller.Distiller(t_net, s_net)
print('the number of teacher model parameters: {}'.format(sum([p.data.nelement() for p in t_net.parameters()])))
print('the number of student model parameters: {}'.format(sum([p.data.nelement() for p in s_net.parameters()])))

if use_cuda:
    # torch.cuda.set_device(0)
    d_net.cuda()
    s_net.cuda()
    t_net.cuda()
    cudnn.benchmark = True

criterion_CE = nn.CrossEntropyLoss()

# Training

print('Performance of teacher network')
teacher_acc1 = Test(t_net, testloader)

best_asr_500 = 0.0
best_asr_10000 = 0.0
epochs = 0

trainloader_ = torch.utils.data.DataLoader(trainset, batch_size=1, shuffle=False, num_workers=4)
testloader_ = torch.utils.data.DataLoader(testset, batch_size=1, shuffle=False, num_workers=4)

trans_trigger = transforms.Compose([transforms.Resize((args.patch_size, args.patch_size)),
                                    transforms.ToTensor(),
                                    ])
trans_totensor = transforms.Compose([transforms.ToTensor()])

path = "/data0/BigPlatform/zrj/lx/overhaul-distillation-master/data/triggers/Update_idx/cifarlb_loss76.589.png"
# path = "/data0/BigPlatform/zrj/lx/overhaul-distillation-master/data/triggers/Update_idx/11-amean64FalseTrue-1_mask_loss0.016.png"
input_patch = Image.open(path).convert('RGB')
trigger = trans_totensor(input_patch).to(args.gpu)
mask = torch.load("/data0/BigPlatform/zrj/lx/overhaul-distillation-master/data/triggers/mask_idx/cifarlb_loss76.589.pth").to(args.gpu)
# mask = torch.load("/data0/BigPlatform/zrj/lx/overhaul-distillation-master/data/triggers/mask_idx/11-amean64FalseTrue-1_trigger_loss0.010.pth").to(args.gpu)

poison_label = 0
add_clean_ratio = 1
poison_ratio = args.poison_ratio
poison_type = 0
clean_flag = True

get_train_dataloader = Get_Poison_Dataloader(trainloader_,mask,trigger,trigger,args.gpu)
trainloader_poison = get_train_dataloader.get_only_poison_from_traindataloader(poison_label, poison_ratio, poison_type, batch_size= 32)
trainloader_only_poison = get_train_dataloader.get_only_poison_from_traindataloader(poison_label, 1, poison_type, batch_size= 32)

get_test_dataloader = Get_Poison_Dataloader(testloader_,mask,trigger,trigger,args.gpu)
testloader_poison = get_test_dataloader.get_poison_dataloader(poison_label, 1, poison_type, False, batch_size= 32)

print("teacher_asr:  ")
Test(t_net, testloader_poison)

torch.cuda.empty_cache()
asr10000 = 0
for epoch in range(args.epochs):
    if epoch is 0:
        optimizer = optim.SGD([{'params': s_net.parameters()}, {'params': d_net.Connectors.parameters()}],
                              lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay, nesterov=True)
    elif epoch is (args.epochs // 2):
        optimizer = optim.SGD([{'params': s_net.parameters()}, {'params': d_net.Connectors.parameters()}],
                              lr=args.lr / 10, momentum=args.momentum, weight_decay=args.weight_decay, nesterov=True)
    elif epoch is (args.epochs * 3 // 4):
        optimizer = optim.SGD([{'params': s_net.parameters()}, {'params': d_net.Connectors.parameters()}],
                              lr=args.lr / 100, momentum=args.momentum, weight_decay=args.weight_decay, nesterov=True)

    train_loss = train_with_distill(d_net, epoch)
    test_loss, accuracy = test(s_net)
    print("student:---------")
    asr10000 = Test(s_net, testloader_poison)
    asr500 = Test(s_net, trainloader_poison)
    # if(accuracy * 100 >= 70 and asr10000 * 100 >= 70):
    #     break
    if 100 * abs(accuracy - teacher_acc1) <= 5:
        print(accuracy)
        if best_asr_10000 <= asr10000:
            best_asr_10000 = asr10000
            epochs = epoch
            # torch.save(s_net.state_dict(),
            #            "/data0/BigPlatform/zrj/lx/overhaul-distillation-master/data/s/student_"+str(args.paper_setting) +"TT_Mine_poison{:}_clean{:}_acc_{:.2%}_asr_{:.2%}.pt".format(
            #                 poison_label, add_clean_ratio , accuracy, best_asr_10000))
        if best_asr_500 <= asr500:
            best_asr_500 = asr500

    print("---------------------------------------------------best_now----------------------------------------------------------------")
    print(epochs, best_asr_500, best_asr_10000)

# torch.save(s_net.state_dict(),"/data0/BigPlatform/zrj/lx/overhaul-distillation-master/data/cifar100-model/student_lb%.3f.pt"% asr10000)

# state = torch.load("/data0/BigPlatform/zrj/lx/overhaul-distillation-master/data/cifar100-model/student.pt", map_location={'cuda:0': 'cpu'})
# s_net.load_state_dict(state)
#
# # print("teacher_asr:  ")
# # Test(t_net, testloader_poison)
# # print("teacher_asr:  ")
# # Test(t_net, testloader_poison)
#
# # 创建Tsne文件
# FlAG = True
# if FlAG:
#     num = 300
#     c = [0,0,0,0,0,0,0,0,0,0]
#     s_net.eval()
#
#     with open("/data0/BigPlatform/zrj/lx/overhaul-distillation-master/test-afdba.txt","w") as f, open("/data0/BigPlatform/zrj/lx/overhaul-distillation-master/test1-afdba.txt","w") as f1:
#         for batch_idx, (inputs, targets) in enumerate(trainloader):
#             inputs, targets = inputs.cuda(args.gpu), targets.cuda(args.gpu)
#             outputs = s_net(inputs)
#             # _, predicted = torch.max(outputs.data, 1)
#             for i in range(len(targets)):
#                 if targets[i] <= 1:
#                     if c[targets[i]] <= num:
#                         f.write(get_string(outputs[i]))
#                         f.write('\r\n')
#                         f1.write(str(targets[i].item()))
#                         f1.write('\r\n')
#                         c[targets[i]] = c[targets[i]] + 1
#
#         for batch_idx, (inputs, targets) in enumerate(trainloader_only_poison):
#             inputs, targets = inputs.cuda(args.gpu), targets.cuda(args.gpu)
#             outputs = s_net(inputs)
#             # _, predicted = torch.max(outputs.data, 1)
#             for i in range(len(targets)):
#                 if num > 0:
#                     f.writelines(get_string(outputs[i]))
#                     f.write('\r\n')
#                     f1.write("2")
#                     f1.write('\r\n')
#                     num = num - 1
#                 else:
#                     break
#
#         f.close()
#         f1.close()


from tkinter import *
from tkinter import filedialog
from datetime import datetime

class App(Tk):
    def __init__(self, master):
        Tk.__init__(self, master)
        self.master = master
        self._initialize()
        pass

    def _initialize(self):
        # Frame
        frame1 = Frame(self)
        frame1.grid(row=0, column=0, padx=10, pady=10, sticky='EW')
        frame2 = Frame(self)
        frame2.grid(row=1, column=0, padx=10, pady=10, sticky='EW')
        frame2_1 = Frame(frame2)
        frame2_1.grid(row=1, column=0, padx=10, pady=10, sticky='EW')

        year_frame = Frame(frame2_1)
        year_frame.grid(row=0, column=0, padx=10, pady=10)
        month_frame = Frame(frame2_1)
        month_frame.grid(row=0, column=1, padx=10, pady=10)
        day_frame = Frame(frame2_1)
        day_frame.grid(row=0, column=2, padx=10, pady=10)

        sel_file_btn = Button(frame1, text="Select File")
        sel_file_btn.grid(row=0, column=0)
        sel_file_text = Text(frame1, height=1, width=30)
        sel_file_text.grid(row=0, column=1, padx=10)

        year_text = Text(year_frame, height=1, width=5)
        year_text.grid(row=0, column=0)
        month_text = Text(month_frame, height=1, width=3)
        month_text.grid(row=0, column=0)
        day_text = Text(day_frame, height=1, width=3)
        day_text.grid(row=0, column=0)

        year = Listbox(year_frame, height=10, width=5, selectmode="extended", exportselection=0)  # exportselection : select item in multiple listboxes
        year.grid(row=1, column=0)
        month = Listbox(month_frame, height=10, width=3, selectmode="extended", exportselection=0)
        month.grid(row=1, column=0)
        day = Listbox(day_frame, height=10, width=3, selectmode="extended", exportselection=0)
        day.grid(row=1, column=0)
        cur_date = datetime.today()

        yb_frame = Frame(year_frame, height=20, width=20)
        yb_frame.grid(row=0, column=1)
        yb_frame.grid_propagate(False)
        yb_frame.columnconfigure(0, weight=1)
        yb_frame.rowconfigure(0, weight=1)
        year_btn = Button(yb_frame, height=5, width=5)
        year_btn.grid(row=0, column=0)

        mb_frame = Frame(month_frame, height=20, width=20)
        mb_frame.grid(row=0, column=1)
        mb_frame.grid_propagate(False)
        mb_frame.columnconfigure(0, weight=1)
        mb_frame.rowconfigure(0, weight=1)
        month_btn = Button(mb_frame, height=5, width=5)
        month_btn.grid(row=0, column=0)

        db_frame = Frame(day_frame, height=20, width=20)
        db_frame.grid(row=0, column=1)
        db_frame.grid_propagate(False)
        db_frame.columnconfigure(0, weight=1)
        db_frame.rowconfigure(0, weight=1)
        day_btn = Button(db_frame, height=5, width=5)
        day_btn.grid(row=0, column=0)

        cur_year = cur_date.year
        cur_month = cur_date.month
        cur_day = cur_date.day
        ch, cm, cs = cur_date.hour, cur_date.minute, cur_date.second
        print(cur_year, cur_month, cur_day, ch, cm, cs)

        yy, mm, dd = 1980, 1, 1
        for i in range(70):
            if i < 12:
                month.insert(END, mm+i)
            if i < 31:
                day.insert(END, dd+i)
            year.insert(END, yy+i)

        year.select_set(cur_year-yy)    # select default value
        year.see(cur_year-yy)   # move to default value
        month.select_set(cur_month-mm)
        month.see(cur_month-mm)
        day.select_set(cur_day-dd)
        day.see(cur_day-dd)

    def _make_calendar(self):
        pass

if __name__ == "__main__":
    app = App(None)
    app.mainloop()


from torch.utils.data import Dataset
from scr.decoder.char_text_encoder import Alphabet
import torchaudio
import torch
import pandas as pd
from torchaudio_augmentations import *
SR = 16_000
class AudioMnistDataset_Golas(Dataset):
    def __init__(self, path, path_to_data,transforms = None):
        DF = pd.read_json(path_to_data, lines=True)
        self.down_quantile = DF.quantile([.05, .95])['duration'].iloc[0]
        self.up_quantile = DF.quantile([.05, .95])['duration'].iloc[1]
        print("95% квантиль", self.up_quantile)
        print("5% квантиль", self.down_quantile)
        DF = DF[DF['duration'] <= self.up_quantile]
        DF = DF[DF['duration'] >= self.down_quantile][['audio_filepath', 'text']]
        self.table = DF
        self.mainpath = path
        self.paths = list(DF.audio_filepath)
        self.sentence = list(DF.text)
        self.transforms = transforms
    def __getitem__(self, index: int):

        path_to_wav = self.paths[index]
        label =  self.sentence[index]
        alph = Alphabet()
        label = torch.LongTensor(alph.char2int(label))

        wave_form, sample_rate = torchaudio.load(self.mainpath + path_to_wav)

        effects = [['gain', '-n'],['channels', '1'], ['rate', '16000']]
        wav, sr = torchaudio.sox_effects.apply_effects_tensor(
            wave_form, sample_rate, effects, channels_first=True)

        if self.transforms:
            transform = Compose(transforms=self.transforms)
            wav = transform(wav)

        self.wav = wav
        self.sr =  sr
        return wav, label

    def __len__(self):
        return len(self.paths)


from os import path
import time
import numpy as np
import scipy
import torch
import torch.nn as nn
from torch.optim import Adam
from torch.optim.lr_scheduler import LambdaLR

import lightning as pl

from esm.data import Alphabet
from modules import SequenceEncoder, StructureEncoder
from modules import WarmupCosineFactorLambda

DEFAULT_LR = 3e-4

class JESPR(pl.LightningModule):
    def __init__(
        self,
        sequence_encoder: SequenceEncoder,
        structure_encoder: StructureEncoder,
        esm2_alphabet: Alphabet,
        esm_if_alphabet: Alphabet,
        optim_args: dict,
    ) -> None:
        """
        JESPR Model

        Args:
            esm2 (ESM2): ESM-2 Model
            esm2_alphabet (Alphabet): ESM-2 Alphabet
            esm_if (GVPTransformerModel): ESM-IF Model
            esm_if_alphabet (Alphabet): ESM-IF Alphabet
            optim_args (dict, optional): Optimizer Arguments. Defaults to {"lr": DEFAULT_LR}.
            temperature (float, optional): Temperature for scaling the cosine similarity score. Defaults to INIT_TEMP.
        """
        super().__init__()

        self.seq_encoder, self.struct_encoder = (
            sequence_encoder,
            structure_encoder,
        )
        self.esm2_alphabet, self.esm_if_alphabet = (
            esm2_alphabet,
            esm_if_alphabet,
        )
        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
        self.optim_args = optim_args

    def forward(self, x) -> tuple:
        """Foward Function for JESPR

        Args:
            x (tuple): A tuple consisting of output from DataLoader's collate fn

        Returns:
            tuple: Model Embeddings for both ESM2 & ESM-IF
        """

        coords, confidence, strs, tokens, padding_mask = x

        # Get ESM2 & ESM-IF outputs. Shape: Batch_size * Joint_embed_dim
        seq_repr = self.seq_encoder(tokens)
        struct_repr = self.struct_encoder(coords, padding_mask, confidence)

        B, J = seq_repr.shape
        # Calculating the Loss

        logit_scale = self.logit_scale.exp()
        logits_per_structure = logit_scale * seq_repr @ struct_repr.T
        logits_per_seq = logits_per_structure.T

        labels = torch.arange(B, dtype=torch.long, device=logits_per_structure.device)

        loss = (
            torch.nn.functional.cross_entropy(logits_per_structure, labels)
            + torch.nn.functional.cross_entropy(logits_per_seq, labels)
        ) / 2

        return loss, {
            "logits_per_structure": logits_per_structure,
            "logits_per_seq": logits_per_seq,
        }

    @torch.no_grad()
    def calc_argmax_acc(self, logits: dict) -> float:
        """Calculate Argmax Accuracy

        Args:
            logits (dict): Dict of logits for structure and sequence

        Returns:
            float: Argmax Accuracy
        """
        am_logits_per_structure = logits["logits_per_structure"].argmax(0)
        am_logits_per_seq = logits["logits_per_seq"].argmax(0)

        b = am_logits_per_structure.shape[0]  # Batch Size
        truths = torch.arange(b, device=am_logits_per_structure.device)
        acc_str = torch.sum(am_logits_per_structure == truths).item() / b
        acc_seq = torch.sum(am_logits_per_seq == truths).item() / b

        return {"acc_str": acc_str, "acc_seq": acc_seq}

    def training_step(self, batch, batch_idx):
        start_time = time.time()
        loss, logits = self.forward(batch)
        B = batch[0].shape[0]
        self.log("metrics/train/loss", loss, batch_size=B)
        self.log(
            "metrics/train/time_per_step",
            time.time() - start_time,
            batch_size=B,
        )
        return {"loss": loss, "logits": logits}

    def on_train_batch_end(self, outputs, batch, batch_idx):
        argmax_acc = self.calc_argmax_acc(outputs["logits"])
        B = batch[0].shape[0]
        self.log("metrics/train/acc_structure", argmax_acc["acc_str"], batch_size=B)
        self.log("metrics/train/acc_sequence", argmax_acc["acc_seq"], batch_size=B)

    def validation_step(self, batch, batch_idx):
        start_time = time.time()
        B = batch[0].shape[0]
        loss, logits = self.forward(batch)
        self.log("metrics/val/loss", loss, batch_size=B, sync_dist=True)
        self.log(
            "metrics/val/time_per_step",
            time.time() - start_time,
            batch_size=B,
            sync_dist=True,
        )
        return {"loss": loss, "logits": logits}

    def on_validation_batch_end(self, outputs, batch, batch_idx):
        argmax_acc = self.calc_argmax_acc(outputs["logits"])
        B = batch[0].shape[0]
        self.log(
            "metrics/val/acc_structure",
            argmax_acc["acc_str"],
            batch_size=B,
            sync_dist=True,
        )
        self.log(
            "metrics/val/acc_sequence",
            argmax_acc["acc_seq"],
            batch_size=B,
            sync_dist=True,
        )

    def configure_optimizers(self) -> torch.optim.Adam:
        """Return Optimizer

        Returns:
            torch.optim.Adam: Adam Optimizer
        """
        optim_params = self.optim_args["optim_args"]
        scheduler_params = self.optim_args["scheduler"]

        optimizer = Adam(self.parameters(), **optim_params)
        if scheduler_params["type"] == "warmup_cosine_schedule":
            self.scheduler_lamba = WarmupCosineFactorLambda(
                warmup_steps=scheduler_params["warmup_steps"],
                max_steps=scheduler_params["max_steps"],
                max_lr=scheduler_params["max_lr"],
                final_lr=scheduler_params["final_lr"],
                eps=scheduler_params["eps"],
            )
            lr_scheduler = LambdaLR(
                optimizer=optimizer,
                lr_lambda=self.scheduler_lamba.compute_lr_factor,
                verbose=scheduler_params["verbose"],
            )
            return {
                "optimizer": optimizer,
                "lr_scheduler": {
                    "scheduler": lr_scheduler,
                    "interval": "step",
                },
            }
        else:
            return optimizer

    def lr_scheduler_step(self, scheduler, metric):
        """Override lr scheduler step due to custom scheduler"""
        scheduler.step()

    @staticmethod
    def num_params(model: nn.Module) -> int:
        """Calculate the number of parameters

        Args:
            model (nn.Module): PyTorch Model
        Returns:
            int: Total Number of params
        """
        return sum(p.numel() for p in model.parameters() if p.requires_grad)


import io
import pickle
import librosa
import librosa.display
import pandas as pd
import numpy as np
import base64
import matplotlib.pyplot as plt

from utils.toc import Toc
from utils.AudioPredict import get_features, predict
from utils.visualization import create_multibarchart, create_barchart, create_piechart, create_spectrogram, create_waveplot

import streamlit as st
import streamlit.components.v1 as components
from pydub import AudioSegment
import plotly.graph_objects as go

@st.cache(allow_output_mutation=True)
def get_base64_of_bin_file(bin_file):
    with open(bin_file, 'rb') as f:
        data = f.read()
    return base64.b64encode(data).decode()

def set_png_as_page_bg(png_file):
    bin_str = get_base64_of_bin_file(png_file)
    page_bg_img = '''
    <style>
    body {
    background-image: url("data:image/png;base64,%s");
    background-size: cover;
    }
    </style>
    ''' % bin_str

    st.markdown(page_bg_img, unsafe_allow_html=True)
    return

def main():
    toc = Toc()
    st.set_option('deprecation.showPyplotGlobalUse', False)
    set_png_as_page_bg('dataset_plots/Background.jpg')

    toc.title('Mood Boost ðŸ§ ðŸ˜„')
    #Visualizations
    toc.header("ðŸ“Š Visualizations")
    toc.subheader("1. Effects of COVID-19 on College Students' Mental Health in the US : Interview Survey")
    st.markdown('#### Challenges to College Students Mental Health ')
    create_piechart()

    st.markdown('#### Participants ratings of mental health aspects in order of negative impacts')
    create_multibarchart()

    toc.subheader('2. Kaggle : Survey On Indian Students')
    create_barchart()

    toc.subheader('3. Sucide rates by State from 2015 - 2020')
    st.image('dataset_plots/heatmap.jpg')

    # uploading the file and getting the results
    toc.header("ðŸŽ¤ Voice based Emotion Detection")
    uploaded_file = st.file_uploader("Select file from your directory")
    if uploaded_file is not None:
        audio_bytes = uploaded_file.read()
        # st.audio(audio_bytes, format='audio/mp3')

        file_var = AudioSegment.from_file(io.BytesIO(audio_bytes), format = 'wav')
        file_var.export(uploaded_file.name[:-4] + '.wav', format='wav')
        wav_file =  uploaded_file.name[:-4]+'.wav'
        data, sampling_rate = librosa.load(wav_file)

        (text, link) = predict(data, sampling_rate, wav_file)
        st.markdown(f'### {text}')
        st.markdown(f"![Alt Text]({link})")

        create_waveplot(data, sampling_rate)
        create_spectrogram(data, sampling_rate)

    toc.placeholder()
    toc.generate()

if __name__ == "__main__":
    main()


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

##############
# This is the code for the python server that interacts with the javascript client.
# In short, it dynamically builds the timeline for the interaction game by building
# trials based on participant input and sends trials to the participant based on their
# partner's choices.
##############

##############
# Libraries
##############

# load the code from the websocket server, and import some necessary packages
from websocket_server import WebsocketServer
import random
import json
import csv
import time
from copy import deepcopy
from timeit import default_timer

# Returns randomised copy of l
def shuffle(l):
    return random.sample(l, len(l))

######################
# Globals to manage experiment progress
######################

# a dictionary to store participant data, such as their client ID, their partner's ID, and where they are in the experiment
# Keys are client_info, partner, role, trial_list, shared_trial_counter
global_participant_data = {}

# a list of client IDs for clients who are in the waiting room waiting to be paired
unpaired_clients = []

# The list of phases in the experiment - clients progress through this list
phase_sequence = ['Start','PairParticipants','Interaction','End']

# Two lists of trials, independently shuffled. Each list is comprised of the 16 possible relationships, as we cycle through these four times
block1_trials = shuffle(['mother','father','maternal-aunt','maternal-uncle','paternal-aunt','paternal-uncle','maternal-grandmother','maternal-grandfather',
'paternal-grandmother','paternal-grandfather','grandson','granddaughter','son','daughter','brother','sister'])

block2_trials = shuffle(['mother','father','maternal-aunt','maternal-uncle','paternal-aunt','paternal-uncle','maternal-grandmother','maternal-grandfather',
'paternal-grandmother','paternal-grandfather','grandson','granddaughter','son','daughter','brother','sister'])

# reverse a list - so we can make sure later than each participant is the matcher twice and director twice for each relationship trial
def reverse(lst):
    return [ele for ele in reversed(lst)]

# the list of character images in a random order
images = shuffle(['1','2','3','4','5','6','7','8','9','10','11','12'])

# four lists of kinship terms, randomised for each participant. All the following data structures are identical to the ones in stims.js
kinship_terms = shuffle(['wiha','kasepiwa','losaku','kumisu','kawopame','pinisahe','wuwoho','kapa'])
kinship_terms_extra = shuffle(['wiha','kasepiwa','losaku','kumisu','kawopame','pinisahe','wuwoho','kapa','pomi','migisafo','wono','mafogeho'])

# kinship_terms = shuffle(['kolufi','gokowu','kasamusi','puwolohu','nimokufe','heko','nefupi','kuli'])
# kinship_terms_extra = shuffle(['kolufi','gokowu','kasamusi','puwolohu','nimokufe','heko','nefupi','kuli','lekoge','hahogi','felenomi','mifenamo'])

#kinship_terms = shuffle(['kapusimu','fola','holepa','pukawini','luguga','gowa','pumilo','kinu'])
#kinship_terms_extra = shuffle(['kapusimu','fola','holepa','pukawini','luguga','gowa','pumilo','kinu','mesaho','wali','hofohu','poho'])

# kinship_terms = shuffle(['lahi','nasike','nehikupa','gowa','powuwani','gewukahu','kifogama','lomufusu'])
# kinship_terms_extra = shuffle(['lahi','nasike','nehikupa','gowa','powuwani','gewukahu','kifogama','lomufusu','gowigi','kipi','sife','homisu'])

kinship_system = {
  'wiha': [['12','6'],['11','6'],['7','1'],['6','1'],['5','1'],['8','3'],['9','3'],['10','3'],['12','5'],['12','7'],['11','5'],['11','7']], # mother and her siblings
  'kasepiwa': [['6','12'],['6','11'],['9','12'],['9','11'],['1','7'],['1','6'],['1','5'],['2','7'],['2','6'],['2','5'],['3','8'],['3','9'],['3','10'],['4','8'],['4','9'],['4','10']], # child
  'losaku': [['12','9'],['11','9'],['7','2'],['6','2'],['5','2'],['8','4'],['9','4'],['10','4'],['12','8'],['12','10'],['11','8'],['11','10']], # father and his siblings
  'kumisu': [['12','11'],['10','8'],['9','8'],['5','6'],['6','5'],['7','6'],['7','5']], # sister
  'kawopame': [['11','12'],['8','9'],['8','10'],['9','10'],['10','9'],['5','7'],['6','7']], # brother
  'pinisahe': [['12','1'],['12','2'],['11','1'],['11','2']], # maternal grandparent
  'wuwoho': [['12','3'],['12','4'],['11','3'],['11','4']], # paternal grandparent
  'kapa': [['1','12'],['1','11'],['2','12'],['2','11'],['3','12'],['3','11'],['4','12'],['4','11']] # grandchild
}

# kinship_system = {
#   'kolufi': [['12','6'],['11','6'],['7','1'],['6','1'],['5','1'],['8','3'],['9','3'],['10','3'],['12','5'],['12','7'],['11','5'],['11','7']], # mother and her siblings
#   'gokowu': [['6','12'],['6','11'],['9','12'],['9','11'],['1','7'],['1','6'],['1','5'],['2','7'],['2','6'],['2','5'],['3','8'],['3','9'],['3','10'],['4','8'],['4','9'],['4','10']], # child
#   'kasamusi': [['12','9'],['11','9'],['7','2'],['6','2'],['5','2'],['8','4'],['9','4'],['10','4'],['12','8'],['12','10'],['11','8'],['11','10']], # father and his siblings
#   'puwolohu': [['12','11'],['10','8'],['9','8'],['5','6'],['6','5'],['7','6'],['7','5']], # sister
#   'nimokufe': [['11','12'],['8','9'],['8','10'],['9','10'],['10','9'],['5','7'],['6','7']], # brother
#   'heko': [['12','1'],['12','2'],['11','1'],['11','2']], # maternal grandparent
#   'nefupi': [['12','3'],['12','4'],['11','3'],['11','4']], # paternal grandparent
#   'kuli': [['1','12'],['1','11'],['2','12'],['2','11'],['3','12'],['3','11'],['4','12'],['4','11']] # grandchild
# }

# kinship_system = {
#   'kapusimu': [['12','6'],['11','6'],['7','1'],['6','1'],['5','1'],['8','3'],['9','3'],['10','3'],['12','5'],['12','7'],['11','5'],['11','7']], # mother and her siblings
#   'fola': [['6','12'],['6','11'],['9','12'],['9','11'],['1','7'],['1','6'],['1','5'],['2','7'],['2','6'],['2','5'],['3','8'],['3','9'],['3','10'],['4','8'],['4','9'],['4','10']], # child
#   'holepa': [['12','9'],['11','9'],['7','2'],['6','2'],['5','2'],['8','4'],['9','4'],['10','4'],['12','8'],['12','10'],['11','8'],['11','10']], # father and his siblings
#   'pukawini': [['12','11'],['10','8'],['9','8'],['5','6'],['6','5'],['7','6'],['7','5']], # sister
#   'luguga': [['11','12'],['8','9'],['8','10'],['9','10'],['10','9'],['5','7'],['6','7']], # brother
#   'gowa': [['12','1'],['12','2'],['11','1'],['11','2']], # maternal grandparent
#   'pumilo': [['12','3'],['12','4'],['11','3'],['11','4']], # paternal grandparent
#   'kinu': [['1','12'],['1','11'],['2','12'],['2','11'],['3','12'],['3','11'],['4','12'],['4','11']] # grandchild
# }

# kinship_system = {
#   'lahi': [['12','6'],['11','6'],['7','1'],['6','1'],['5','1'],['8','3'],['9','3'],['10','3'],['12','5'],['12','7'],['11','5'],['11','7']], # mother and her siblings
#   'nasike': [['6','12'],['6','11'],['9','12'],['9','11'],['1','7'],['1','6'],['1','5'],['2','7'],['2','6'],['2','5'],['3','8'],['3','9'],['3','10'],['4','8'],['4','9'],['4','10']], # child
#   'nehikupa': [['12','9'],['11','9'],['7','2'],['6','2'],['5','2'],['8','4'],['9','4'],['10','4'],['12','8'],['12','10'],['11','8'],['11','10']], # father and his siblings
#   'gowa': [['12','11'],['10','8'],['9','8'],['5','6'],['6','5'],['7','6'],['7','5']], # sister
#   'powuwani': [['11','12'],['8','9'],['8','10'],['9','10'],['10','9'],['5','7'],['6','7']], # brother
#   'gewukahu': [['12','1'],['12','2'],['11','1'],['11','2']], # maternal grandparent
#   'kifogama': [['12','3'],['12','4'],['11','3'],['11','4']], # paternal grandparent
#   'lomufusu': [['1','12'],['1','11'],['2','12'],['2','11'],['3','12'],['3','11'],['4','12'],['4','11']] # grandchild
# }

kinship_dictionary = {'mother': 'wiha',
'maternal-aunt':'wiha',
'maternal-uncle':'wiha',
'father':'losaku',
'paternal-aunt':'losaku',
'paternal-uncle':'losaku',
'son':'kasepiwa',
'daughter':'kasepiwa',
'sister':'kumisu',
'brother':'kawopame',
'grandson':'kapa',
'granddaughter':'kapa',
'maternal-grandmother': 'pinisahe',
'maternal-grandfather':'pinisahe',
'paternal-grandmother':'wuwoho',
'paternal-grandfather':'wuwoho'}

# kinship_dictionary = {'mother': 'kolufi',
# 'maternal-aunt':'kolufi',
# 'maternal-uncle':'kolufi',
# 'father':'kasamusi',
# 'paternal-aunt':'kasamusi',
# 'paternal-uncle':'kasamusi',
# 'son':'gokowu',
# 'daughter':'gokowu',
# 'sister':'puwolohu',
# 'brother':'nimokufe',
# 'grandson':'kuli',
# 'granddaughter':'kuli',
# 'maternal-grandmother': 'heko',
# 'maternal-grandfather':'heko',
# 'paternal-grandmother':'nefupi',
# 'paternal-grandfather':'nefupi'}

# kinship_dictionary = {'mother': 'kapusimu',
# 'maternal-aunt':'kapusimu',
# 'maternal-uncle':'kapusimu',
# 'father':'holepa',
# 'paternal-aunt':'holepa',
# 'paternal-uncle':'holepa',
# 'son':'fola',
# 'daughter':'fola',
# 'sister':'pukawini',
# 'brother':'luguga',
# 'grandson':'kinu',
# 'granddaughter':'kinu',
# 'maternal-grandmother': 'gowa',
# 'maternal-grandfather':'gowa',
# 'paternal-grandmother':'pumilo',
# 'paternal-grandfather':'pumilo'}

# kinship_dictionary = {'mother': 'lahi',
# 'maternal-aunt':'lahi',
# 'maternal-uncle':'lahi',
# 'father':'nehikupa',
# 'paternal-aunt':'nehikupa',
# 'paternal-uncle':'nehikupa',
# 'son':'nasike',
# 'daughter':'nasike',
# 'sister':'gowa',
# 'brother':'powuwani',
# 'grandson':'lomufusu',
# 'granddaughter':'lomufusu',
# 'maternal-grandmother': 'gewukahu',
# 'maternal-grandfather':'gewukahu',
# 'paternal-grandmother':'kifogama',
# 'paternal-grandfather':'kifogama'}

relationships = {'mother': [['12','6'],['11','6'],['7','1'],['6','1'],['5','1'],['8','3'],['9','3'],['10','3']],
'maternal-aunt': [['12','5'],['11','5']],
'maternal-uncle': [['12','7'],['11','7']],
'father': [['12','9'],['11','9'],['7','2'],['6','2'],['5','2'],['8','4'],['9','4'],['10','4']],
'paternal-aunt': [['12','8'],['11','8']],
'paternal-uncle': [['12','10'],['11','10']],
'son': [['1','7'],['2','7'],['3','10'],['4','10'],['3','9'],['4','9'],['6','12'],['9','12']],
'daughter': [['1','5'],['1','6'],['2','5'],['2','6'],['3','8'],['4','8'],['6','11'],['9','11']],
'sister': [['5','6'],['6','5'],['7','5'],['7','6'],['9','8'],['10','8'],['12','11']],
'brother': [['5','7'],['6','7'],['8','9'],['8','10'],['9','10'],['10','9'],['11','12']],
'grandson': [['1','12'],['2','12'],['3','12'],['4','12']],
'granddaughter': [['1','11'],['2','11'],['3','11'],['4','11']],
'maternal-grandmother': [['12','1'],['11','1']],
'maternal-grandfather': [['12','2'],['11','2']],
'paternal-grandmother': [['12','3'],['11','3']],
'paternal-grandfather': [['12','4'],['11','4']]
}

# picks a random pair of characters from a list of such pairs
def pick_pair(relationships,relationship):
    pair = random.choice(relationships[relationship])
    return pair

##############
# Utility functions
##############

# Converts message string to JSON string and sends to client_id.
def send_message_by_id(client_id,message):
    client = global_participant_data[client_id]['client_info']
    server.send_message(client,json.dumps(message))

# Checks that all clients listed in list_of_ids are still connected to the server
def all_connected(list_of_ids):
    connected_status = [id in global_participant_data for id in list_of_ids]
    if sum(connected_status)==len(list_of_ids):
        return True
    else:
        return False

# Called when a client drops out to notify any clients who are still connected
def notify_stranded(list_of_ids):
    for id in list_of_ids:
        if id in global_participant_data:
            #this will notify the participant and cause them to disconnect
            send_message_by_id(id,{"command_type":"PartnerDropout"})

######################
# Handling clients connecting, disconnecting, sending messages
######################

# Called for every client connecting, adds that client's data to global_participant_data,
# then sends them to the "Start" phase of the experiment
def new_client(client, server):
    client_id = client['id']
    print("New client connected and was given id %d" % client_id)
    global_participant_data[client_id] = {'client_info':client}
    #give them the instructions for the first phase
    enter_phase(client_id,"Start")

# Called for every client disconnecting
# Remove the client from unpaired_clients if appropriate and remove the client from global_participant_data
def client_left(client, server):
    client_id = client['id']
    print("Client(%d) disconnected" % client['id'])
    if client_id in unpaired_clients:
        unpaired_clients.remove(client_id)
    # If they have a partner, and if you are not leaving because you are at the End state,
    # notify partner that they have been stranded
    if 'partner' in global_participant_data[client_id]:
        if global_participant_data[client_id]['phase']!='End':
            partner = global_participant_data[client_id]['partner']
            notify_stranded([partner])
    del global_participant_data[client_id]

# Called when the server receives a message from the client.
# Prses the message to a dictionary using json.loads, reads off the response_type, and passes to handle_client_response
def message_received(client, server, message):
	print("Client(%d) said: %s" % (client['id'], message))
	#OK, now we have to handle the various possible responses
	response = json.loads(message)
	response_code =  response['response_type']
	handle_client_response(client['id'],response_code,response)

##########################
# Management of phases
##########################

# Look up the client's current phase and moves them to the next phase
def progress_phase(client_id):
	current_phase = global_participant_data[client_id]['phase']
	current_phase_i = phase_sequence.index(current_phase)
	next_phase_i = current_phase_i+1
	next_phase = phase_sequence[next_phase_i]
	enter_phase(client_id,next_phase)

# Triggers actions associated with each phase.
def enter_phase(client_id,phase):
    # Update the phase info for this client in the global dictionary
    global_participant_data[client_id]['phase']=phase

    # Nothing actually happens here, but in some experiments we will need to set stuff up
    # when the participant starts the experiment
    if phase=='Start':
        progress_phase(client_id)

    # Attempts to pair this client with anyone already in the waiting room
    elif phase=='PairParticipants':
        #send message to the client sending them to waiting room
        send_message_by_id(client_id,{"command_type":"WaitingRoom"})
        unpaired_clients.append(client_id) #add to the unpaired clients list
        start = time.time()
        end = time.time() - start
        while (len(unpaired_clients)%2 == 1):
            end = time.time() - start
            if end > 300.0: # if the time at the end of the trial is greater than 300 seconds, it's because the participant wasn't paired. Remove them from unpaired clients.
                stranded = unpaired_clients[0]
                unpaired_clients.remove(stranded)
        # If they can be immediately paired, do so and progress to next phase
        if (len(unpaired_clients)%2 == 0) and len(unpaired_clients) >= 2: #If there are exactly 2 people now in unpaired_clients, pair them
            unpaired_one = unpaired_clients[0]
            unpaired_two = unpaired_clients[1]
            #remove from unpaired list
            unpaired_clients.remove(unpaired_one)
            unpaired_clients.remove(unpaired_two)
            # Link them - mark them as each others' partner in global_participant_data
            global_participant_data[unpaired_one]['partner']=unpaired_two
            global_participant_data[unpaired_two]['partner']=unpaired_one

            # Both participants will work through a shared target list, so store that info
            # with both clients, then move them to the next phase
            shuffled_trials = block1_trials + reverse(block1_trials) + block2_trials + reverse(block2_trials)
            for c in [unpaired_one,unpaired_two]:
                global_participant_data[c]['trial_list'] = shuffled_trials
                global_participant_data[c]['shared_trial_counter'] = 0
                progress_phase(c)
    # Once paired with a partner clients will end up here; Interaction phase starts with instructions,
    # so just send those instructions to the client
    elif phase=='Interaction':
        print('Initalising for interaction')
        send_instructions(client_id,phase)

    # When they hit the end phase, the EndExperiment command will instruct the clients to end the experiment.
    elif phase=='End':
        send_message_by_id(client_id,{"command_type":"EndExperiment"})

#################
# Client loop, handling various client responses
#################

def handle_client_response(client_id,response_code,full_response):
    print('handle_client_response',client_id,response_code,full_response)

    # client is passing in a unique ID, simply associate that with this client
    if response_code=='CLIENT_INFO':
        global_participant_data[client_id]['participantID']=full_response['client_info']

    #interaction, instructions complete, can initiate actual interaction
    elif response_code=='INTERACTION_INSTRUCTIONS_COMPLETE':
        initiate_interaction(client_id)

    #response returned from director or matcher respectively
    elif response_code=='RESPONSE' and full_response['role']=='Director':
        handle_director_response(client_id,full_response)
    elif response_code=='RESPONSE' and full_response['role']=='Matcher':
        handle_matcher_response(client_id,full_response)

    #interaction feedback complete, next trial please
    elif response_code=='FINISHED_FEEDBACK':
        swap_roles_and_progress(client_id)

    #client reporting a non-responsive partner
    elif response_code=='NONRESPONSIVE_PARTNER':
        pass #not doing anything special with this - the participant reporting
		#the problem leaves, so for their partner it will be as if they have
		#dropped out

#################
# Interaction, handles trial progression etc
#################

# Runs when participants click to continue to the interaction phase.
def initiate_interaction(client_id):
    partner_id = global_participant_data[client_id]['partner']
    list_of_participants = [client_id,partner_id]
    #checking both players are still connected
    if not(all_connected(list_of_participants)):
        notify_stranded(list_of_participants)
    else:
        send_message_by_id(client_id,{"command_type":"WaitForPartner"})
        partner_role = global_participant_data[partner_id]['role']
        # check if your partner is ready to go before starting
        if partner_role=='ReadyToInteract':
            print('Starting interaction')
            # allocate random director and matcher, and run start_interaction_trial for both clients
            for client, role in zip(list_of_participants,shuffle(["Director", "Matcher"])):
                global_participant_data[client]['role'] = role
            start_interaction_trial(list_of_participants)
        else: # else mark you as ready to go, so you will wait for partner
            global_participant_data[client_id]['role']='ReadyToInteract'

# Interaction trial - sends director trial instruction to director and wait instruction to matcher
def start_interaction_trial(list_of_participants):
    #check everyone is still connected!
    if not(all_connected(list_of_participants)):
        notify_stranded(list_of_participants)
    else:
        #figure out who is the director
        director_id = [id for id in list_of_participants if global_participant_data[id]['role']=='Director'][0]
        #retrieve their trial list and trial counter
        trial_counter = global_participant_data[director_id]['shared_trial_counter']
        trial_list = global_participant_data[director_id]['trial_list']
        ntrials = len(trial_list)
        #check that the director has more trials to run - if not, move to next phase
        if trial_counter>=ntrials:
            for c in list_of_participants:
                progress_phase(c)
        else: #otherwise, if there are still trials to run
            # work out if we're in block 1 or block 2, to make the data handling easier later
            if trial_counter < 32:
                block = '1'
            else:
                block = '2'
            matcher_id = global_participant_data[director_id]['partner']
            matcher_participant_id = global_participant_data[matcher_id]['participantID']
            # work out some variables to build the trial with
            relationship = trial_list[trial_counter] # the relationship for this trial
            correct_label = kinship_dictionary[relationship] # the correct label for this relationship
            pair = pick_pair(relationships,relationship) # a random pair who have this relationship
            person1 = pair[0]
            person2 = pair[1]
            for c in list_of_participants:
                this_role = global_participant_data[c]['role']
                if this_role=='Director': #send the appropriate instruction to the Director
                    instruction_string = {"command_type":"Director",
                                            "relationship": relationship,
                                            "person1":person1,
                                            "person2":person2,
                                            "correct_label":correct_label,
                                            "kinship_terms_extra":kinship_terms_extra,
                                            "partner_id":matcher_participant_id,
                                            "block": block}
                else: #and tell the matcher to wait
                    instruction_string = {"command_type":"WaitForPartner"}
                send_message_by_id(c,instruction_string)

# sends matcher trial instruction to matcher and wait instruction to director
def handle_director_response(director_id,director_response):
    print('handle_director_response',director_response)
    matcher_id = global_participant_data[director_id]['partner']
    if not(all_connected([matcher_id])):
        notify_stranded([director_id])
    else:
        #note that director_response['response'] is the clue word the director sent
        director_participant_id=global_participant_data[director_id]['participantID']
        send_message_by_id(director_id,{"command_type":"WaitForPartner"})
        person1 = director_response['person1'] # person 1 is the same greeter from the director's trial
        person2 = director_response['person2'] # and person 2 is the same greetee
        choices = [x for x in images if x != person1] # remove the greetee from the button choices
        send_message_by_id(matcher_id,
                            {"command_type":"Matcher",
                                "relationship": director_response['relationship'],
                                "person1":person1,
                                "correct_person":person2,
                                "images":choices,
                                "label": director_response['response'],
                                "partner_id":director_participant_id,
                                "block": director_response['block']})

# When the matcher responds with their guess, send feedback to matcher + director.
def handle_matcher_response(matcher_id,matcher_response):
    print("in handle_matcher_response")
    director_id = global_participant_data[matcher_id]['partner']
    if not(all_connected([director_id,matcher_id])):
        notify_stranded([director_id,matcher_id])
    else:
        #easiest way to access what the target was is to look it up in the director's trial list
        trial_n = global_participant_data[director_id]['shared_trial_counter']
        target = global_participant_data[director_id]['trial_list'][trial_n]
        #info on the clue the director provided and the matcher's guess are included in the matcher's response
        target = matcher_response['correct_person']
        guess = matcher_response['response']
        person1 = matcher_response['person1']
        if target==guess:
            score=1
        else:
            score=0
        # send the necessary information to build the feedback trials
        matcher_feedback = {"command_type":"MatcherFeedback",
                            "score":score,
                            "person1":person1,
                            "correct_person":matcher_response["correct_person"],
                            "kin_term":matcher_response['label']}
        director_feedback = {"command_type":"DirectorFeedback",
                            "score":score,
                            "person1":person1,
                            "matcher_choice":guess,
                            "kin_term":matcher_response['label']}
        send_message_by_id(matcher_id,matcher_feedback)
        send_message_by_id(director_id,director_feedback)

# swaps the roles after the feedback trial ends
def swap_roles_and_progress(client_id):
    print('swap roles',client_id)
    partner_id = global_participant_data[client_id]['partner']
    if not(all_connected([client_id,partner_id])):
        notify_stranded([client_id,partner_id])
    else:
        #increment global counter - both participants will do this independently when they reach this point
        global_participant_data[client_id]['shared_trial_counter']+=1

        this_client_role = global_participant_data[client_id]['role']
        partner_role = global_participant_data[partner_id]['role']
        #If your partner is already ready, then switch roles and progress
        if partner_role=='WaitingToSwitch':
            if this_client_role=='Director': #if you were director for this trial then
                global_participant_data[client_id]['role'] = "Matcher" #next time you will be Matcher...
                global_participant_data[partner_id]['role'] = "Director" #..and your partner will be Director
            else:
                global_participant_data[client_id]['role'] = "Director" #otherwise the opposite
                global_participant_data[partner_id]['role'] = "Matcher"
            #next trial
            start_interaction_trial([client_id,partner_id])
        #Otherwise your partner is not yet ready, so just flag up that you are
        else:
            global_participant_data[client_id]['role'] = "WaitingToSwitch"

####################
# Instructions between blocks
####################

# send over a command_type Instructions message to the client, with instructon_type set to "Interaction"
def send_instructions(client_id,phase):
    if phase=='Interaction':
        #set role
        global_participant_data[client_id]['role'] = "ReadingInstructions"
        send_message_by_id(client_id,{"command_type":"Instructions","instruction_type":"Interaction"})

#######################
# Start up server
#######################

PORT=9020 #this will run on port 9020

#standard stuff here from the websocket_server code
print('starting up')
server = WebsocketServer(PORT,'0.0.0.0')
server.set_fn_new_client(new_client)
server.set_fn_client_left(client_left)
server.set_fn_message_received(message_received)
server.run_forever()


"""A custom hash function implementation that properly supports pytorch."""
import contextlib
import io
import os
import pickle
import struct
import sys
from pathlib import Path

from joblib.hashing import Hasher, NumpyHasher

Pickler = pickle._Pickler

class NoMemoizeHasher(Hasher):
    """A joblib hasher with all memoize features disabled."""

    def memoize(self, obj):  # noqa: ARG002
        # We skip memoization here, as it can cause some issues with nested objects.
        # https://github.com/joblib/joblib/issues/1283
        # In particular with the way cloning is implemented in tpcp, such bugs might occur more often than in other
        # applications.
        # My understanding is that memoization is used to reduce the size of the output.
        # Not using memoization is actually faster, but will fail with self referential objects.
        # (https://docs.python.org/3/library/pickle.html#pickle.Pickler.fast).
        # In these cases hashing will throw an `RecursionError`.
        # So it seems like it is a tradeoff between the two issues.
        # For now, we accept the recursion issue, as I think this might happen less often by accident.
        return

    def hash(self, obj, return_digest=True):  # noqa: A003
        """Get hash while handling some edgecases.

        Namely, this implementation fixes the following issues:

        - Because we skip memoization, we need to handle the case where the object is self-referential.
          We just catch the error and raise a more descriptive error message.

        """
        try:
            return super().hash(obj, return_digest)
        except RecursionError as e:
            raise ValueError(
                "The custom hasher used in tpcp does not support hashing of self-referential objects."
            ) from e

    def save_global(self, obj, name=None, pack=struct.pack):
        """We overwrite this method to fix the issue with objects defined in the `__main__` module.

        We need to handle the case where the object is defined in the `__main__` module.
        For some reason, this can lead to pickle issues.
        Based on some information I found, this should not happen, but it still does...
        To fix it, we detect, when an object is defined in `__main__` and temporarily add it to the "real" module
        representing the main function.
        Afterwards we do some cleanup.
        Not sure if really required, but it seems to work.
        Overall very hacky, but I don't see a better way to fix this.

        The implementation provided by the parent class (joblib.hashing.Hasher) should work, but it does not.
        I think the second call to `save_global` there is causing the issue.
        In our version we first fix the module before calling save_global again.

        Note that we are not using super() here, as we want to skip the faulty implementation in the parent class.

        We also need to modify the dispatch table to force our custom save_global method.
        See below this method for the implementation.
        """
        kwargs = {"name": name, "pack": pack}
        del kwargs["pack"]
        try:
            Pickler.save_global(self, obj, **kwargs)
        except pickle.PicklingError:
            modules_modified = []
            if getattr(obj, "__module__", None) == "__main__":
                try:
                    name = obj.__qualname__
                    to_add_obj = obj
                except AttributeError:
                    name = obj.__class__.__qualname__
                    to_add_obj = obj.__class__
                mod = sys.modules["__main__"]
                if not hasattr(mod, name):
                    modules_modified.append((mod, name))
                    setattr(mod, name, to_add_obj)
            try:
                Pickler.save_global(self, obj, **kwargs)
            finally:
                # Remove all new entries made to the main module.
                for mod, name in modules_modified:
                    delattr(mod, name)

    # We also need to rewrite the dispatch table to force our custom save_global method.
    dispatch = Hasher.dispatch.copy()
    # builtin
    dispatch[type(len)] = save_global
    # type
    dispatch[type(object)] = save_global
    # classobj
    dispatch[type(Hasher)] = save_global
    # function
    dispatch[type(pickle.dump)] = save_global

class NoMemoizeNumpyHasher(NoMemoizeHasher, NumpyHasher):
    """A joblib numpy hasher with all memoize features disabled."""

class NNHasher(NoMemoizeNumpyHasher):
    """A hasher that can handle torch models.

    Under the hood this uses the new implementation of `torch.save` to serialize the model.
    This produces consistent output.

    Note: I never did any performance checks with large models.
    """

    def __init__(self, hash_name="md5", coerce_mmap=False) -> None:
        super().__init__(hash_name, coerce_mmap)
        try:
            import torch  # pylint: disable=import-outside-toplevel

            self.torch = torch
        except ImportError:
            self.torch = None
        try:
            import tensorflow

            self.tensorflow = tensorflow
        except ImportError:
            self.tensorflow = None

    def save(self, obj):
        if self.torch and isinstance(obj, (self.torch.nn.Module, self.torch.Tensor)):
            b = b""
            buffer = io.BytesIO(b)
            self.torch.save(obj, buffer)
            self._hash.update(b)
            return

        if self.tensorflow and isinstance(obj, (self.tensorflow.keras.Model,)):
            # The normal tensorflow objects don't have a consistent hash.
            # Therefore, we need to serialize all relevant information.
            # I hope by using `serialize_keras_object` we get all important parts.
            # The `SharedObjectSavingScope` is required to make sure that shared objects are properly serialized.
            from tensorflow.python.keras.utils.generic_utils import SharedObjectSavingScope, serialize_keras_object

            with SharedObjectSavingScope():
                NumpyHasher.save(
                    self,
                    [obj.__class__.__name__, serialize_keras_object(obj), obj.get_weights()],
                )
            return
        NumpyHasher.save(self, obj)

# This function is modified based on
# https://github.com/joblib/joblib/blob/4dafaff788a3b5402acfed091558b4c511982959/joblib/hashing.py#L244
def custom_hash(obj, hash_name="md5", coerce_mmap=False):
    """Quick calculation of a hash to identify uniquely Python objects containing numpy arrays and torch models.

    This function is modified based on `joblib.hash` so that it can properly handle torch models.

    Parameters
    ----------
    obj
        The object to be hashed
    hash_name: 'md5' or 'sha1'
        Hashing algorithm used. sha1 is supposedly safer, but md5 is faster.
    coerce_mmap: boolean
        Make no difference between np.memmap and np.ndarray

    """
    valid_hash_names = ("md5", "sha1")
    if hash_name not in valid_hash_names:
        raise ValueError(f"Valid options for 'hash_name' are {valid_hash_names}. Got hash_name={hash_name!r} instead.")
    if "torch" in sys.modules or "tensorflow" in sys.modules:
        hasher = NNHasher(hash_name=hash_name, coerce_mmap=coerce_mmap)
    elif "numpy" in sys.modules:
        hasher = NoMemoizeNumpyHasher(hash_name=hash_name, coerce_mmap=coerce_mmap)
    else:
        hasher = NoMemoizeHasher(hash_name=hash_name)
    with Path(os.devnull).open("w") as devnull, contextlib.redirect_stdout(devnull):
        # Some object decide to print stuff to stdout when pickling.
        # As we potentially pickle a lot of objects, we don't want to see this.
        return hasher.hash(obj)

{'torch': 'torch', 'tensorflow': 'tensorflow', 'SharedObjectSavingScope': 'tensorflow.python.keras.utils.generic_utils.SharedObjectSavingScope', 'serialize_keras_object': 'tensorflow.python.keras.utils.generic_utils.serialize_keras_object'}