""" timed rotating text file """
# pylint: disable=too-many-instance-attributes,too-many-arguments
import os
import re
import time
from datetime import datetime
from io import TextIOWrapper
from stat import ST_MTIME

class TimedRotatingTextFile(TextIOWrapper):
    """ Rotating the text file at certain timed intervals. """

    def __init__(
            self,
            filename,
            when="d",
            interval=1,
            backup_count=0,
            mode="ab+",
            delay=False,
            utc=False,
            **kwargs,
    ):
        self.filename = filename
        self.when = when.upper()
        self.backup_count = backup_count
        self.mode = mode
        self.delay = delay
        self.utc = utc
        self.file = open(filename, self.mode)

        TextIOWrapper.__init__(self, self.file, **kwargs)

        if self.when in ["D", "MIDNIGHT"]:
            self.interval = 60 * 60 * 24  # One day
            self.suffix = "%Y-%m-%d"
            self.regex_match = r"^\d{4}-\d{2}-\d{2}(\.\w+)?$"
        elif self.when == "H":
            self.interval = 60 * 60  # One hour
            self.suffix = "%Y-%m-%d_%H"
            self.regex_match = r"^\d{4}-\d{2}-\d{2}_\d{2}(\.\w+)?$"
        elif self.when == "M":
            self.interval = 60  # one minute
            self.suffix = "%Y-%m-%d_%H-%M"
            self.regex_match = r"^\d{4}-\d{2}-\d{2}_\d{2}-\d{2}(\.\w+)?$"
        elif self.when == 'S':
            self.interval = 1  # one second
            self.suffix = "%Y-%m-%d_%H-%M-%S"
            self.regex_match = r"^\d{4}-\d{2}-\d{2}_\d{2}-\d{2}-\d{2}(\.\w+)?$"
        else:
            raise ValueError("Invalid rollover interval specified: %s" % self.when)

        self.interval = self.interval * interval  # multiply by units requested
        self.ext_match = re.compile(self.regex_match, re.ASCII)

        file_name = os.fspath(self.filename)
        # keep the absolute path, otherwise derived classes which use this
        # may come a cropper when the current directory changes
        self.base_file_name = os.path.abspath(file_name)
        if os.path.exists(self.base_file_name):
            modified_time = os.stat(self.base_file_name)[ST_MTIME]
        else:
            modified_time = int(time.time())

        self.rollover_at = self.compute_rollover(modified_time)

    def compute_rollover(self, current_time):
        """ Work out the rollover time based on the specified time. """
        result = current_time + self.interval
        if self.when == "MIDNIGHT":
            rotate_ts = 24 * 60 * 60
            if self.utc:
                struct_time = time.gmtime(current_time)
            else:
                struct_time = time.localtime(current_time)
            current_hour = struct_time[3]
            current_minute = struct_time[4]
            current_second = struct_time[5]

            rotate_time = rotate_ts - ((current_hour * 60 + current_minute) * 60 + current_second)

            if rotate_time < 0:
                rotate_time += 24 * 60 * 60
            result = current_time + rotate_time
        return result

    def should_rollover(self):
        """ Determine if rollover shoud occur. """
        current_time = int(time.time())
        if current_time >= self.rollover_at:
            return 1
        return 0

    def get_files_to_delete(self):
        """ Determine the files to delete when rolling over. """
        dir_name, base_name = os.path.split(self.base_file_name)
        files_name = os.listdir(dir_name)
        result = []
        prefix = base_name + "."
        prefix_len = len(prefix)
        for filename in files_name:
            if filename[:prefix_len] == prefix:
                suffix = filename[prefix_len:]
                if self.ext_match.match(suffix):
                    result.append(os.path.join(dir_name, filename))
        if len(result) < self.backup_count:
            return []
        result.sort()
        return result[: len(result) - self.backup_count]

    def do_rollover(self):
        """do a rollover; in this case, a date/time stamp is appended to the filename
        when the rollover happens.  However, you want the file to be named for the
        start of the interval, not the current time.  If there is a backup count,
        then we have to get a list of matching filenames, sort them and remove
        the one with the oldest suffix."""
        current_time = int((time.time()))
        dst_now = time.localtime(current_time)[-1]
        timed = self.rollover_at - self.interval
        if self.utc:
            time_tuple = time.gmtime(timed)
        else:
            time_tuple = time.localtime(timed)
            dst_then = time_tuple[-1]
            if dst_now != dst_then:
                addend = 3600 if dst_now else -3600
                time_tuple = time.localtime(timed + addend)
        dfn = self.base_file_name + "." + time.strftime(self.suffix, time_tuple)
        if os.path.exists(dfn):
            os.remove(dfn)
        if os.path.exists(self.base_file_name):
            os.rename(self.base_file_name, dfn)
        if self.backup_count > 0:
            for file_path in self.get_files_to_delete():
                os.remove(file_path)
        if not self.delay:
            self.close()
            self.file = open(self.base_file_name, self.mode)
            TextIOWrapper.__init__(self, self.file)
        new_rollover_at = self.compute_rollover(current_time)
        while new_rollover_at <= current_time:
            new_rollover_at = new_rollover_at + self.interval
        if self.when == "MIDNIGHT" and not self.utc:
            dst_at_rollover = time.localtime(new_rollover_at)[-1]
            if dst_now != dst_at_rollover:
                addend = -3600 if not dst_now else 3600
                new_rollover_at += addend
        self.rollover_at = new_rollover_at

    def write(self, line):
        if self.should_rollover():
            self.do_rollover()

        self.file.write(line.encode())
        # super().write(line)

    def close(self):
        self.file.close()
        super().close()

################
# MAIN SECTION #
################
if __name__ == "__main__":
    with TimedRotatingTextFile("/tmp/tmp.log", when="MIDNIGHT", backup_count=5) as filepath:
        filepath.write(f"Report Generated on {datetime.now().strftime('%A %d %B %Y %H:%M:%S')}\n")
        time.sleep(5)
        filepath.write("Another line\n")


# This module is responsible for rendering a scene.

from geometry import *
from itertools import chain

MAX_RECURSION_DEPTH = 2
BACKGROUND_COLOR = np.array([0.2, 0.2, 0.2])
# evaluate flags
FLAG_DEFAULT = 0x00
FLAG_LIGHTS_ONLY = 0x01
FLAG_LIGHTS_SKIP = 0x02

MAX_DISTANCE = 10 ** 8

class Ray:
    # ray has information about
    def __init__(self, point, vector):
        self.point = point
        self.vector = vector * (1 / np.linalg.norm(vector))

    def intersect_with_primitives(self, primitives):
        # we use intersection testing algorithms from
        # http://www.iquilezles.org/www/articles/intersectors/intersectors.htm

        nearest_dist = MAX_DISTANCE + 1
        nearest_prim = None

        for prim in primitives:
            if prim.kind == KIND_SPHERE:

                origin_center = self.point - prim.points[0]
                b = np.dot(origin_center, self.vector)
                c = np.linalg.norm(origin_center) ** 2 - prim.radius ** 2
                h = b ** 2 - c
                if h >= 0.0:  # TODO check if this algorithm fails when ray origin is inside sphere
                    h = np.sqrt(h)
                    if nearest_dist > -b - h > 0 and not np.isclose(-b - h, 0):
                        nearest_dist = -b - h
                        nearest_prim = prim
                    if nearest_dist > -b + h > 0 and not np.isclose(-b + h, 0):
                        nearest_dist = -b + h
                        nearest_prim = prim

            elif prim.kind == KIND_TRIANGLE:

                v0v1 = prim.points[4]  # point_b - point_a
                v0v2 = prim.points[5]  # point_c - point_a

                rov0 = self.point - prim.points[0]

                n = prim.points[6]  # cross(v0v1,v0v2)

                q = np.cross(rov0, self.vector)

                a = np.dot(self.vector, n)
                if not a == 0:
                    d = 1.0 / a

                    u = d * np.dot(-q, v0v2)  # first barycentric coordinate
                    v = d * np.dot(q, v0v1)  # second barycentric coordinate
                    t = d * np.dot(-n, rov0)  # distance

                    if 1.0 >= u >= 0.0 and 1.0 >= v >= 0.0 and (
                            u + v) <= 1.0 and nearest_dist > t > 0 and not np.isclose(t, 0):
                        nearest_dist = t
                        nearest_prim = prim

        if nearest_dist > MAX_DISTANCE:  # no intersection
            return None

        return nearest_dist, nearest_prim

    def intersect(self, scene):
        # calculate the earliest intersection with a primitive of the scene

        bounding_primitives = set(filter(lambda key: key is not None, scene.primitives.keys()))

        nearest_bounding_primitive_intersection = self.intersect_with_primitives(bounding_primitives)

        primitives_to_check = set()

        if scene.primitives.get(None):
            primitives_to_check = primitives_to_check.union(scene.primitives.get(None))

        if nearest_bounding_primitive_intersection:
            primitives_to_check = primitives_to_check.union(scene.primitives.get(nearest_bounding_primitive_intersection[1]))

        intersection = self.intersect_with_primitives(primitives_to_check)

        if not intersection:
            return None

        nearest_dist, nearest_prim = intersection

        nearest_point = nearest_dist * self.vector + self.point
        normal = np.array([0.0, 0.0, 0.0])
        if nearest_prim.kind == KIND_SPHERE:
            normal = nearest_point - nearest_prim.points[0]
            normal /= np.linalg.norm(normal)
        elif nearest_prim.kind == KIND_TRIANGLE:
            normal = nearest_prim.normal

        return nearest_prim, nearest_point, nearest_dist, normal

    def evaluate(self, scene, recursion_depth, flag=FLAG_DEFAULT):
        # evaluates color and brightness of a given ray recursively for a set depth of recursion
        # [1.0, 1.0, 1.0] is white

        if recursion_depth > MAX_RECURSION_DEPTH:
            # rays die after exceeding depth of recursion
            return np.array([0.0, 0.0, 0.0])

        r = self.intersect(scene)

        if not r:
            return BACKGROUND_COLOR

        prim, coord, length, normal = r

        if prim.is_light_source:
            if flag & FLAG_LIGHTS_SKIP:
                return np.array([0.0, 0.0, 0.0])
            # if the ray hits a light source we can stop iterating
            return prim.color / (length + 1) ** 2

        if flag & FLAG_LIGHTS_ONLY:
            return np.array([0.0, 0.0, 0.0])

        if prim.is_reflective:
            f = Ray(coord, self.vector - 2 * np.dot(self.vector, normal) * normal).evaluate(scene, recursion_depth + 1)
            f = np.multiply(f, prim.color)
            return f / (length + 1) ** 2

        incoming_lightrays = []

        # diffuse lightrays heading to source
        for light_prim in scene.lights:
            if light_prim.kind == KIND_SPHERE:
                light_point = light_prim.points[0]
            elif light_prim.kind == KIND_TRIANGLE:
                light_point = light_prim.points[3]
            else:
                continue

            light_vector = light_point - coord
            light_vector /= np.linalg.norm(light_vector)
            c = Ray(coord, light_vector).evaluate(scene, MAX_RECURSION_DEPTH, flag=FLAG_LIGHTS_ONLY)

            # intensity has to be scaled down relative to the angle
            cos_light_normal = np.abs(np.dot(normal, light_vector))

            # diffuse light seeking lightsource
            e = c * cos_light_normal

            incoming_lightrays.append(np.multiply(e, prim.color))
            # incoming_lightrays.append(c)

            if prim.specular > 1:
                middle = -self.vector + light_vector
                middle /= np.linalg.norm(middle)
                cos_middle_normal = np.abs(np.dot(normal, middle))
                s = prim.shininess * c * cos_middle_normal ** prim.specular
                incoming_lightrays.append(s)

        # diffuse lightray in random direction (introduces noise)
        delta = np.random.randn(3)
        delta /= np.linalg.norm(delta)
        diffuse_vector = normal + delta
        d = Ray(coord, diffuse_vector).evaluate(scene, MAX_RECURSION_DEPTH - 1, flag=FLAG_LIGHTS_SKIP)
        d = d * np.abs(np.dot(normal, diffuse_vector)) / (np.linalg.norm(diffuse_vector))
        incoming_lightrays.append(np.multiply(d, prim.color))
        # incoming_lightrays.append(d)

        result = np.sum(incoming_lightrays, axis=0)
        # result = np.multiply(np.sum(incoming_lightrays, axis=0), prim.color)
        return result / (length + 1) ** 2
        # if prim.shininess == 0.0:
        #    pass

        # TODO:
        # depending on the texture of the primitive, calculate more rays and combine result
        # light_reflection = Ray(coords, self.vector - 2*np.dot(self.vector,normal)*normal)
        #                       .evaluate(scene,recursion_depth+1)
        # light_source = Ray(coords, scene.light.coords - coords).eval...
        # light_diffuse = ...
        # light_refraction...
        # return np.quick_maths(light...)

def ray_iterator(camera, resolution, start_row, end_row, start_column, end_column):
    # this function enables iteration through every pixel on the FOV for the specified part to be rendered
    # yield ray,(x,y)...

    half_fov_width = np.tan(camera.fov_width_angle / 2)
    ratio = resolution[1] / resolution[0]
    half_fov_height = half_fov_width * ratio
    top_left = np.array([1, -half_fov_width, half_fov_height])
    down = np.array([0, 0, -2 * half_fov_height])
    right = np.array([0, 2 * half_fov_width, 0])

    for y in range(end_row - start_row + 1):
        for x in range(resolution[0] - start_column if y == 0 else (end_column + 1 if y == end_row else resolution[0])):
            yield Ray(camera.point,
                      top_left + (start_column + x if y == 0 else x) / resolution[0] * right + (start_row + y) /
                      resolution[1] * down), (x, y)

def render_scene(scene: Scene, **config):
    # Iterates through all pixels in viewport, traces rays and draws to bitmap. For the documentation of the config
    # parameters, look at main.py. The method is expected to return a 2D array containing the exact block as
    # requested. The size of the first and last column may vary - therefore take the parameters "firstColumn" and
    # "lastColumn" into account

    w = config["width"]
    h = config["height"]

    start_row = config["start_row"]
    start_column = config["start_column"]

    end_row = config["end_row"]
    end_column = config["end_column"]

    pixels = []  # Don't build the array in advance, because the pixel array can have a weird shape (half columns and such things)

    for ray, pixel in ray_iterator(scene.camera, [w, h], start_row, end_row, start_column, end_column):
        color = np.uint8(np.clip(ray.evaluate(scene, 0) * 255, a_min=0, a_max=255))
        x = pixel[0]
        y = pixel[1]
        if len(pixels) <= y:
            pixels.append([])
        row = pixels[y]
        if (len(row)) <= x:
            row.append([])
        row[x] = color

    return pixels


from enum import Enum

class BookingClass:
    def __init__(self, bselem):
        self.value = bselem.attrs["value"]
        self.string = bselem.text.replace('\xa0', ' ')

    def print(self):
        print(self.value, self.string)

class BookingDate:
    def __init__(self, data):
        self.value = data["DATE_VALUE"]
        self.string = data["DATE_STRING"].replace('\xa0', ' ')

    def print(self):
        print(self.value, self.string)

class BookingBranch:
    def __init__(self, data):
        self.name = data["BRANCH_NAME"].replace('\xa0', ' ')
        self.expiredate = data["EFFECTIVE_DATE"]
        self.value = data["LOCATION"]

    def print(self):
        print(self.value, self.name, self.expiredate)

class BookingTime:
    def __init__(self, data):
        self.value = data["SEGMENT"]
        self.seatnum = data["ALLSEAT"]
        self.remaining = data["SEATNUM"]
        self.openrate = data["OPEN_RATE"]
        self.starttime = data["INIT_TIME"]
        self.endtime = data["END_TIME"]
        self.stop = True if data["OFFCLASS"] == "1" else False
        self.reserved = True if data["HASCLASS"] == "1" else False

    def print(self):
        print(self.value, self.starttime, "~", self.endtime,
              (str(int(self.seatnum) - int(self.remaining)) + "/" + self.seatnum), ("Close" if self.stop else ""),
              ("Reserved" if self.reserved else ""))

class BookingResponse(Enum):
    OK = "預約成功"
    NOCONTINOUS = "無連續座位"
    FAIL = "預約失敗"
    NOTAVAILABLE = "該時間不可預約"
    TIMEOUT = "Timeout錯誤"

class CheckingResponse(Enum):
    OK = "檢查OK"
    CLASSNOTREACHABLE = "可上課程取得發生錯誤"
    BRANCHNOTREACHABLE = "可上課地點取得發生錯誤"
    DATENOTREACHABLE = "可上課日期取得發生錯誤"
    TIMENOTREACHABLE = "可上課時段取得發生錯誤"

class TokenResponse(Enum):
    OK = "成功獲得Token"
    FAIL = "獲得TOKEN失敗"

class AsciiArt:
    logo = """

////////////////////////////////////////////////////////////////////////////
_/_/_/_/_/  _/    _/  _/_/_/                             author: JasonLiTW
   _/      _/  _/    _/    _/                                      2018.07
  _/      _/_/      _/_/_/
 _/      _/  _/    _/    _/
_/      _/    _/  _/_/_/

      _/_/    _/    _/  _/_/_/_/_/    _/_/
   _/    _/  _/    _/      _/      _/    _/
  _/_/_/_/  _/    _/      _/      _/    _/
 _/    _/  _/    _/      _/      _/    _/
_/    _/    _/_/        _/        _/_/

    _/_/_/      _/_/      _/_/    _/    _/  _/_/_/  _/      _/    _/_/_/
   _/    _/  _/    _/  _/    _/  _/  _/      _/    _/_/    _/  _/
  _/_/_/    _/    _/  _/    _/  _/_/        _/    _/  _/  _/  _/  _/_/
 _/    _/  _/    _/  _/    _/  _/  _/      _/    _/    _/_/  _/    _/
_/_/_/      _/_/      _/_/    _/    _/  _/_/_/  _/      _/    _/_/_/
////////////////////////////////////////////////////////////////////////////

    """
    success = """

////////////////////////////////////////////////////////////////////////////
      _/_/_/  _/    _/    _/_/_/    _/_/_/  _/_/_/_/    _/_/_/    _/_/_/  _/
   _/        _/    _/  _/        _/        _/        _/        _/        _/
    _/_/    _/    _/  _/        _/        _/_/_/      _/_/      _/_/    _/
       _/  _/    _/  _/        _/        _/              _/        _/
_/_/_/      _/_/      _/_/_/    _/_/_/  _/_/_/_/  _/_/_/    _/_/_/    _/
////////////////////////////////////////////////////////////////////////////

    """


# noteParse contains methods that are used in breaking
#  apart evernote .enex files and giving them a parsed
#  and more accessible format

import xml.etree.ElementTree as ET
import note
import book
import os

def parseNote(fileName):
    tree = ET.parse(fileName)
    enExport = tree.getroot()
    noteTree = enExport[0]
    title = noteTree.findall('title')[0].text
    author = noteTree.findall('note-attributes')[0].findall('author')[0].text
    created = noteTree.findall('created')[0].text
    updated = noteTree.findall('updated')[0].text

    contentStr = noteTree.findall('content')[0].text
    contentStr = contentStr.replace("""<?xml version="1.0" encoding="UTF-8" standalone="no"?>""", "")
    contentStr = contentStr.replace("""<!DOCTYPE en-note SYSTEM "http://xml.evernote.com/pub/enml2.dtd">""","")
    contentTree = ET.fromstring(contentStr)

    resourceListTree = noteTree.findall('resource')

    noteObj = note.Note()
    noteObj.author = author
    noteObj.title = title
    noteObj.created = created
    noteObj.updated = updated
    noteObj.contentTree = contentTree
    noteObj.resources = resourceListTree

    for tagNode in noteTree.findall('tag'):
        noteObj.addTag(tagNode.text)

    return noteObj

def parseNotebook(fileName):
    tree = ET.parse(fileName)
    enExport = tree.getroot()
    noteTree = enExport[0]

    bookObj = book.Book()
    bookObj.title = os.path.splitext(os.path.basename(fileName))[0]

    for noteTree in enExport.findall('note'):
        title = noteTree.findall('title')[0].text
        try:
            author = noteTree.findall('note-attributes')[0].findall('author')[0].text
        except:
            author = "Joey Kilgore"
        created = noteTree.findall('created')[0].text
        updated = noteTree.findall('updated')[0].text

        contentStr = noteTree.findall('content')[0].text
        contentStr = contentStr.replace("""<?xml version="1.0" encoding="UTF-8" standalone="no"?>""", "")
        contentStr = contentStr.replace("""<!DOCTYPE en-note SYSTEM "http://xml.evernote.com/pub/enml2.dtd">""","")
        contentStr = contentStr.replace(" ", "")
        contentStr = ''.join([i if ord(i) < 128 else '' for i in contentStr]) # remove non ascii
        contentTree = ET.fromstring(contentStr.strip())

        resourceListTree = noteTree.findall('resource')

        noteObj = note.Note()
        noteObj.author = author
        noteObj.title = title
        noteObj.created = created
        noteObj.updated = updated
        noteObj.contentTree = contentTree
        noteObj.resources = resourceListTree

        bookObj.notes.append(noteObj)
        bookObj.author = noteObj.author

    return bookObj

def recursivePrintTree(prefix, node):
    print(prefix + "|-->" + node.tag)
    for child in node:
        recursivePrintTree(prefix+"    ", child)

def printTree(fileName):
    tree = ET.parse(fileName)
    root = tree.getroot()
    print(fileName)
    recursivePrintTree("", root)


import glob

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import scipy.stats as stats
import seaborn as sns
import statsmodels.api as sm
import tqdm
from matplotlib.font_manager import FontProperties
from matplotlib.ticker import MaxNLocator
from numpy.linalg import norm
from sklearn import preprocessing
from sklearn.decomposition import PCA

class GenerateTestPlots:
    """Generates plots using data collected during testing."""

    def __init__(
        self,
        block_type="default",
        mode="default",
        num_actions=3,
        num_states=42,
        num_hidden_units=128,
        gamma=0.96,
    ):
        """Initialize the plot generator.

        Args:
            block_type (str, optional): [description]. Argument stating if block length
            is fixed or not. Defaults to "default".
            mode (str, optional): Argument stating if the trials can be stimulated.
            Defaults to "default".
            num_actions (int, optional): Number of allowed actions. Defaults to 3.
            num_states (int, optional): Number of states. Defaults to 42.
            num_hidden_units (int, optional): Number of hidden units of the LSTMS.
            Defaults to 128.
            gamma (float, optional): Discount factor. Defaults to 0.96.
        """

        testing_data_path = f"testing_data/block_type-{block_type}/mode-{mode}/"

        testing_files = [
            f for f in glob.glob(testing_data_path + "*.npz", recursive=True)
        ]

        self.current_trial_times = []
        self.trial_count_in_blocks = []
        self.actions = []
        self.choices = []
        self.actor_logits = []
        self.reward_probs = []
        self.rewards = []
        self.is_rewardeds = []
        self.values = []
        self.rpes = []
        self.critic_hs = []
        self.critic_cs = []
        self.actor_hs = []
        self.actor_cs = []
        self.is_stimulateds = []

        for file in tqdm.tqdm(testing_files):
            testing_data = np.load(file)
            self.current_trial_times.append(testing_data["current_trial_times"])
            self.trial_count_in_blocks.append(testing_data["trial_count_in_blocks"])
            self.actions.append(testing_data["actions"])
            self.choices.append(testing_data["choices"])
            self.actor_logits.append(np.float32(testing_data["actor_logits"]))
            self.reward_probs.append(np.float32(testing_data["reward_probs"]))
            self.rewards.append(np.float32(testing_data["rewards"]))
            self.is_rewardeds.append(testing_data["is_rewardeds"])
            self.values.append(np.float32(testing_data["values"]))
            self.rpes.append(np.float32(testing_data["rpes"]))
            self.critic_hs.append(np.float32(testing_data["critic_hs"]))
            self.critic_cs.append(np.float32(testing_data["critic_cs"]))
            self.actor_hs.append(np.float32(testing_data["actor_hs"]))
            self.actor_cs.append(np.float32(testing_data["actor_cs"]))
            self.is_stimulateds.append(testing_data["is_stimulateds"])

        self.current_trial_times = np.array(self.current_trial_times).reshape(-1)
        self.trial_count_in_blocks = np.array(self.trial_count_in_blocks).reshape(-1)
        self.actions = np.array(self.actions).reshape(-1)
        self.choices = np.array(self.choices).reshape(-1)
        self.actor_logits = np.array(self.actor_logits).reshape(-1, num_actions)
        self.reward_probs = np.array(self.reward_probs).reshape(-1)
        self.rewards = np.array(self.rewards).reshape(-1)
        self.is_rewardeds = np.array(self.is_rewardeds).reshape(-1)
        self.values = np.array(self.values).reshape(-1)
        self.rpes = np.array(self.rpes).reshape(-1)
        self.critic_hs = np.array(self.critic_hs).reshape(-1, num_hidden_units)
        self.critic_cs = np.array(self.critic_cs).reshape(-1, num_hidden_units)
        self.actor_hs = np.array(self.actor_hs).reshape(-1, num_hidden_units)
        self.actor_cs = np.array(self.actor_cs).reshape(-1, num_hidden_units)
        self.is_stimulateds = np.array(self.is_stimulateds).reshape(-1)
        self.num_states = num_states
        self.num_hidden_units = num_hidden_units
        self.gamma = gamma

    def behavior(self, start_num_trials=500, num_tr=200, save=None):

        print(
            f"reward rate = "
            f"{np.mean(self.is_rewardeds[self.current_trial_times == 1])}"
        )

        end_num_trials = start_num_trials + num_tr

        high_prob = self.reward_probs[self.current_trial_times == 2]
        rewarded_trials = self.is_rewardeds[self.current_trial_times == 2]
        chosen_side = 2 * self.choices[self.current_trial_times == 2] - 1
        rew_side = chosen_side * rewarded_trials
        logits = (
            self.actor_logits[:, 2][self.current_trial_times == 1]
            - self.actor_logits[:, 1][self.current_trial_times == 1]
        )
        logits = np.roll(logits, 1)
        max_logit, mid_logit = np.max(logits), (np.max(logits) + np.min(logits)) / 2

        switch_high = np.where(np.diff(high_prob) != 0)[0] + 1

        print(
            f"mean block length ="
            f"{np.mean(np.diff(switch_high))}"
            "\u00B1"
            f"{np.std(np.diff(switch_high))}"
        )

        plt.figure(figsize=(16, 5))
        plt.scatter(
            np.arange(start_num_trials, end_num_trials),
            high_prob[start_num_trials:end_num_trials] * (max_logit - mid_logit) * 1.5
            + mid_logit,
            s=40,
            color="black",
            label="high probability side",
        )
        plt.scatter(
            np.arange(start_num_trials, end_num_trials),
            chosen_side[start_num_trials:end_num_trials] * (max_logit - mid_logit) * 1.3
            + mid_logit,
            s=20,
            color="blue",
            label="chosen side",
        )
        plt.scatter(
            np.arange(start_num_trials, end_num_trials),
            rew_side[start_num_trials:end_num_trials] * (max_logit - mid_logit) * 1.1
            + mid_logit,
            s=30 * np.abs(rew_side[start_num_trials:end_num_trials]),
            color="green",
            label="rewarded side",
        )
        plt.plot(
            np.arange(start_num_trials, end_num_trials),
            logits[start_num_trials:end_num_trials],
            color="aqua",
            label="value averaged over time",
        )

        fontP = FontProperties()
        fontP.set_size("xx-large")
        plt.legend(
            bbox_to_anchor=(0.0, 1.02, 1.0, 0.102),
            loc="lower left",
            ncol=3,
            mode="expand",
            borderaxespad=0.0,
            prop=fontP,
        )
        plt.xlabel("Trial number", fontsize=20)

        if save:
            plt.savefig(save, bbox_inches="tight")

    def stay_probability(self, mode=None, save=None):
        if mode == "optogenetic":

            rewarded_trials = self.is_rewardeds[self.current_trial_times == 2]
            chosen_side = 2 * self.choices[self.current_trial_times == 2] - 1
            stim_trials = self.is_stimulateds[self.current_trial_times == 2]

            unrewarded_trials = np.abs(rewarded_trials - 1)
            next_choice = np.roll(chosen_side, -1)
            next_id = np.roll(stim_trials, -1)
            norm_chosen_side = chosen_side[next_id == 0]
            norm_next_choice = next_choice[next_id == 0]
            norm_rewarded_trials = rewarded_trials[next_id == 0]
            norm_unrewarded_trials = unrewarded_trials[next_id == 0]
            norm_return_rew = norm_rewarded_trials * (
                (norm_next_choice == norm_chosen_side) * 1
            )
            norm_prob_rew_return = np.sum(norm_return_rew) / np.sum(norm_rewarded_trials)
            norm_return_unrew = norm_unrewarded_trials * (
                (norm_next_choice == norm_chosen_side) * 1
            )
            norm_prob_unrew_return = np.sum(norm_return_unrew) / np.sum(
                norm_unrewarded_trials
            )

            opto_chosen_side = chosen_side[next_id == 1]
            opto_next_choice = next_choice[next_id == 1]
            opto_rewarded_trials = rewarded_trials[next_id == 1]
            opto_unrewarded_trials = unrewarded_trials[next_id == 1]
            opto_return_rew = opto_rewarded_trials * (
                (opto_next_choice == opto_chosen_side) * 1
            )
            opto_prob_rew_return = np.sum(opto_return_rew) / np.sum(opto_rewarded_trials)
            opto_return_unrew = opto_unrewarded_trials * (
                (opto_next_choice == opto_chosen_side) * 1
            )
            opto_prob_unrew_return = np.sum(opto_return_unrew) / np.sum(
                opto_unrewarded_trials
            )

            fig = plt.figure(figsize=(9, 3))
            plt.subplot(1, 2, 1)
            plt.bar(
                [1, 2, 4, 5],
                [
                    norm_prob_rew_return,
                    opto_prob_rew_return,
                    norm_prob_unrew_return,
                    opto_prob_unrew_return,
                ],
                width=0.6,
                color=["k", "aqua", "k", "aqua"],
            )
            plt.xticks(
                [1.5, 4.5], ["current trial\n reward", "current trial\n no reward"]
            )
            plt.ylim([0.2, 1.05])
            plt.ylabel("probability of return")

            norm_prob_rew_return = 0.0
            norm_prob_unrew_return = 0.0

            opto_prob_rew_return = 0.0
            opto_prob_unrew_return = 0.0

            unrewarded_trials = np.abs(rewarded_trials - 1)
            next_choice = np.roll(chosen_side, -1)
            next_id = stim_trials
            norm_chosen_side = chosen_side[next_id == 0]
            norm_next_choice = next_choice[next_id == 0]
            norm_rewarded_trials = rewarded_trials[next_id == 0]
            norm_unrewarded_trials = unrewarded_trials[next_id == 0]
            norm_return_rew = norm_rewarded_trials * (
                (norm_next_choice == norm_chosen_side) * 1
            )
            norm_prob_rew_return += np.sum(norm_return_rew) / np.sum(norm_rewarded_trials)
            norm_return_unrew = norm_unrewarded_trials * (
                (norm_next_choice == norm_chosen_side) * 1
            )
            norm_prob_unrew_return += np.sum(norm_return_unrew) / np.sum(
                norm_unrewarded_trials
            )

            opto_chosen_side = chosen_side[next_id == 1]
            opto_next_choice = next_choice[next_id == 1]
            opto_rewarded_trials = rewarded_trials[next_id == 1]
            opto_unrewarded_trials = unrewarded_trials[next_id == 1]
            opto_return_rew = opto_rewarded_trials * (
                (opto_next_choice == opto_chosen_side) * 1
            )
            opto_prob_rew_return += np.sum(opto_return_rew) / np.sum(opto_rewarded_trials)
            opto_return_unrew = opto_unrewarded_trials * (
                (opto_next_choice == opto_chosen_side) * 1
            )
            opto_prob_unrew_return += np.sum(opto_return_unrew) / np.sum(
                opto_unrewarded_trials
            )

            # fig = plt.figure(figsize=(3,3))
            plt.subplot(1, 2, 2)
            plt.bar(
                [1, 2, 4, 5],
                [
                    norm_prob_rew_return,
                    opto_prob_rew_return,
                    norm_prob_unrew_return,
                    opto_prob_unrew_return,
                ],
                width=0.6,
                color=["k", "aqua", "k", "aqua"],
            )
            plt.xticks(
                [1.5, 4.5], ["previous trial\n reward", "previous trial\n no reward"]
            )
            plt.ylim([0.2, 1.05])
            plt.ylabel("probability of return")

        else:
            prob_rew_return = 0.0
            prob_unrew_return = 0.0

            rewarded_trials = self.is_rewardeds[self.current_trial_times == 2]
            chosen_side = 2 * self.choices[self.current_trial_times == 2] - 1

            unrewarded_trials = (rewarded_trials - 1) * -1
            next_choice = np.roll(chosen_side, -1)
            return_rew = rewarded_trials * ((next_choice == chosen_side) * 1)
            prob_rew_return += np.sum(return_rew) / np.sum(rewarded_trials)
            return_unrew = unrewarded_trials * ((next_choice == chosen_side) * 1)
            prob_unrew_return += np.sum(return_unrew) / np.sum(unrewarded_trials)

            plt.subplot(1, 2, 1)
            plt.bar(
                [1, 1.5],
                [prob_rew_return, prob_unrew_return],
                width=0.2,
                color=["green", "red"],
            )
            plt.xticks([1, 1.5], ["previously\n rewarded", "previously\n unrewarded"])
            plt.ylabel("probability of return")
            plt.ylim([0.0, 1.1])

        if save:
            plt.savefig(save, bbox_inches="tight")

    def choice_regression(self, mode=None, trials_back=11, save=None):

        if mode == "optogenetic":
            rewarded_trials = self.is_rewardeds[
                (self.current_trial_times == 2) & (self.choices != -1)
            ]
            chosen_side = (
                2 * self.choices[(self.current_trial_times == 2) & (self.choices != -1)]
                - 1
            )
            stim_trials = self.is_stimulateds[
                (self.current_trial_times == 2) & (self.choices != -1)
            ]

            reward_mat = np.zeros((trials_back - 1, len(rewarded_trials)))
            reward_vect = chosen_side * rewarded_trials
            for i in np.arange(1, trials_back):
                reward_mat[i - 1, :] = np.roll(reward_vect, i)

            # makes unreward matrix
            unrewarded_trials = np.abs(rewarded_trials - 1)
            unreward_mat = np.zeros((trials_back - 1, len(unrewarded_trials)))
            unreward_vec = chosen_side * unrewarded_trials
            for i in np.arange(1, trials_back):
                unreward_mat[i - 1, :] = np.roll(unreward_vec, i)

            # makes laser matrix
            laser_mat = np.zeros((trials_back - 1, len(rewarded_trials)))
            for i in np.arange(1, trials_back):
                laser_mat[i - 1, :] = np.roll(stim_trials, i)

            y = chosen_side
            x = np.concatenate(
                (
                    np.ones([1, len(y)]),
                    reward_mat,
                    unreward_mat,
                    reward_mat * laser_mat,
                    unreward_mat * laser_mat,
                    laser_mat,
                ),
                axis=0,
            )

            y_new = np.asarray((y + 1) / 2, dtype=int)

            log_reg = sm.Logit(y_new, x.T).fit()

            # Plots regression
            reward_coefs = log_reg.params[1:trials_back]
            unreward_coefs = log_reg.params[trials_back : int(trials_back * 2 - 1)]
            rewlaser_coefs = log_reg.params[
                int(trials_back * 2 - 1) : int(trials_back * 3 - 2)
            ]
            norewlaser_coefs = log_reg.params[
                int(trials_back * 3 - 2) : int(trials_back * 4 - 3)
            ]

            fig = plt.figure(figsize=(12, 6))
            plt.plot(reward_coefs, "b", label="rewarded trials no stimulation")
            plt.plot(unreward_coefs, "r", label="unrewarded trials no stimulation")
            plt.plot(
                reward_coefs + rewlaser_coefs,
                linestyle="dotted",
                color="b",
                label="rewarded trials with stimulation",
            )
            plt.plot(
                unreward_coefs + norewlaser_coefs,
                linestyle="dotted",
                color="r",
                label="unrewarded trials with stimulation",
            )
            # plt.plot(laser_coefs,'k')
            plt.axhline(y=0, linestyle="dotted", color="gray")
            plt.xticks(
                np.arange(0, trials_back - 1, 2),
                [str(i) for i in np.arange(-1, -trials_back, -2)],
            )
            plt.xlabel("trials back")
            plt.ylabel("regression coefficients")
            plt.legend()

        else:

            rewarded_trials = self.is_rewardeds[
                (self.current_trial_times == 2) & (self.choices != -1)
            ]
            chosen_side = (
                2 * self.choices[(self.current_trial_times == 2) & (self.choices != -1)]
                - 1
            )

            reward_mat = np.zeros((trials_back - 1, len(rewarded_trials)))
            reward_vect = chosen_side * rewarded_trials
            for i in np.arange(1, trials_back):
                reward_mat[i - 1, :] = np.roll(reward_vect, i)

            # makes unreward matrix
            unrewarded_trials = np.abs(rewarded_trials - 1)
            unreward_mat = np.zeros((trials_back - 1, len(unrewarded_trials)))
            unreward_vec = chosen_side * unrewarded_trials
            for i in np.arange(1, trials_back):
                unreward_mat[i - 1, :] = np.roll(unreward_vec, i)

            y = chosen_side
            x = np.concatenate((np.ones([1, len(y)]), reward_mat, unreward_mat), axis=0)

            y_new = np.asarray((y + 1) / 2, dtype=int)

            log_reg = sm.Logit(y_new, x.T).fit()

            fig = plt.figure(figsize=(10, 3))
            plt.plot(
                np.arange(0, trials_back - 1),
                log_reg.params[1:trials_back],
                color="blue",
                linewidth=1,
                label="model - rewarded trials",
            )
            plt.plot(
                np.arange(0, trials_back - 1),
                log_reg.params[trials_back:None],
                color="red",
                linewidth=1,
                label="model - unrewarded trials",
            )
            plt.axhline(y=0, linestyle="dotted", color="gray")
            plt.legend()
            plt.xticks(
                np.arange(0, trials_back - 1, 2),
                [str(i) for i in np.arange(-1, -trials_back, -2)],
            )
            plt.xlabel("trials back")
            plt.ylabel("regression coefficients")

        if save:
            plt.savefig(save, bbox_inches="tight")

    def block_value_plot(self, save=None):

        high_prob = self.reward_probs[(self.current_trial_times == 2)]
        rewarded_trials = self.is_rewardeds[(self.current_trial_times == 2)]
        chosen_side = 2 * self.choices[(self.current_trial_times == 2)] - 1
        value = np.array(
            [
                self.values[(self.current_trial_times == (i + 1) % self.num_states)]
                for i in range(self.num_states)
            ]
        ).T

        switch_high = np.where(np.diff(high_prob) != 0)[0] + 1

        early_trials = switch_high[0 : len(switch_high) - 2]
        late_trials = early_trials + 14  # switch_high[1:-1] - 1
        # middle_trials=(early_trials+late_trials)/2
        middle_trials = early_trials + 4

        print(
            f"reward rate on last trial of a block = "
            f"{np.round(np.mean(rewarded_trials[early_trials - 1]), 2)}"
        )
        print(
            f"reward rate on first trial of a block = "
            f"{np.round(np.mean(rewarded_trials[early_trials]), 2)}"
        )
        # block identity of each trial
        block_iden = high_prob[early_trials]

        # indeces for block and actual choice
        left_left_early = early_trials[
            (block_iden == 1) & (chosen_side[early_trials] == 1)
        ]
        left_right_early = early_trials[
            (block_iden == 1) & (chosen_side[early_trials] == -1)
        ]
        right_right_early = early_trials[
            (block_iden == -1) & (chosen_side[early_trials] == -1)
        ]
        right_left_early = early_trials[
            (block_iden == -1) & (chosen_side[early_trials] == 1)
        ]

        left_left_middle = middle_trials[
            (block_iden == 1) & (chosen_side[middle_trials] == 1)
        ]
        left_right_middle = middle_trials[
            (block_iden == 1) & (chosen_side[middle_trials] == -1)
        ]
        right_right_middle = middle_trials[
            (block_iden == -1) & (chosen_side[middle_trials] == -1)
        ]
        right_left_middle = middle_trials[
            (block_iden == -1) & (chosen_side[middle_trials] == 1)
        ]

        left_left_late = late_trials[(block_iden == 1) & (chosen_side[late_trials] == 1)]
        left_right_late = late_trials[
            (block_iden == 1) & (chosen_side[late_trials] == -1)
        ]
        right_right_late = late_trials[
            (block_iden == -1) & (chosen_side[late_trials] == -1)
        ]
        right_left_late = late_trials[
            (block_iden == -1) & (chosen_side[late_trials] == 1)
        ]

        left_left_early_val = value[left_left_early, :]
        left_right_early_val = value[left_right_early, :]
        right_right_early_val = value[right_right_early, :]
        right_left_early_val = value[right_left_early, :]

        left_left_middle_val = value[left_left_middle, :]
        left_right_middle_val = value[left_right_middle, :]
        right_right_middle_val = value[right_right_middle, :]
        right_left_middle_val = value[right_left_middle, :]

        left_left_late_val = value[left_left_late, :]
        left_right_late_val = value[left_right_late, :]
        right_right_late_val = value[right_right_late, :]
        right_left_late_val = value[right_left_late, :]

        fig = plt.figure(figsize=(8, 8))

        y_lim_min = -0.5
        y_lim_max = np.max(value) - 2

        t = np.arange(-2.1, 2.1, 0.1)  # range(self.num_states+1)

        plt.subplot(3, 2, 1)
        plt.title("Left block")
        val_trace_l = np.mean(left_left_early_val, axis=0)
        val_trace_r = np.mean(left_right_early_val, axis=0)
        val_sem_l = stats.sem(left_left_early_val, axis=0)
        val_sem_r = stats.sem(left_right_early_val, axis=0)
        plt.errorbar(
            t, val_trace_l, val_sem_l, color="blue", ecolor="skyblue", linewidth=1.0
        )
        plt.errorbar(
            t, val_trace_r, val_sem_r, color="green", ecolor="lime", linewidth=1.0
        )
        plt.ylim([y_lim_min, y_lim_max])

        plt.subplot(3, 2, 2)
        plt.title("Right block")
        val_trace_l = np.mean(right_left_early_val, axis=0)
        val_trace_r = np.mean(right_right_early_val, axis=0)
        val_sem_l = stats.sem(right_left_early_val, axis=0)
        val_sem_r = stats.sem(right_right_early_val, axis=0)
        plt.errorbar(
            t, val_trace_l, val_sem_l, color="blue", ecolor="skyblue", linewidth=1.0
        )
        plt.errorbar(
            t, val_trace_r, val_sem_r, color="green", ecolor="lime", linewidth=1.0
        )
        plt.ylim([y_lim_min, y_lim_max])
        plt.text(4, y_lim_max / 2, "1st trial", fontsize=15, color="k")

        plt.subplot(3, 2, 3)
        val_trace_l = np.mean(left_left_middle_val, axis=0)
        val_trace_r = np.mean(left_right_middle_val, axis=0)
        val_sem_l = stats.sem(left_left_middle_val, axis=0)
        val_sem_r = stats.sem(left_right_middle_val, axis=0)
        plt.errorbar(
            t, val_trace_l, val_sem_l, color="blue", ecolor="skyblue", linewidth=1.0
        )
        plt.errorbar(
            t, val_trace_r, val_sem_r, color="green", ecolor="lime", linewidth=1.0
        )
        plt.ylim([y_lim_min, y_lim_max])
        plt.ylabel("Value", fontsize=15)

        plt.subplot(3, 2, 4)
        val_trace_l = np.mean(right_left_middle_val, axis=0)
        val_trace_r = np.mean(right_right_middle_val, axis=0)
        val_sem_l = stats.sem(right_left_middle_val, axis=0)
        val_sem_r = stats.sem(right_right_middle_val, axis=0)
        plt.errorbar(
            t, val_trace_l, val_sem_l, color="blue", ecolor="skyblue", linewidth=1.0
        )
        plt.errorbar(
            t, val_trace_r, val_sem_r, color="green", ecolor="lime", linewidth=1.0
        )
        plt.ylim([y_lim_min, y_lim_max])
        plt.text(4, y_lim_max / 2, "5th trial", fontsize=15, color="k")

        plt.text(6, y_lim_max / 3, "Left Press", fontsize=15, color="blue")
        plt.text(6, y_lim_max / 9, "Right Press", fontsize=15, color="green")

        plt.subplot(3, 2, 5)
        val_trace_l = np.mean(left_left_late_val, axis=0)
        val_trace_r = np.mean(left_right_late_val, axis=0)
        val_sem_l = stats.sem(left_left_late_val, axis=0)
        val_sem_r = stats.sem(left_right_late_val, axis=0)
        plt.errorbar(
            t, val_trace_l, val_sem_l, color="blue", ecolor="skyblue", linewidth=1.0
        )
        plt.errorbar(
            t, val_trace_r, val_sem_r, color="green", ecolor="lime", linewidth=1.0
        )
        plt.ylim([y_lim_min, y_lim_max])

        plt.subplot(3, 2, 6)
        val_trace_l = np.mean(right_left_late_val, axis=0)
        val_trace_r = np.mean(right_right_late_val, axis=0)
        val_sem_l = stats.sem(right_left_late_val, axis=0)
        val_sem_r = stats.sem(right_right_late_val, axis=0)
        plt.errorbar(
            t, val_trace_l, val_sem_l, color="blue", ecolor="skyblue", linewidth=1.0
        )
        plt.errorbar(
            t, val_trace_r, val_sem_r, color="green", ecolor="lime", linewidth=1.0
        )
        plt.ylim([y_lim_min, y_lim_max])
        plt.text(4, y_lim_max / 2, "15th trial", fontsize=15, color="k")

        plt.text(-3, -2.5, "Time(s)", fontsize=15)

        if save:
            plt.savefig(save, bbox_inches="tight")

    def dopamine_regression(self, plt_sec=None, trials_back=6, save=None):

        rewarded_trials = self.is_rewardeds[self.current_trial_times == 2]
        chosen_side = 2 * self.choices[self.current_trial_times == 2] - 1
        if plt_sec:
            y_vec_cs = np.mean(
                np.array(
                    [self.values[self.current_trial_times == i] for i in range(2, 7)]
                ),
                axis=0,
            )
        else:
            y_vec_cs = np.mean(
                np.array(
                    [
                        self.rpes[self.current_trial_times == i]
                        for i in range(self.num_states // 2, self.num_states // 2 + 10)
                    ]
                ),
                axis=0,
            )

        # makes x matrix of reward identity on previous trials
        x_mat = np.zeros((trials_back, len(rewarded_trials)))
        for i in np.arange(0, trials_back):
            x_mat[i, :] = np.roll(rewarded_trials, i)

        y = np.reshape(y_vec_cs, [len(y_vec_cs), 1])
        x = np.concatenate((np.ones([1, len(y_vec_cs)]), x_mat), axis=0)

        results = sm.OLS(y, x.T).fit()

        plt.figure(figsize=(11, 4))
        if plt_sec:
            plt.suptitle(
                "Regressing DA activity at beginning against outcome", fontsize=20
            )
        else:
            plt.suptitle(
                "Regressing DA activity at reward time against outcome", fontsize=20
            )

        plt.scatter(np.arange(trials_back), results.params[1:None])
        plt.axhline(y=0, linestyle="dashed", color="k")
        plt.xlabel("trials back", fontsize=15)
        plt.ylabel("regression coefficients", fontsize=15)
        print(results.params[1:None])

        if save:
            plt.savefig(save, bbox_inches="tight")

    def block_switch(self, trial_back=10, save=None):

        high_prob = self.reward_probs[
            (self.current_trial_times == 2) & (self.choices != -1)
        ]
        chosen_side = (
            2 * self.choices[(self.current_trial_times == 2) & (self.choices != -1)] - 1
        )

        switch_high = np.where(np.diff(high_prob) != 0)[0] + 1
        early_trials = switch_high[0 : len(switch_high) - 1]
        block_iden = high_prob[early_trials]

        # finds times of left to right
        block_switch = early_trials[block_iden == -1] - 1

        time_window = np.arange(-trial_back, trial_back + 1)
        r_choice_mat = np.zeros([len(block_switch), len(time_window)])
        for i in np.arange(1, len(block_switch)):
            r_choice_mat[i, :] = chosen_side[time_window + block_switch[i]]
            r_choice_mat[i, :] = (r_choice_mat[i, :] + 1) / 2

        # same except right to left
        block_switch = early_trials[block_iden == 1] - 1

        time_window = np.arange(-trial_back, trial_back + 1)
        l_choice_mat = np.zeros([len(block_switch), len(time_window)])
        for i in np.arange(1, len(block_switch)):
            l_choice_mat[i, :] = chosen_side[time_window + block_switch[i]] * -1
            l_choice_mat[i, :] = (l_choice_mat[i, :] + 1) / 2

        final_choice_mat = np.concatenate([l_choice_mat, r_choice_mat], axis=0)

        plot_trace = np.mean(final_choice_mat, axis=0)
        sem_trace = stats.sem(final_choice_mat, axis=0)

        plot_trace2 = np.mean(-1 * (final_choice_mat - 1), axis=0)
        sem_trace2 = stats.sem(-1 * (final_choice_mat - 1), axis=0)

        ax = plt.figure(figsize=(8, 4)).gca()
        ax.axvline(x=0, linestyle="dotted", color="gray")
        ax.errorbar(time_window, plot_trace, sem_trace)
        ax.errorbar(time_window, plot_trace2, sem_trace2)
        ax.set_ylim(0, 1)
        ax.xaxis.set_major_locator(MaxNLocator(integer=True, min_n_ticks=10))
        ax.text(
            -10,
            1.05,
            "Blue - pre-switch \nhigh probability\nchoice  ",
            fontsize=15,
            color="dodgerblue",
        )
        ax.text(
            1,
            1.05,
            "Orange - post-switch \nhigh probability\nchoice  ",
            fontsize=15,
            color="orange",
        )
        plt.xlabel("Trials from block switch")
        plt.ylabel("p(choice)")
        # plt.axhline(y=plot_trace[time_window==1],color='k',linestyle='dotted')
        print(plot_trace[time_window == 1])

        if save:
            plt.savefig(save, bbox_inches="tight")

    def rpe_plot(self, save=None):

        rpe_r = np.array(
            [
                np.mean(
                    self.rpes[
                        (self.current_trial_times == (i + 2) % self.num_states)
                        & (self.choices != -1)
                        & (self.is_rewardeds == 1)
                    ]
                )
                for i in range(self.num_states)
            ]
        )
        rpe_ur = np.array(
            [
                np.mean(
                    self.rpes[
                        (self.current_trial_times == (i + 2) % self.num_states)
                        & (self.choices != -1)
                        & (self.is_rewardeds == 0)
                    ]
                )
                for i in range(self.num_states)
            ]
        )
        val_r = np.mean(
            self.values[
                (self.current_trial_times == 1)
                & (self.choices != -1)
                & (self.is_rewardeds == 1)
            ]
        )
        val_ur = np.mean(
            self.values[
                (self.current_trial_times == 1)
                & (self.choices != -1)
                & (self.is_rewardeds == 0)
            ]
        )

        rew_rpe = np.concatenate([[0], [val_r], rpe_r])
        unrew_rpe = np.concatenate([[0], [val_ur], rpe_ur])

        plt.figure(figsize=(16, 6))
        plt.subplot(1, 2, 1)
        plt.title("rewarded trials", fontsize=20)
        plt.plot(np.arange(-2.2, 2.1, 0.1), rew_rpe, color="green")
        plt.ylim(-2.5, 2.5)
        plt.xlabel("time", fontsize=20)
        plt.ylabel("reward prediction error", fontsize=20)
        plt.subplot(1, 2, 2)
        plt.title("unrewarded trials", fontsize=20)
        plt.plot(np.arange(-2.2, 2.1, 0.1), unrew_rpe, color="grey")
        plt.ylim(-2.5, 2.5)
        plt.xlabel("time", fontsize=20)

        if save:
            plt.savefig(save, bbox_inches="tight")

    def logit_rpe(self, save=None):

        mean_rpe_r = np.zeros(30)
        mean_rpe_l = np.zeros(30)
        mean_logit_r = np.zeros(30)
        mean_logit_l = np.zeros(30)

        prob_vec = self.actor_logits[:, 1] - self.actor_logits[:, 2]

        for i in range(1, 31):

            if i == 30:
                mean_logit_r[i - 1] = np.mean(
                    prob_vec[
                        (self.current_trial_times == 1)
                        & (self.reward_probs == -1)
                        & (self.trial_count_in_blocks == 1)
                    ]
                )
                mean_logit_l[i - 1] = np.mean(
                    prob_vec[
                        (self.current_trial_times == 1)
                        & (self.reward_probs == 1)
                        & (self.trial_count_in_blocks == 1)
                    ]
                )
            else:
                mean_logit_r[i - 1] = np.mean(
                    prob_vec[
                        (self.current_trial_times == 1)
                        & (self.reward_probs == 1)
                        & (self.trial_count_in_blocks == i + 1)
                    ]
                )
                mean_logit_l[i - 1] = np.mean(
                    prob_vec[
                        (self.current_trial_times == 1)
                        & (self.reward_probs == -1)
                        & (self.trial_count_in_blocks == i + 1)
                    ]
                )

            mean_rpe_r[i - 1] = np.mean(
                self.rpes[
                    (self.current_trial_times > self.num_states // 2)
                    & (self.reward_probs == 1)
                    & (self.trial_count_in_blocks == i)
                ]
            )
            mean_rpe_l[i - 1] = np.mean(
                self.rpes[
                    (self.current_trial_times > self.num_states // 2)
                    & (self.reward_probs == -1)
                    & (self.trial_count_in_blocks == i)
                ]
            )

        color_range = np.arange(1, 31)
        cbar_name = "Trials in a left to right block"

        plt.figure(figsize=(10, 10))
        plt.plot(
            mean_rpe_r,
            mean_logit_r,
            color="grey",
        )

        fig1 = plt.scatter(
            mean_rpe_r,
            mean_logit_r,
            c=color_range,
            cmap="cool",
            edgecolors="black",
            s=100,
        )
        plt.plot(
            mean_rpe_l,
            mean_logit_l,
            color="grey",
        )

        fig2 = plt.scatter(
            mean_rpe_l,
            mean_logit_l,
            c=color_range,
            cmap="hot",
            edgecolors="black",
            s=100,
        )
        cbar = plt.colorbar(fig1)
        cbar.set_label("Trials in a left block", fontsize=15)
        cbar2 = plt.colorbar(fig2)
        cbar2.set_label("Trials in a right block", fontsize=15)
        plt.ylabel("Decision Variable", fontsize=20)
        plt.xlabel("RPE", fontsize=20)
        plt.title("Decision variable vs RPE plot", fontsize=15)

        if save:
            plt.savefig(save, bbox_inches="tight")

    def value_plot(self, save=None):

        val_r = np.array(
            [
                np.mean(
                    self.values[
                        (self.current_trial_times == (i + 1) % self.num_states)
                        & (self.choices == 1)
                    ]
                )
                for i in range(self.num_states)
            ]
        )
        val_l = np.array(
            [
                np.mean(
                    self.values[
                        (self.current_trial_times == (i + 1) % self.num_states)
                        & (self.choices == 0)
                    ]
                )
                for i in range(self.num_states)
            ]
        )

        plt.figure(figsize=(12, 4))
        plt.subplot(1, 2, 1)
        plt.plot(np.arange(-2.1, 2.1, 0.1), val_l)
        plt.xlabel("time")
        plt.ylabel("value left")

        plt.subplot(1, 2, 2)
        plt.plot(np.arange(-2.1, 2.1, 0.1), val_r)
        plt.xlabel("time")
        plt.ylabel("value right")

        if save:
            plt.savefig(save, bbox_inches="tight")

    def lstm_dynamic_across_blocks_3d(
        self,
        lstm_type="actor",
        state_type="hidden",
        plot_elevation_angle=15,
        plot_azimuth_angle=70,
        save=None,
    ):

        if lstm_type == "actor":
            if state_type == "hidden":
                lstm_state = self.actor_hs
            elif state_type == "cell":
                lstm_state = self.actor_cs
        else:
            if state_type == "hidden":
                lstm_state = self.critic_hs
            elif state_type == "cell":
                lstm_state = self.critic_cs

        num_pca_components = 3

        pca = PCA(n_components=num_pca_components)

        pca.fit(lstm_state[self.current_trial_times == 2, :])

        print(
            f"Fraction of variance explained by the PCA components are "
            f"{np.round(pca.explained_variance_ratio_, 3)}"
        )

        trial_count_lowest = np.min(self.trial_count_in_blocks)
        trial_count_highest = np.max(self.trial_count_in_blocks)

        left_pca = np.zeros(
            (trial_count_highest - trial_count_lowest + 1, num_pca_components)
        )
        right_pca = np.zeros(
            (trial_count_highest - trial_count_lowest + 1, num_pca_components)
        )

        for i in range(trial_count_lowest, trial_count_highest + trial_count_lowest):
            left_pca[i - trial_count_lowest, :] = np.mean(
                pca.transform(
                    lstm_state[
                        (self.current_trial_times == 2)
                        & (self.reward_probs == -1)
                        & (self.trial_count_in_blocks == i)
                    ]
                ),
                axis=0,
            )

            right_pca[i - trial_count_lowest, :] = np.mean(
                pca.transform(
                    lstm_state[
                        (self.current_trial_times == 2)
                        & (self.reward_probs == 1)
                        & (self.trial_count_in_blocks == i)
                    ]
                ),
                axis=0,
            )

        color_range = np.arange(
            trial_count_lowest, trial_count_lowest + trial_count_highest
        )

        plt.figure(figsize=(30, 30))
        ax = plt.axes(projection="3d")
        cb1 = ax.scatter3D(
            left_pca[:, 0],
            left_pca[:, 1],
            left_pca[:, 2],
            c=color_range,
            cmap="cool",
            edgecolors="black",
            alpha=1.0,
            s=500,
        )
        ax.plot3D(left_pca[:, 0], left_pca[:, 1], left_pca[:, 2], color="black")
        cbar1 = plt.colorbar(cb1)
        cbar1.set_label("Trials in a left block", fontsize=20)

        cb2 = ax.scatter3D(
            right_pca[:, 0],
            right_pca[:, 1],
            right_pca[:, 2],
            c=color_range,
            cmap="hot",
            edgecolors="black",
            alpha=1.0,
            s=500,
        )
        ax.plot3D(
            right_pca[:, 0],
            right_pca[:, 1],
            right_pca[:, 2],
            color="black",
        )
        cbar2 = plt.colorbar(cb2)
        cbar2.set_label(f"Trials in a right block", fontsize=20)
        ax.set_xlabel("PC 1")
        ax.set_ylabel("PC 2")
        ax.set_zlabel("PC 3")
        ax.xaxis.labelpad = 30
        ax.yaxis.labelpad = 30
        ax.zaxis.labelpad = 30
        ax.xaxis.set_tick_params(pad=10)
        ax.yaxis.set_tick_params(pad=10)
        ax.zaxis.set_tick_params(pad=10)

        ax.view_init(plot_elevation_angle, plot_azimuth_angle)
        plt.draw()

        if save:
            plt.savefig(save, bbox_inches="tight")

    def regress_lstm_pca_components(
        self,
        lstm_type="actor",
        state_type="hidden",
        num_pca_components=3,
        num_trials_back=7,
        num_pca_components_to_plot=1,
        save=None,
    ):

        if lstm_type == "actor":
            if state_type == "hidden":
                lstm_state = self.actor_hs[self.current_trial_times == 2, :]
            elif state_type == "cell":
                lstm_state = self.actor_cs[self.current_trial_times == 2, :]
        else:
            if state_type == "hidden":
                lstm_state = self.critic_hs[self.current_trial_times == 2, :]
            elif state_type == "cell":
                lstm_state = self.critic_cs[self.current_trial_times == 2, :]

        pca = PCA(n_components=num_pca_components)

        transformed_lstm_state = pca.fit_transform(lstm_state)

        print(
            f"Fraction of variance explained by the PCA components are "
            f"{np.round(pca.explained_variance_ratio_, 3)}"
        )

        dependent_variables = pd.DataFrame(
            transformed_lstm_state,
            columns=[f"PC {i}" for i in range(1, num_pca_components + 1)],
        )

        standard_scaler = preprocessing.StandardScaler()

        # fit and transform the data
        dependent_variables = pd.DataFrame(
            standard_scaler.fit_transform(dependent_variables),
            columns=dependent_variables.columns,
        )

        independent_variables = pd.DataFrame()
        for i in range(1, num_trials_back + 1):
            independent_variables[f"{i}\ntrials\nback"] = (
                2 * np.roll(self.choices[(self.current_trial_times == 3)], i) - 1
            )
        for i in range(1, num_trials_back + 1):
            independent_variables[f"{i}\ntrials\nback\nrpe"] = np.roll(
                np.mean(
                    np.array(
                        [
                            self.rpes[self.current_trial_times == j]
                            for j in range(
                                self.num_states // 2, self.num_states // 2 + 15
                            )
                        ]
                    ),
                    axis=0,
                ),
                i,
            )
        for i in range(1, num_trials_back + 1):
            independent_variables[f"{i}\ntrials\nback\nchoice\nrpe"] = np.roll(
                np.mean(
                    np.array(
                        [
                            (2 * self.choices[self.current_trial_times == j] - 1)
                            * self.rpes[self.current_trial_times == j]
                            for j in range(
                                self.num_states // 2, self.num_states // 2 + 15
                            )
                        ]
                    ),
                    axis=0,
                ),
                i,
            )

        columns_to_scale = []
        for i in range(1, num_trials_back + 1):
            columns_to_scale.append(f"{i}\ntrials\nback\nrpe")
            columns_to_scale.append(f"{i}\ntrials\nback\nchoice\nrpe")
        standard_scaler = preprocessing.StandardScaler()
        independent_variables[columns_to_scale] = standard_scaler.fit_transform(
            independent_variables[columns_to_scale]
        )

        independent_variables = sm.add_constant(independent_variables)

        plt.figure(figsize=(12, 8 * num_pca_components_to_plot))

        for component_num in range(1, num_pca_components_to_plot + 1):
            pc_name = f"PC {component_num}"

            regression_results = sm.OLS(
                dependent_variables[pc_name], independent_variables
            ).fit()

            plt.subplot(num_pca_components_to_plot, 1, component_num)

            plt.title(
                f"{pc_name} (R-squared = "
                f""
                f"{np.round(regression_results.rsquared_adj, 3)})"
            )

            plt.plot(
                range(1, num_trials_back + 1),
                regression_results.params.values[
                    0 * num_trials_back + 1 : 1 * num_trials_back + 1
                ],
                label="Choice",
                color="green",
            )
            plt.plot(
                range(1, num_trials_back + 1),
                regression_results.params.values[
                    1 * num_trials_back + 1 : 2 * num_trials_back + 1
                ],
                label="RPE",
                color="red",
            )
            plt.plot(
                range(1, num_trials_back + 1),
                regression_results.params.values[
                    2 * num_trials_back + 1 : 3 * num_trials_back + 1
                ],
                label="Choice x RPE",
                color="blue",
            )

            plt.scatter(
                y=regression_results.params.values[0],
                x=1,
                color="black",
                label="constant",
            )

            plt.axhline(y=0, color="grey", linestyle="dotted")

            plt.xticks(range(1, num_trials_back + 1))
            plt.xlim([1 - 0.2, num_trials_back])

            plt.ylabel("regression coefficients")
            plt.xlabel("Trials back")

            plt.legend()

        if save:
            plt.savefig(save, bbox_inches="tight")

    def actor_lstm_PC_overlap_with_output(self, num_pca_components=3, save=None):

        from environment import TestingEnvironment
        from model_arch import ActorCriticLSTM
        from MetaRLModel import MetaRLModel

        MetaRLModel = MetaRLModel(
            ActorCriticLSTM(),
            TestingEnvironment(reward_width=0.3),
            np.load("recorded_data/sequential_activity.npz"),
        )

        actor_weights = MetaRLModel.load_weights(load_model="saved_model/model-62000")[
            "ActorOutput"
        ]

        lstm_state = self.actor_hs[self.current_trial_times == 2, :]

        pca = PCA(n_components=num_pca_components)
        pca.fit(lstm_state)

        print(
            f"Fraction of variance explained by the PCA components are "
            f"{np.round(pca.explained_variance_ratio_, 3)}"
        )

        weights, _ = actor_weights
        weight_vector = (weights[:, 2] - weights[:, 1]).reshape(-1)

        cosines = np.zeros(num_pca_components)
        for pca_component in range(num_pca_components):
            PC_vector = pca.components_[pca_component, :].reshape(-1)
            cosines[pca_component] = np.dot(PC_vector, weight_vector) / (
                norm(PC_vector) * norm(weight_vector)
            )

        plt.figure(figsize=(8, 8))
        plt.title("overlap of output activity PCs\n with actor readout weights")
        sns.barplot(y=cosines, x=["PC 1", "PC 2", "PC 3"], color="Blue")
        plt.ylabel("cosine similarity")
        plt.ylim([-0.2, 1.0])

        if save:
            plt.savefig(save, bbox_inches="tight")

def plot_model_performances(
    input_types=[
        "sequential",
        "persistent",
    ],
    labels=[
        "choice-selective \nPL-NAc activity",
        "choice-selective \npersistent activity",
    ],
    episodes=np.arange(10000, 72000, 2000),
    learning_rate=0.001,
    num_hidden_units=128,
    save=None,
):
    """Generate the average reward received plot for model at different training stages.

    Args:
        input_types (list, optional): List of PL-NAc input types.
        Defaults to [ "sequential", "persistent", ].
        labels (list, optional): Label corresponding to the PL-NAc input types.
        Defaults to [ "choice-selective \nPL-NAc activity", "choice-selective \npersistent activity", ].
        episodes ([type], optional): List of training stages. Defaults to
        np.arange(10000, 72000, 2000).
        learning_rate (float, optional): Learning rate used. Defaults to 0.001.
        num_hidden_units (int, optional): Number of hidden units of the LSTMS. Defaults
        to 128.
        save ([type], optional): Location to save the plot. Defaults to None.
    """

    performance_results = {input_type: [] for input_type in input_types}

    fig = plt.figure(figsize=(10, 6))

    for input_type in input_types:
        test_destination = (
            f"training_data/{input_type}_input/learning_rate="
            f"{learning_rate}_hidden_units={num_hidden_units}/"
            f"model_performances/"
        )
        for episode in episodes:
            testing_file = test_destination + f"/model-{episode}/" + "test_0.npz"
            testing_data = np.load(testing_file)
            is_rewardeds = testing_data["is_rewardeds"]
            current_trial_times = testing_data["current_trial_times"]
            rewarded_trials = is_rewardeds[current_trial_times == 2]
            performance_results[input_type].append(np.mean(rewarded_trials))
        plt.plot(episodes, performance_results[input_type], label=input_type)

    plt.xlabel("episodes")
    plt.ylabel("reward rate")
    plt.axhline(y=0.4, linestyle="dotted", color="k")
    plt.ylim(0, 0.6)
    plt.legend(labels)

    if save:
        plt.savefig(save, bbox_inches="tight")


import uuid

from django.conf import settings
from django.db import models

from .utils import format_temperature
from .validators import validate_human_body_temperature

class TemperatureRecord(models.Model):
    id = models.UUIDField(
        editable=False, default=uuid.uuid4, primary_key=True, verbose_name="ID"
    )
    person = models.ForeignKey("people.Person", on_delete=models.PROTECT)
    body_temperature = models.DecimalField(
        max_digits=4,
        decimal_places=2,
        validators=[validate_human_body_temperature],
        help_text="The person's body temperature in degrees celsius.",
    )
    created_by = models.ForeignKey(
        to=settings.AUTH_USER_MODEL,
        on_delete=models.SET_NULL,
        null=True,
        help_text="The user who created this record.",
    )
    created_at = models.DateTimeField(auto_now_add=True)
    last_modified = models.DateTimeField(auto_now=True)

    class Meta:  # noqa
        db_table = "records_temperature"
        ordering = ["person__username", "created_at"]

    def __str__(self):
        temp = format_temperature(self.body_temperature)
        return f"{self.person} was {temp} at {self.created_at}"


#Build as of 6.5.23
import nltk
import openai
import logging
from nltk.sentiment.vader import SentimentIntensityAnalyzer
from typing import Any, Dict, List, Union

class ChatCompletionResponse:
    choices: List[Dict[str, Any]]

    def __init__(self, **kwargs):
        self.choices = kwargs.get('choices', [])

class SentimentAnalysis:
    def __init__(self, openai_api_key: str):
        nltk.download('vader_lexicon')
        self.vader = SentimentIntensityAnalyzer()
        openai.api_key = openai_api_key
        # Set up logging
        logging.basicConfig(filename='sentiment_analysis.log', level=logging.ERROR)

    def analyze_vader(self, text: str) -> float:
        try:
            scores = self.vader.polarity_scores(text)
            return scores['compound']
        except Exception as e:
            logging.error(f"Error analyzing text with Vader: {str(e)}")
            return 0  # Return neutral sentiment score in case of error

    def analyze_openai(self, text: str) -> str:
        try:
            response = openai.ChatCompletion.create(
                model="text-curie-0001",
                messages=[
                    {"role": "system", "content": "This is a sentiment analysis task. Please categorize the following news article's sentiment as 'Good', 'Bad', or 'Unknown'."},
                    {"role": "user", "content": text}
                ]
            )
            if 'choices' in response and len(response.choices) > 0:
                return response.choices[0].message['content']
            else:
                logging.error("OpenAI API returned no choices or the list itself was empty.")
                return 'Unknown'
        except Exception as e:
            logging.error(f"Error analyzing text with OpenAI: {str(e)}")
            return 'Unknown'  # Return 'Unknown' sentiment in case of error


from posixpath import split

class Carro:
    def __init__(self, matricula):
        self.matricula = matricula
        self.emUso = False

    def setEmUso(self):
        self.emUso = True

    def getIsEmUso(self):
        return self.emUso

    def getMatricula(self):
        return self.matricula

class ParqueEstacionamento:
    def __init__(self):
        self.carrosDentro = []
        self.limite = 5
        self.carrosTotal = {}

    def getLugaresEmUso(self):
        return len(self.carrosDentro)

    def receberCarro(self, matricula):
        if matricula not in self.carrosTotal.keys():
            carro = Carro(matricula)
        else:
            carro = self.carrosTotal[matricula]

        if self.limite > len(self.carrosDentro) and matricula not in self.carrosDentro:
            self.carrosDentro.append(matricula)
            carro.emUso = True
        self.carrosTotal[matricula] = carro

    def retirarCarro(self,matricula):
        if matricula in self.carrosDentro:
            self.carrosDentro.remove(matricula)
            self.carrosTotal[matricula].emUso = False

if __name__ == '__main__':
    parque = ParqueEstacionamento()
    lines = int(input(""))
    for _ in range(lines):
        line = input("").split(" - ")
        if line[0] == "aceitarCarro":
            parque.receberCarro(line[1])
        elif line[0] == "sairCarro":
            parque.retirarCarro(line[1])
    print(parque.getLugaresEmUso())
    for carro in sorted(parque.carrosDentro):
        print(carro)


import datetime
import logging.config
import argparse

from scrapers import *
from util.latest_round import CALENDAR_DICT
from util.latest_round import find_latest_round
from logging_config import LOGGING_CONFIG

logging.config.dictConfig(LOGGING_CONFIG)
LOGGER = logging.getLogger(__name__)

def main(from_year, to_year, from_round, to_round):
    LOGGER.info("Starting Full DataScrape")
    playerScrape.main()
    injuryScrape.main()
    gameScrape.scrape_games(from_year, to_year)
    fantasyScrape.main(from_year, to_year, from_round, to_round)
    match_statsScrape.main(from_year, to_year, from_round, to_round)
    LOGGER.info("Finished Full DataScrape")

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="Scrape AFL data")
    parser.add_argument('--from_year', help='scrape data from this year onwards',
                        type=int, choices=range(1990, datetime.datetime.now().year + 1), default=2021)
    parser.add_argument('--to_year', help='scrape data up until this year (inclusive)',
                        type=int, choices=range(1990, datetime.datetime.now().year + 1), default=datetime.datetime.now().year)
    parser.add_argument('--from_round', help='scrape data from this round onwards',
                        type=int, choices=range(1, 30), default=1)
    parser.add_argument('--to_round', help='scrape data up until this round (inclusive)',
                        type=int, choices=range(1, 30), default=30)
    parser.add_argument('-l', '--latest', help='scrape the latest data only', action="store_true")
    args = parser.parse_args()
    if args.latest:
        LOGGER.info("\"Latest\" flag provided: Ignoring any from/to arguments and scraping for the latest rounds")
        args.from_year = datetime.datetime.now().year
        args.to_year = datetime.datetime.now().year
        args.from_round = find_latest_round(CALENDAR_DICT, datetime.datetime.now())
        args.to_round = args.from_round

    main(from_year=args.from_year, to_year=args.to_year, from_round=args.from_round, to_round=args.to_round)