import datetime

from sqlalchemy import (
    ForeignKey,
    Table,
    Column,
    Integer,
    DateTime,
    Enum,
    Text,
    create_engine,
)
from sqlalchemy.orm import relationship, registry

from app.domain.model import User, App, AppAuthorization

mapper_registry = registry()

applications = Table(
    "authorizations",
    mapper_registry.metadata,
    Column("id", Integer, primary_key=True),
    Column("name", Enum(App)),
    Column("user_id", ForeignKey("users.id")),
    Column(
        "updated_at",
        DateTime,
        default=datetime.datetime.now,
        onupdate=datetime.datetime.now,
    ),
    Column("created_at", DateTime, default=datetime.datetime.now),
)

users = Table(
    "users",
    mapper_registry.metadata,
    Column("id", Text, primary_key=True),
    Column("password", Text),
    Column("first_name", Text),
    Column("last_name", Text),
    Column("email", Text, unique=True),
    Column(
        "updated_at",
        DateTime,
        default=datetime.datetime.now,
        onupdate=datetime.datetime.now,
    ),
    Column("created_at", DateTime, default=datetime.datetime.now),
)

def start_mappers():
    operation_mapper = mapper_registry.map_imperatively(AppAuthorization, applications)
    account_mapper = mapper_registry.map_imperatively(
        User,
        users,
        properties={
            "_applications_auth": relationship(operation_mapper, collection_class=set)
        },
    )

def set_up_db(uri: str) -> None:
    engine = create_engine(uri)
    mapper_registry.metadata.create_all(engine)


from typing import List
from unittest.mock import Mock

import pytest
from openai_pygenerator import Completion

from llm_cooperation import DEFAULT_MODEL_SETUP, Payoffs
from llm_cooperation.experiments.dilemma import (
    Cooperate,
    Defect,
    DilemmaChoice,
    DilemmaEnum,
    P,
    R,
    S,
    T,
    extract_choice_pd,
    get_prompt_pd,
    payoffs_pd,
    strategy_cooperate,
    strategy_defect,
    strategy_t4t_cooperate,
    strategy_t4t_defect,
)
from llm_cooperation.gametypes import simultaneous
from llm_cooperation.gametypes.repeated import GameSetup, play_game
from tests.common import make_completion

@pytest.mark.parametrize("condition", [True, False])
def test_get_instruction_prompt(condition: bool):
    role_prompt = "You are a helpful assistant."
    prompt = get_prompt_pd(condition, role_prompt)
    assert ("Explanation:" in prompt) == condition
    assert "Choice:" in prompt
    for payoff in [R, S, T, P]:
        assert f"${payoff}.00" in prompt
    assert role_prompt in prompt

@pytest.mark.parametrize(
    "text, expected_move",
    [
        ("project green", Cooperate),
        ("project blue", Defect),
        ("Project GREEN", Cooperate),
        ("Project BLUE", Defect),
        ("'project green'", Cooperate),
        ("Choice: 'project green'", Cooperate),
        (
            """Explanation: Because project green is bad for me.
Choice: Project BLUE""",
            Defect,
        ),
    ],
)
def test_extract_choice_pd(text: str, expected_move: DilemmaChoice):
    move = extract_choice_pd(make_completion(text))
    assert move == expected_move

@pytest.mark.parametrize(
    "user_choice, partner_choice, expected_payoffs",
    [
        (Defect, Cooperate, (T, S)),
        (Cooperate, Cooperate, (R, R)),
        (Defect, Defect, (P, P)),
        (Cooperate, Defect, (S, T)),
    ],
)
def test_payoffs(
    user_choice: DilemmaChoice,
    partner_choice: DilemmaChoice,
    expected_payoffs: Payoffs,
):
    user_payoff, partner_payoff = payoffs_pd(user_choice, partner_choice)
    assert (user_payoff, partner_payoff) == expected_payoffs

@pytest.mark.parametrize(
    "strategy, index, expected",
    [
        (strategy_cooperate, 5, Cooperate),
        (strategy_cooperate, 3, Cooperate),
        (strategy_defect, 5, Defect),
        (strategy_defect, 3, Defect),
        (strategy_t4t_cooperate, 5, Defect),
        (strategy_t4t_cooperate, 3, Cooperate),
        (strategy_t4t_cooperate, 2, Cooperate),
        (strategy_t4t_defect, 2, Defect),
    ],
)
def test_strategy(strategy, index, expected, conversation):
    state = Mock(spec=["messages"])
    state.messages = conversation[:index]
    assert strategy(state) == expected

def test_dilemma_choice():
    c1 = DilemmaChoice(DilemmaEnum.C)
    c2 = DilemmaChoice(DilemmaEnum.C)
    d = DilemmaChoice(DilemmaEnum.D)
    assert c1 == c2
    assert c1 != d
    assert Cooperate != Defect

def test_run_repeated_game(mocker):
    completions = [
        {"role": "assistant", "content": "project green"},
    ]
    mocker.patch(
        "openai_pygenerator.openai_pygenerator.generate_completions",
        return_value=completions,
    )
    conversation: List[Completion] = list(
        play_game(
            partner_strategy=strategy_defect,
            role_prompt="You are a participant in a psychology experiment",
            participant_condition=False,
            game_setup=GameSetup(
                num_rounds=3,
                generate_instruction_prompt=get_prompt_pd,
                next_round=simultaneous.next_round,
                rounds=simultaneous.rounds_setup,
                payoffs=payoffs_pd,
                extract_choice=extract_choice_pd,
                model_setup=DEFAULT_MODEL_SETUP,
            ),
        )
    )
    assert len(conversation) == 7
    # pylint: disable=unsubscriptable-object
    assert Defect.description in conversation[-1]["content"]


from ftplib import FTP
from io import StringIO
from datetime import datetime
from abc import ABCMeta, abstractmethod
import numpy as np
import pandas as pd
from bokeh.models.sources import ColumnDataSource

class timeseries_dataset(metaclass=ABCMeta):
    #This is the abstract base class for time series datasets. They need a name, the info whether one is able to obtain
    #the data without a token or password, and whether users need to upload the file on their own
    ##TODO are there any cases I haven't thought about yet?
    ##TODO think about how to use needs_token and needs_local_dataset in the GUI
    def __init__(self, name, bool_token, bool_local):
        self.name = name ##TODO create names that are unique and userfriendly, or create userfriendly names and ids
        self.needs_token = bool_token
        self.needs_local_dataset = bool_local
        self.is_loaded = False
        self.version = None ##TODO think about how to maintain datasets

    @abstractmethod
    def get_citation(self):
        #This function must return a string that users can simply paste into their documents as correct citation
        pass

    @abstractmethod
    def load_data(self): #This function needs to load a ColumnDataSource into self.data and set self.is_loaded to true
        pass

    def __str__(self):
        return self.name

class dwd_dataset(timeseries_dataset):
    def __init__(self, init_arguments):
        self.ftp_path, self.ftp_file = init_arguments ##If more than one argument is passed, arguments need to be separated afterwards
        print(self.ftp_file)
        print(self.ftp_path)
        self.y_columns = ['y']
        self.y_labels = ['Temperature (Celsius)']
        ##TODO Especially for these datasets, make geolocation input-able and use this information to choose the
        # correct dataset

        super().__init__("Tmp_DWD", False, False)

    def load_data(self):
        month_str_to_int = {"Jan": 1, "Feb": 2, "Mrz": 3, "Apr": 4, "Mai": 5, "Jun": 6, "Jul": 7, "Aug": 8, "Sep": 9,
                            "Okt": 10, "Nov": 11, "Dez": 12}

        ftp = FTP("opendata.dwd.de")
        ftp.login()
        ftp.cwd(self.ftp_path)
        r = StringIO()
        def customWriter(line):
            r.write(line)
            r.write("\n")
        ftp.retrlines("RETR " + self.ftp_file, customWriter)

        labels = []
        values = []
        for i, line in enumerate(r.getvalue().split("\n")):
            if i == 0:
                tmp_labels = line.split(";")[1:]
            else:
                if len(line.split(";")) > 1:
                    labels += [np.datetime64(datetime(
                        int(line.split(";")[0]), month_str_to_int[lab.strip()], 1))
                                    for lab in tmp_labels]
                    values += [float(part) for part in line.split(";")[1:]]

        self.data = ColumnDataSource({'x': labels, 'y': values})
        self.is_loaded = True

    def get_citation(self):
        return("Deutscher Wetterdienst")

class global_carbon_project_dataset(timeseries_dataset):
    def __init__(self, gcp_global_file):
        self.gcp_global_file = gcp_global_file
        self.y_columns = ['Fossil-Fuel-And-Industry', 'Land-Use-Change-Emissions', 'Atmospheric-Growth',
                          'Ocean-Sink', 'Land-Sink', 'Budget-Imbalance']
        self.y_labels = ['Fossil-Fuel-And-Industry (Gt)', 'Land-Use-Change-Emissions (Gt)', 'Atmospheric-Growth (Gt)',
                         'Ocean-Sink (Gt)', 'Land-Sink (Gt)', 'Budget-Imbalance (Gt)']
        super().__init__("Global Carbon Budget 2019", False, False)

    def load_data(self):
        csv_data = pd.read_csv(self.gcp_global_file, error_bad_lines=False)
        csv_data.columns = ["x", 'Fossil-Fuel-And-Industry', 'Land-Use-Change-Emissions',
                                  'Atmospheric-Growth', 'Ocean-Sink', 'Land-Sink', 'Budget-Imbalance']
        csv_data["x"] = [np.datetime64(datetime(f, 1, 1)) for f in csv_data["x"]]
        self.data = ColumnDataSource(csv_data)
        self.is_loaded = True

    def get_citation(self):
        return("Global Carbon Project. (2019). Supplemental data of Global Carbon Budget 2019 (Version 1.0) [Data set]. "
               "Global Carbon Project. https://doi.org/10.18160/gcp-2019")

class data_collection(object):
    def __init__(self):
        self.datasets = []

        #Set up variables in order to be able to pipe information between tabs
        self.latpos = None
        self.lonpos = None
        self.perimeter = None

    def add_dataset(self, new_dataset):
        #TODO this might be necessary if a user is adding its own dataset. Make this possible!
        None

    def load(self, savefile):
        for i, line in enumerate(open(savefile)):
            if len(line.strip().split("\t")) == 2:
                self.datasets.append(globals()[line.split("\t")[0]](line.strip().split("\t")[1]))
            else:
                self.datasets.append(globals()[line.split("\t")[0]](line.strip().split("\t")[1:]))


import customtkinter as ctk
from customtkinter import filedialog
import winreg
import zipfile
from os import listdir,walk,makedirs
from os.path import join,relpath,isdir,exists
import img2pdf
import tempfile
TYPE=['.cbz','.pdf']

F_type=['Folder','Folder > Folder','Single']

def get_windows_downloads_folder():
    try:
        with winreg.OpenKey(winreg.HKEY_CURRENT_USER, r"Software\\Microsoft\Windows\\CurrentVersion\\Explorer\Shell Folders") as key:
            downloads_guid = '{374DE290-123F-4565-9164-39C4925E467B}' # every single built it folder has a GUID
            downloads_folder, _ = winreg.QueryValueEx(key, downloads_guid)
            return downloads_folder
    except Exception as e:
        print("Error:", e)
        return None

class App(ctk.CTk):
    global r
    r=0
    class MyScrollableCheckboxFrame(ctk.CTkScrollableFrame):
        def __init__(self, master, title, values):
            super().__init__(master, label_text=title)
            self.grid_columnconfigure(0, weight=5)
            self.grid_columnconfigure(1, weight=1)
            self.grid_columnconfigure(2, weight=1)
            self.grid_columnconfigure(3, weight=1)
            self.values = values
            self.rows = []

            for row,data in enumerate(values):
                combobox = ctk.CTkOptionMenu(self,values=TYPE)
                combobox2 = ctk.CTkOptionMenu(self,values=F_type)
                button=ctk.CTkButton(self,text="X",command= lambda: self.removedata())
                label = ctk.CTkLabel(self, text=data)
                label.grid(row=row, column=0, padx=10, pady=(10, 0), sticky="w")
                combobox.grid(row=row, column=1, padx=10, pady=(10, 0), sticky="w")
                combobox2.grid(row=row, column=2, padx=10, pady=(10, 0), sticky="w")
                button.grid(row=row, column=3, padx=10, pady=(10, 0), sticky="w")
                self.rows.append({'label':label,'combo':combobox,'folder':combobox2})

        def get(self):
            all_data = []
            for row in self.rows:
                all_data.append({'dst':row['label'].cget("text"),"type":row['combo'].get(),'f_type':row['folder'].get()})
            return all_data

        def removedata(self):
            # self.rows.pop(index)
            # TODO no idea
            # self.grid_remove(row=index)
            # self.add(combobox)
            pass

    def __init__(self):
        super().__init__()
        self.title("Image to PDF,CBZ")
        self.geometry("600x600")
        self.grid_columnconfigure(0, weight=1)
        self.grid_columnconfigure(1, weight=1)
        self.grid_rowconfigure(0, weight=1)
        self.grid_rowconfigure(1, weight=5)
        self.grid_rowconfigure(2, weight=1)
        self.grid_rowconfigure(3, weight=1)

        self.button=ctk.CTkButton(self,text="Add Folder", command=self.button_callbck)
        self.button.grid(row=0, column=0, sticky="w",padx=10, pady=(10, 0))
        self.values = []
        self.scrollable_checkbox_frame = self.MyScrollableCheckboxFrame(self, title="Folder list", values=self.values)
        self.scrollable_checkbox_frame.grid(row=1, column=0, padx=10, pady=(10, 0), sticky="nsew",columnspan=2)
        self.outputlfolder=ctk.CTkLabel(self, text=get_windows_downloads_folder())
        self.outputlfolder.grid(row=2, column=0, sticky="ns",padx=10, pady=(10, 0))
        self.button3=ctk.CTkButton(self,text="Output Folder",command=self.set_output)
        self.button3.grid(row=2, column=1, sticky="w",padx=10, pady=(10, 0))
        self.button2=ctk.CTkButton(self,text="Convert", command=self.convert)
        self.button2.grid(row=3, column=0, sticky="w",padx=10, pady=(10, 0))

    def button_callbck(self):
        folder_selected = filedialog.askdirectory()
        if len(folder_selected)!=0:
            # self.scrollable_checkbox_frame.destroy()
            self.values.append(folder_selected)
            self.scrollable_checkbox_frame = self.MyScrollableCheckboxFrame(self, title="Folder list", values=self.values)
            self.scrollable_checkbox_frame.grid(row=1, column=0, padx=10, pady=(10, 0), sticky="nsew",columnspan=2)

    def set_output(self):
        folder_selected = filedialog.askdirectory()
        self.outputlfolder.configure(text=folder_selected)

    def create_cbz(self,source_folder, output_cbz):
            with zipfile.ZipFile(output_cbz, 'w') as cbz_file:
                for root, _, files in walk(source_folder):
                    for file in files:
                        if file.lower().endswith(('.jpg', '.jpeg', '.png', '.gif')):
                           image_path = join(root, file)
                           if not isdir(image_path):
                                cbz_file.write(image_path, relpath(image_path, source_folder))
            print('\033[91m'+'output_cbz: ' + '\033[92m', output_cbz)

    def create_pdf(self,source_folder, output_pdf):
        imgs=[]
        a4inpt = (img2pdf.mm_to_pt(210),img2pdf.mm_to_pt(297))
        layout_fun = img2pdf.get_layout_fun(a4inpt)
        for fname in listdir(source_folder):
            if fname.endswith(('.jpg', '.jpeg', '.png', '.gif')):
                path=join(source_folder,fname)
                if not isdir(path):
                    imgs.append(path)
            # elif fname.endswith(('.cbz')):
            #     self.cbztopdf(source_folder,output_pdf)

        with open(output_pdf,"wb") as f:#take folder name as pdf name
            f.write(img2pdf.convert(imgs,layout_fun=layout_fun))
        print('\033[91m'+'output_pdf: ' + '\033[92m', output_pdf)
        imgs.clear()

    def cbztopdf(self,source_folder, output_pdf):
        temp_dir = tempfile.TemporaryDirectory()
        print('\033[91m'+'temp_dir.name: ' + '\033[92m', temp_dir.name)
        with zipfile.ZipFile(source_folder, 'r') as zip_ref:
            zip_ref.extractall(temp_dir.name)
            l=listdir(temp_dir)
            print('\033[91m'+'l: ' + '\033[92m', l)

    def folderInsideFolder(self,source_folder:str,output:str,c_type:str):
        for i in listdir(source_folder):
            source_folder=join(source_folder,i)
            if isdir(source_folder):
                self.onlyFolder(source_folder,output,i,c_type)
            else:
                print('\033[91m'+'single_file: ' + '\033[92m', source_folder)

    def onlyFolder(self,source_folder:str,output:str,filename:str,c_type:str):
        output=self.makeFile(c_type,filename,output)
        print('\033[91m'+'output: ' + '\033[92m', output)
        if c_type==".cbz":
            self.create_cbz(source_folder,output)
        elif c_type==".pdf":
            self.create_pdf(source_folder,output)

    def SingleFile(self,source_folder:str,output:str,filename:str,c_type:str):
        output=self.makeFile(c_type,filename,output)
        print('\033[91m'+'output: ' + '\033[92m', output)
        if c_type==".cbz":
            self.create_cbz(source_folder,output)
        elif c_type==".pdf":
            self.create_pdf(source_folder,output)

    def makeFile(self,filetype,filename,output):
        full_path=join(output,filename+filetype)
        while(True):
            if not exists(full_path):
                return full_path
            else:
                full_path=join(output,filename+' -copy'+filetype)
                print('\033[91m'+'full_path: ' + '\033[92m', full_path)

    def makeFolder(self,folder_name,output):
        full_path=join(output,folder_name)
        while(True):
            if not exists(full_path):
                makedirs(full_path)
                return full_path
            else:
                full_path+=' -copy'
                print('\033[91m'+'full_path: ' + '\033[92m', full_path)

    def convert(self):
       data=self.scrollable_checkbox_frame.get()
       print('\033[91m'+'data: ' + '\033[92m', data)
       output=self.outputlfolder.cget("text")
       for i in data:
            dst=i['dst']
            c_type=i['type']
            f_type=i['f_type']
            title=dst.split('/')[-1].split('.')[0]
            if f_type=='Folder > Folder':
                output=self.makeFolder(title,output)
                self.folderInsideFolder(dst,output,c_type)
            elif f_type=='Folder':
                self.onlyFolder(dst,output,title,c_type)

if __name__=='__main__':
    app = App()
    app.mainloop()


from django import forms

from crispy_forms.helper import FormHelper
from crispy_forms.layout import Layout, Div, Submit, HTML, Button, Row, Field
from crispy_forms.bootstrap import AppendedText, PrependedText, FormActions

class PollForm(forms.Form):
    question = forms.CharField(
        label='Question',
        max_length=200
    )

    choice1 = forms.CharField(
        label='Choice',
        max_length=100
    )

    choice2 = forms.CharField(
        label='Choice',
        max_length=100
    )

    choice3 = forms.CharField(
        label='Choice',
        max_length=100,
        required=False
    )

    choice4 = forms.CharField(
        label='Choice',
        max_length=100,
        required=False
    )

    helper = FormHelper()
    helper.form_class = 'form-horizontal'
    helper.layout = Layout(
        Field('question', css_class='input-lg form-control', placeholder='Enter Question'),
        Field('choice1', css_class='form-control', placeholder='Enter Choice'),
        Field('choice2', css_class='form-control', placeholder='Enter Choice'),
        Field('choice3', css_class='form-control', placeholder='Enter Choice'),
        Field('choice4', css_class='form-control', placeholder='Enter Choice'),
        FormActions(
            Submit('save_changes', 'Submit', css_class="btn-primary btn-lg", style="margin-top: 30px"),
        )
   )

class SearchForm(forms.Form):
    query = forms.CharField(max_length=100)


import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import tensorflow as tf
try:
    import tensorflow.keras as keras
    import tensorflow.keras.backend as K
except:
    import keras
    import keras.backend as K
import importlib
import gc

class History(object):
    """
    Custom class to help get log data from keras.callbacks.History objects.

    :param history: a ``keras.callbacks.History object`` or ``None``.
    """
    def __init__(self, history=None):
        if history is not None:
            self.epoch = history.epoch
            self.history = history.history
        else:
            self.epoch = []
            self.history = {}

def concatenate_history(hlist, reindex_epoch=False):
    """
    A helper function to concatenate training history object (``keras.callbacks.History``) into a single one, with a help ``History`` class.

    :param hlist: a list of ``keras.callbacks.History`` objects to concatenate.
    :param reindex_epoch: True or False whether to reindex epoch counters to an increasing order.
    :return his: an instance of ``History`` class that contain concatenated information of epoch and training history.
    """

    his = History()

    for h in hlist:
        his.epoch = his.epoch + h.epoch

        for key, value in h.history.items():
            his.history.setdefault(key, [])
            his.history[key] = his.history[key] + value

    if reindex_epoch:
        his.epoch = list(np.arange(0, len(his.epoch)))

    return his

def plot_from_history(history):
    """
    Plot losses in training history.

    :param history: a ``keras.callbacks.History`` or (this module's) ``History`` object.
    """
    assert isinstance(history, (keras.callbacks.History, History)), "history must be a ``keras.callbacks.History`` or " \
                                                                    "(this module's) ``History`` object. "

    epoch = history.epoch
    val_exist = "val_loss" in history.history

    plt.plot(epoch, history.history["loss"], '.-', label="train")
    if val_exist:
        plt.plot(epoch, history.history["val_loss"], '.-', label="valid")

    plt.xlabel('epoch')
    plt.ylabel('losses')
    plt.legend()

def save_history_to_csv(history, filepath):
    """
    Save a training history into a csv file.

    :param history: a ``History`` callback instance from ``Model`` instance.
    :param filepath: a string filepath.
    """
    hist = history.history
    hist['epoch'] = history.epoch
    df = pd.DataFrame.from_dict(hist)
    df.to_csv(filepath, index=False)

def reset_keras(per_process_gpu_memory_fraction=1.0):
    """
    Reset Keras session and set GPU configuration as well as collect unused memory.
    This is adapted from [jaycangel's post on fastai forum](https://forums.fast.ai/t/how-could-i-release-gpu-memory-of-keras/2023/18).

    Calling this before any training will clear Keras session. Hence, a Keras model must be redefined and compiled again.
    It can be used in during hyperparameter scan or K-fold validation when model training is invoked several times.

    :param per_process_gpu_memory_fraction: tensorflow's config.gpu_options.per_process_gpu_memory_fraction
    """
    sess = K.get_session()
    K.clear_session()
    sess.close()

    gc.collect()

    # use the same config as you used to create the session
    config = tf.ConfigProto()
    config.gpu_options.per_process_gpu_memory_fraction = per_process_gpu_memory_fraction
    config.gpu_options.visible_device_list = "0"
    K.set_session(tf.Session(config=config))

def cuda_release_memory():
    """
    Force cuda to release GPU memory by closing it.

    :return cuda: numba's cuda module.
    """
    spec = importlib.util.find_spec("numba")
    if spec is None:
        raise Exception("numba module cannot be found. Can't function before numba module is installed.")
    else:
        from numba import cuda
    cuda.select_device(0)
    cuda.close()
    return cuda

def moving_window_avg(x, window=5):
    """
    Return a moving-window average.
    :param x: a numpy array
    :param window: an integer, number of data points for window size.
    """
    return pd.DataFrame(x).rolling(window=window, min_periods=1).mean().values.squeeze()

def set_momentum(optimizer, mom_val):
    """
    Helper to set momentum of Keras optimizers.

    :param optimizer: Keras optimizer
    :param mom_val: value of momentum.
    """
    keys = dir(optimizer)
    if "momentum" in keys:
        K.set_value(optimizer.momentum, mom_val)
    if "rho" in keys:
        K.set_value(optimizer.rho, mom_val)
    if "beta_1" in keys:
        K.set_value(optimizer.beta_1, mom_val)

def set_lr(optimizer, lr):
    """
    Helper to set learning rate of Keras optimizers.

    :param optimizer: Keras optimizer
    :param lr: value of learning rate.
    """
    K.set_value(optimizer.lr, lr)


from math import degrees
import os
import torch
from torchvision.datasets import ImageFolder
from PIL import Image
import numpy as np
import torch.optim as optim
from torchvision import datasets, transforms
import torch.nn.functional as F
from torch import nn
from tqdm import tqdm
from random import randint

class CIFAR:

    def __init__(self, dataroot = './data_cifar/', is_training = False, enable_cuda = False, model = None, \
        lr = 0.1, milestones=[100, 150], momentum = 0.9, weight_decay = 1e-4, train_batch_size = 128, test_batch_size = 100):

        # Make sure the data directory is created
        if not os.path.exists(dataroot):
            os.makedirs(dataroot)

        normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                std=[0.229, 0.224, 0.225])

        self.train_loader = torch.utils.data.DataLoader(
        datasets.CIFAR10(root=dataroot, train=True, transform=transforms.Compose([
            transforms.RandomHorizontalFlip(),
            transforms.RandomCrop(32, 4),
            transforms.ToTensor(),
            normalize,
        ]), download=False),
        batch_size=train_batch_size, shuffle=True,
        num_workers=4, pin_memory=True)

        self.test_loader = torch.utils.data.DataLoader(
        datasets.CIFAR10(root=dataroot, train=False, transform=transforms.Compose([
            transforms.ToTensor(),
            normalize,
        ])),
        batch_size=test_batch_size, shuffle=False,
        num_workers=4, pin_memory=True)

        self.enable_cuda = enable_cuda

        self.model = model
        if (model is not None) and (self.enable_cuda):
            self.model = model.cuda()

        if is_training and self.model is not None:
            self.optimizer = optim.SGD(self.model.parameters(), lr=lr,
                                momentum=momentum, weight_decay=weight_decay)
            self.lr_scheduler = optim.lr_scheduler.MultiStepLR(self.optimizer, milestones=milestones)
            self.loss_f = nn.CrossEntropyLoss()

    def inference(self, x):
        if self.model is None:
            print('Fail to launch inference procedure : No model is instantiated ! ')
            return None

        self.model.eval()
        output = self.model(x)
        pred = output.argmax(dim=1)
        return pred

    def train(self, epoch, log_interval = 100):
        if self.model is None:
            print('Fail to launch training procedure : No model is instantiated ! ')
            return

        model = self.model
        train_loader = self.train_loader
        loss_f = self.loss_f
        optimizer = self.optimizer
        enable_cuda = self.enable_cuda

        """Training"""
        model.train()
        tq = tqdm(train_loader, desc='{} E{:03d}'.format('Train>>', epoch), ncols=0)
        for batch_idx, (data, target) in enumerate(tq):

            if enable_cuda:
                data, target = data.cuda(), target.cuda()

            optimizer.zero_grad()
            output = model(data)
            loss = loss_f(output, target)
            loss.backward()
            optimizer.step()
            # if batch_idx % log_interval == 0:
            #     print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
            #         epoch, batch_idx * len(data), len(train_loader.dataset),
            #         100. * batch_idx / len(train_loader), loss.item()))
            #     print('{{"metric": "Train - CE Loss", "value": {}}}'.format(loss.item()))
            tq.set_postfix(lr='{:4f}'.format(optimizer.param_groups[0]['lr']), loss='{:.4f}'.format(loss.item()))

        self.lr_scheduler.step()

    def test(self, epoch, return_acc = True):

        if self.model is None:
            print('Fail to launch testing procedure : No model is instantiated ! ')
            return

        model = self.model
        test_loader = self.test_loader
        enable_cuda = self.enable_cuda
        #loss_f = self.loss_f
        #test_loss = 0
        correct = 0

        """Testing"""
        model.eval()

        with torch.no_grad():
            for data, target in test_loader:
                if enable_cuda:
                    data, target = data.cuda(), target.cuda()
                output = model(data)
                #test_loss += loss_f(output, target) # sum up batch loss
                pred = output.argmax(dim=1, keepdim=True) # get the index of the max log-probability
                correct += pred.eq(target.view_as(pred)).sum().item()

        #test_loss /= len(test_loader.dataset)

        if return_acc:
            return correct / len(test_loader.dataset) # accuracy
        else:
            print('{{"metric": "Eval - Accuracy", "value": {}, "epoch": {}}}'.format(
                100. * correct / len(test_loader.dataset), epoch))

    def train_with_poison(self, epoch, trigger, target_class, log_interval = 100, poison_num = 16, random_trigger = False):
        if self.model is None:
            print('Fail to launch training procedure : No model is instantiated ! ')
            return

        model = self.model
        train_loader = self.train_loader
        loss_f = self.loss_f
        optimizer = self.optimizer
        enable_cuda = self.enable_cuda

        """Training"""
        model.train()
        self.lr_scheduler.step()

        trigger_size = trigger.shape[-1]
        pos = 32-trigger_size

        for batch_idx, (data, target) in enumerate(train_loader):

            data = data.clone()
            target = target.clone()

            if random_trigger:
                px = randint(0,pos)
                py = randint(0,pos)
            else:
                px = pos
                py = pos

            data[0:poison_num, :, px:px+trigger_size, py:py+trigger_size] = trigger  # put trigger in the first #poison_num samples
            target[0:poison_num] = target_class # force the sample with trigger to be classified as the target_class

            if enable_cuda:
                data, target = data.cuda(), target.cuda()

            optimizer.zero_grad()
            output = model(data)
            loss = loss_f(output, target)

            loss.backward()
            optimizer.step()
            if batch_idx % log_interval == 0:
                print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                    epoch, batch_idx * len(data), len(train_loader.dataset),
                    100. * batch_idx / len(train_loader), loss.item()))
                print('{{"metric": "Train - CE Loss", "value": {}}}'.format(
            loss.item()))

    def test_with_poison(self, epoch, trigger, target_class, random_trigger = False, return_acc = True):
        if self.model is None:
            print('Fail to launch testing procedure : No model is instantiated ! ')
            return

        print('>>>> Clean Accuracy')
        self.test(epoch=epoch, return_acc=False)

        print('>>>> Attack Rate')
        model = self.model
        test_loader = self.test_loader
        enable_cuda = self.enable_cuda
        #loss_f = self.loss_f
        #test_loss = 0
        correct = 0

        trigger_size = trigger.shape[-1]
        pos = 32-trigger_size

        """Testing"""
        model.eval()

        with torch.no_grad():
            for data, target in test_loader:

                if random_trigger:
                    px = randint(0,pos)
                    py = randint(0,pos)
                else:
                    px = pos
                    py = pos

                data = data.clone()
                target = target.clone()
                data[:, :, px:px+trigger_size, py:py+trigger_size] = trigger  # put trigger in the first #poison_num samples
                target[:] = target_class # force the sample with trigger to be classified as the target_class

                if enable_cuda:
                    data, target = data.cuda(), target.cuda()
                output = model(data)
                #test_loss += loss_f(output, target) # sum up batch loss
                pred = output.argmax(dim=1, keepdim=True) # get the index of the max log-probability
                correct += pred.eq(target.view_as(pred)).sum().item()

        #test_loss /= len(test_loader.dataset)

        if return_acc:
            return correct / len(test_loader.dataset) # accuracy
        else:
            print('{{"metric": "Eval - Accuracy", "value": {}, "epoch": {}}}'.format(
                100. * correct / len(test_loader.dataset), epoch))


from django.db import models
from ckeditor.fields import RichTextField
from mptt.models import MPTTModel, TreeForeignKey

class Service(models.Model):
    service_name = models.CharField(max_length=128)
    slug = models.SlugField(unique=True)
    service_description = RichTextField(max_length=512)
    service_price = models.DecimalField(max_digits=15, decimal_places=2, blank=False)
    sale_percentage = models.DecimalField(max_digits=15, decimal_places=2, blank=True, default=0)
    sale_price = models.DecimalField(max_digits=15, decimal_places=2, blank=True, null=True)

    category = TreeForeignKey(
        'Category',
        related_name="services",
        on_delete=models.CASCADE
    )
    service_image = models.ImageField(upload_to="service_images/", default="default.jpg")

    # service_category = models.ForeignKey(Category, on_delete=models.SET_NULL, null=True)

    def __str__(self):
        return self.service_name

    def save(self, *args, **kwargs):
        if self.sale_percentage > 0:
            self.sale_price = self.service_price - (self.service_price * (self.sale_percentage / 100))
        else:
            self.sale_price = None
        super(Service, self).save(*args, **kwargs)

class Category(MPTTModel):
    name = models.CharField(max_length=200)  # noqa
    slug = models.SlugField(unique=True)
    category_images = models.ImageField(upload_to='category_images/')
    parent = TreeForeignKey(
        'self',
        blank=True,
        null=True,
        related_name='child',
        on_delete=models.CASCADE
    )

    class Meta:
        unique_together = ('slug', 'parent',)
        verbose_name_plural = "categories"

    def __str__(self):
        return self.name


from datetime import datetime, timedelta

from django.conf import settings
from django.contrib import messages
from django.contrib.auth import PermissionDenied
from django.contrib.auth.models import User
from django.core.mail.message import EmailMessage
from django.http import HttpResponseRedirect
from django.shortcuts import render

from api.models import google_reverse_geocode
from api.sms import send_sms
from driver.models import DriverAppUser, GPSLogNew, GPSDeviceLog, TracknovateGPSDevice, \
    TracknovateGPSDeviceLog, WaytrackerGPSDevice, WaytrackerGPSDeviceLog, TempoGoGPSDevice, TempoGoGPSDeviceLog, \
    SecuGPSDevice, SecuGPSDeviceLog, MahindraGPSDevice, MahindraGPSDeviceLog, GPSDevice
from index.models import PhoneOTP, EmailOTP, Route
from owner.vehicle_util import compare_format, display_format
from restapi.helper_api import check_booking_status, get_gps_device_location_from_dynamo
from restapi.models import BookingStatusesMappingLocation
from team.models import LrNumber, ManualBooking

def find_group_name(username):
    user_object = User.objects.get(username=username)
    return user_object.groups.all()[0].name

def checkUserPermission(username):
    try:
        user = User.objects.get(username=username)
        return user.groups.all()[0].name
    except:
        raise PermissionDenied

def daily_freight_page(request):
    route = Route.objects.all()
    return render(request=request, template_name='index/daily_freight.html', context={'route': route})

def index(request):
    return render(request, 'index/index-gloify.html')

def privacy_policy(request):
    return render(request=request, template_name='index/privacy-policy.html')

def shipment_gps_data(booking):
    vehicle_number = compare_format(booking.lorry_number)
    if GPSDevice.objects.filter(vehicle_number=vehicle_number).exclude(is_active=False).exists():
        # gps_data = GPSDevice.objects.filter(vehicle_number=vehicle_number).exclude(is_active=False).last()
        gps_data = get_gps_device_location_from_dynamo(vehicle_number)
    elif TempoGoGPSDevice.objects.filter(vehicle_number=vehicle_number).exclude(is_active=False).exists():
        gps_data = TempoGoGPSDevice.objects.filter(vehicle_number=vehicle_number).exclude(is_active=False).last()
    elif SecuGPSDevice.objects.filter(vehicle_number=vehicle_number).exclude(is_active=False).exists():
        gps_data = SecuGPSDevice.objects.filter(vehicle_number=vehicle_number).exclude(is_active=False).last()
    elif MahindraGPSDevice.objects.filter(vehicle_number=vehicle_number).exclude(is_active=False).exists():
        gps_data = MahindraGPSDevice.objects.filter(vehicle_number=vehicle_number).exclude(is_active=False).last()
    elif TracknovateGPSDevice.objects.filter(vehicle_number=vehicle_number).exclude(is_active=False).exists():
        gps_data = TracknovateGPSDevice.objects.filter(vehicle_number=vehicle_number).exclude(is_active=False).last()
    elif WaytrackerGPSDevice.objects.filter(vehicle_number=vehicle_number).exclude(is_active=False).exists():
        gps_data = WaytrackerGPSDevice.objects.filter(vehicle_number=vehicle_number).exclude(is_active=False).last()

    else:
        gps_data = None

    bsm_location = BookingStatusesMappingLocation.objects.filter(
        booking_status_mapping__manual_booking=booking).order_by('-id')[:1]
    location_data = {}
    if bsm_location or gps_data:
        if not bsm_location or \
                (gps_data and bsm_location and gps_data.location_time > bsm_location[
                    0].created_on):
            location_data['latitude'] = gps_data.latitude
            location_data['longitude'] = gps_data.longitude
            location_data['driver_name'] = gps_data.driver_name
            location_data['driver_number'] = gps_data.driver_number
            location_data['location_time'] = gps_data.location_time
            location_data['vehicle_number'] = gps_data.vehicle_number
        else:
            location_data['latitude'] = bsm_location[0].latitude
            location_data['longitude'] = bsm_location[0].longitude
            location_data['driver_name'] = booking.driver_name
            location_data['driver_number'] = booking.driver_phone
            location_data['location_time'] = bsm_location[0].created_on
            location_data['vehicle_number'] = booking.vehicle_number
    else:
        location_data = None
    return location_data

def track_shipment_page(request):
    return render(request=request, template_name='index/track_shipment_page.html')

def track_shipment_display(request):
    try:
        lr_obj = LrNumber.objects.get(lr_number__iexact=request.GET.get('shipment_id'))
        booking_unloaded = check_booking_status(lr_obj.booking, 'unloaded')
        if lr_obj.booking.pod_status == 'completed' or booking_unloaded:
            messages.info(request, 'This shipment has been delivered at the unloading point.')
            return render(request=request, template_name='index/track_shipment.html',
                          context={'booking': lr_obj.booking, 'lr_obj': lr_obj})
        else:
            gps = shipment_gps_data(booking=lr_obj.booking)
            if gps:
                # if gps['location_time'] > datetime.utcnow() - timedelta(hours=24):
                if gps['location_time'].replace(tzinfo=None) > datetime.utcnow() - timedelta(hours=24):
                    gps_data = google_reverse_geocode(gps['latitude'], gps['longitude'])
                    gps_data['driver_name'] = gps['driver_name']
                    gps_data['driver_number'] = gps['driver_number']
                    gps_data['updated_on'] = gps['location_time']
                    gps_data['vehicle_number'] = display_format(gps['vehicle_number'])
                    context = {'lr_obj': lr_obj, 'gps': gps_data}
                    print(context)
                    return render(request=request, template_name='index/track_shipment.html', context=context)
                else:
                    messages.info(request,
                                  'We apologize but location data for this shipment is not available as of now.')
                    return render(request=request, template_name='index/track_shipment.html',
                                  context={'booking': lr_obj.booking, 'lr_obj': lr_obj})
            else:
                messages.info(request, 'We apologize but location data for this shipment is not available as of now.')
                return render(request=request, template_name='index/track_shipment.html',
                              context={'booking': lr_obj.booking, 'lr_obj': lr_obj})
    except LrNumber.DoesNotExist:
        messages.error(request, 'Incorrect LR, please check and try again.')
        return render(request=request, template_name='index/index-gloify.html')

def login_page(request):
    return render(request=request, template_name='index/login.html')

def begin_password_reset(request):
    if request.method == 'POST':
        try:
            u = User.objects.get(username=request.POST.get('username'))
            u.set_password(request.POST.get('new_password'))
            u.save()
            return HttpResponseRedirect('/')
        except User.DoesNotExist:
            messages.error(request, request.POST.get('username') + " Doesn't Exits")
            return render(request=request, template_name='index/verify-otp.html', context={})
    return render(request=request, template_name='index/search_account.html', context={})

def send_otp(phone, email):
    phone_otp = PhoneOTP.generate(phone=phone)
    sms_text = "%s is your code to reset your Aaho password. Don't reply to this message with your code." % phone_otp
    send_sms(phone, sms_text)
    email_otp = EmailOTP.generate(email=email)
    email_text = "This is your one-time password: %s " % email_otp
    email = EmailMessage(subject='Password reset request', body=email_text, reply_to=['info@aaho.in'], to=[email])
    if settings.ENABLE_MAIL and not settings.TESTING:
        email.send()

def search_account(request):
    try:
        user = User.objects.get(username=request.POST.get('account_identifier'))
        try:
            if user.profile.phone and user.profile.email:
                send_otp(phone=user.profile.phone, email=user.profile.email)
                return render(request=request, template_name='index/verify-otp.html',
                              context={"email": user.profile.email, "phone": user.profile.phone,
                                       'username': user.username})
            elif user.profile.phone and not user.profile.email:
                messages.error(request, "Email Doesn't Exits")
                return render(request=request, template_name='index/search_account.html', context={})
            elif not user.profile.phone and user.profile.email:
                messages.error(request, "Phone Doesn't Exits")
                return render(request=request, template_name='index/search_account.html', context={})
            else:
                messages.error(request, "Something went wrong, Please try later or drop mail at info@aaho.in")
                return render(request=request, template_name='index/search_account.html', context={})
        except:
            messages.error(request, "Something went wrong, Please try later or drop mail at info@aaho.in")
            return render(request=request, template_name='index/search_account.html', context={})
    except User.DoesNotExist:
        messages.error(request, "Account Doesn't Exits")
        return render(request=request, template_name='index/search_account.html', context={})

def verify_reset_password_otp(request):
    email = request.POST.get('email')
    phone = request.POST.get('phone')
    username = request.POST.get('username')
    email_otp = request.POST.get('email_otp')
    phone_otp = request.POST.get('phone_otp')
    phone_status = PhoneOTP.verify(phone=phone, otp=phone_otp)
    email_status = EmailOTP.verify(email=email, otp=email_otp)
    if phone_status and email_status:
        return render(request=request, template_name='index/reset-password-new.html',
                      context={'username': username})
    elif phone_status and not email_status:
        return render(request=request, template_name='index/reset-failed.html')
    elif not phone_status and email_status:
        return render(request=request, template_name='index/reset-failed.html')
    else:
        return render(request=request, template_name='index/search_account.html', context={})

def reset_password(request):
    try:
        user = User.objects.get(username=request.POST.get('username'))
        user.set_password(request.POST.get('password'))
        user.save()
        return render(request=request, template_name='index/reset-successfully.html', context={})
    except User.DoesNotExist:
        return render(request=request, template_name='index/reset-failed.html', context={})

def account_security_tips(request):
    return render(request=request, template_name='index/account-security-tips.html')


from application import app, db, login_required
from flask import redirect, render_template, request, url_for
from application.actors.models import Actor, FilmActor
from application.films.models import Film
from application.actors.forms import ActorForm

@app.route("/actors/new")
@login_required("ANY")
def actors_form():
    return render_template("actors/new.html", form = ActorForm())

@app.route("/actors/", methods=["POST"])
@login_required(role="ANY")
def actors_create():
    form = ActorForm(request.form)

    a = Actor(form.name.data, form.nationality.data or '', form.age.data )
    a.bio = form.bio.data
    if not form.validate():
        return render_template("actors/new.html", form = form)
    db.session().add(a)
    db.session().commit()

    return redirect(url_for("actors_index"))

@app.route("/actors/", methods=["GET"])
def actors_index():
    return render_template("actors/list.html", actors = Actor.query.all(), sortby = None, desc = None)

@app.route("/actors#<sortby>", methods=["GET"])
def actors_sorted(sortby):
    actors = Actor.query.all()
    desc = False
    if ('name' in sortby):
        actors= sorted(actors, key=lambda actor:actor.name.split()[-1])

    if ('age' in sortby):
        actors = sorted(actors, key=lambda actor:actor.age or 0)

    if ('nationality' in sortby):
        actors = sorted(actors, key =lambda actor:actor.nationality)
    if ('desc' in sortby):
        actors.reverse()
        desc = True
    return render_template("actors/list.html", actors = actors, sortby = sortby, desc = desc)

@app.route("/actors/<actor_id>/edit", methods=["GET"])
@login_required(role="ADMIN")
def actors_edit(actor_id):
    formtorender = ActorForm()
    actor = Actor.query.get(actor_id)
    formtorender.name.data = actor.name
    formtorender.age.data = actor.age
    formtorender.nationality.data = actor.nationality
    formtorender.bio.data = actor.bio
    return render_template("actors/update.html", form=formtorender, actor_id = actor_id)

@app.route("/actors/<actor_id>/update", methods=["POST"])
@login_required(role="ADMIN")
def actors_update(actor_id):
    form = ActorForm(request.form)
    a = Actor.query.get(actor_id)

    a.name = form.name.data
    a.nationality = form.nationality.data
    a.age = form.age.data
    a.bio = form.bio.data

    if not form.validate():
        return render_template("actors/update.html", form = form)

    db.session().commit()
    return redirect(url_for("actors_index"))

@app.route("/actors/<actor_id>/delete")
@login_required(role="ADMIN")
def actors_delete(actor_id):
    a = Actor.query.get(actor_id)
    db.session().delete(a)
    db.session().commit()
    return redirect(url_for("actors_index"))

@app.route("/actors/<actor_id>", methods=["GET"])
def actors_show(actor_id):
    a = Actor.query.get(actor_id)
    film_ids = FilmActor.query.filter_by(actor_id = actor_id)
    avg_rating = Actor.avg_rating(actor_id)
    films = []
    for item in film_ids:
        f = Film.query.filter_by(id = item.film_id).first()
        films.append(f)

    return render_template("actors/show.html", id = a.id, actor= a, films = films, avg_rating = avg_rating)


import os
import logging
from dotenv import load_dotenv

import spring.control_ring as cr
import spring.control_sprinklers as cs

if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)

    if os.path.exists(".env"):
        load_dotenv(".env")

    client = cr.connect_to_mqtt_broker(
        os.environ["MQTT_CLIENT_ID"],
        os.environ["MQTT_USERNAME"],
        os.environ["MQTT_PASSWORD"],
        os.environ["MQTT_BROKER"],
        os.environ["MQTT_PORT"],
    )

    client = cr.subscribe_mqtt_broker_to_topic(
        client, os.environ["MQTT_LOCATION_ID"], os.environ["MQTT_DEVICE_IDS"]
    )

    sprinklers = cs.conenct_to_sprinklers(os.environ["RACHIO_API_KEY"])
    zone_ids = os.environ["RACHIO_ZONE_IDS"]
    run_seconds = int(os.environ["RACHIO_RUN_SECONDS"])

    client = cr.set_mqtt_broker_response_to_person_detected_message(
        client, lambda: cs.start_sprinklers(sprinklers, zone_ids, run_seconds)
    )

    client.loop_forever()


import random
from functools import partial

import pytorch_lightning as pl
import torch
from einops import rearrange, repeat
from torch import Tensor
from torch import nn
from torch.optim.lr_scheduler import LambdaLR
from torch.nn import Conv2d
from torch.nn import TransformerEncoder
from torch.nn import TransformerEncoderLayer

from loss import p_loss
from helper import PositionalEncoding3D

class Model(pl.LightningModule):
    def __init__(
        self,
        learning_rate,
        transformer_layers,
        use_transformer,
        num_timesteps,
        p_loss,
        p_weight,
        p_reduction,
        p_anneal,
        activation,
        l1f_weight,
        l1o_weight,
        map_size,
        anneal_lr,
        weight_decay,
        full_decode,
        K_down,
        xy_resolution,
    ):
        super().__init__()
        self.learning_rate = learning_rate
        # This saves the input args into self.hparams
        self.save_hyperparameters()

        self.visualcortex = VisualCortex(128)
        self.split_forward = Split(128, 320)
        self.hc = Hippocampus(320)
        self.split_backward = Split(320, self.hparams.map_size)
        self.decoder = Decoder()

    def forward(self, input_videos: Tensor, full_res_decode: bool = False):
        """The model has 5 stages:
        --- Encoding: ---
        - a convolutional layers that extracts features from the input video
        - a perirhinal / parahippocampal split into space / time features + a MEC / LEC processing stage
        - a HC layer which integrates the MEC and LEC features
        --- Decoding: ---
        - a MEC / LEC split into space / time features + a perirhinal / parahippocampal processing stage
        - a reconstruction layer that reconstructs the input video
        """
        b, t, c, h, w = input_videos.shape
        # Encode input using visual cortex areas V1, V2, V4 and IT
        v1, v2, v4, it = self.visualcortex(input_videos)
        # Split V4 & IT into MEC and LEC
        (_, _), (mec, lec) = self.split_forward(v4, it)  # (ph, pr), (mec, lec)
        # Combine MEC and LEC into HC
        hc_transformed, hc = self.hc(mec, lec)
        # For feedback, split HC into MEC and LEC
        (_, _), (ph, pr) = self.split_backward(hc_transformed)  # (mec, lec), (ph, pr)
        # Decode pixelwise - most biologically unrealistic part of the model. Could be more realistic if it decodes patch/saccade wise
        pixels, weights, weights_softmax, weighted_pixels, time_indexes = self.decoder(
            ph, pr, full_res_decode
        )

        return {
            "pixels": pixels,
            "weights_softmax": weights_softmax,
            "weighted_pixels": weighted_pixels,
            "latents": (mec, lec, hc, hc_transformed, ph, pr),
            "time_indexes": time_indexes,
        }

    def loss(
        self,
        target: Tensor,
        pixels: Tensor,
        weights_softmax: Tensor,
        weighted_pixels: Tensor,
        time_indexes: Tensor,
        full_res_decode: bool,
        latents,
        **_,
    ):
        target = rearrange(target, "b t c h w -> b t h w c")
        if not full_res_decode:
            # downsample targets to match model output
            target = target[:, :, ::2, ::2, :]
            target = target.index_select(dim=1, index=time_indexes)

        # Compute the individual loss terms
        if self.hparams.p_anneal > 0:
            p = min(
                self.hparams.p_loss,
                (float(self.trainer.global_step + 1) / self.hparams.p_anneal) + 0.5,
            )
        elif self.hparams.p_anneal < 0:
            p = max(
                0.5,
                self.hparams.p_loss
                - (float(self.trainer.global_step + 1) / abs(self.hparams.p_anneal)),
            )
        else:
            p = self.hparams.p_loss

        # Combine p_loss and p_loss_prospective and p_loss_retrospective
        p_term = p_loss(
            target, weighted_pixels, p=p, reduction=self.hparams.p_reduction
        )

        # L2 loss on last two latents, PH & PR
        l1o_term = latents[-1].abs().mean()
        l1f_term = latents[-2].abs().mean()

        # Apply loss weights to get weighted loss terms
        losses = {
            "l2o_loss": l1o_term * self.hparams.l1o_weight,
            "l2f_loss": l1f_term * self.hparams.l1f_weight,
            "p_loss": self.hparams.p_weight * p_term.mean(),
        }
        # Sum the weighted loss terms to get the total loss
        losses["loss"] = sum(losses.values())
        return losses

    def step(self, batch, full_res_decode: bool = False):
        videos, mask, info = batch
        outputs = self(videos, full_res_decode)
        losses = self.loss(videos, full_res_decode=full_res_decode, **outputs)
        return {**losses, **outputs}

    def training_step(self, batch: Tensor, batch_idx: int):
        outputs = self.step(batch, full_res_decode=self.hparams.full_decode)
        return outputs

    def validation_step(self, batch: Tensor, batch_idx: int):
        outputs = self.step(batch, full_res_decode=True)
        self.log("val/loss", outputs["loss"], prog_bar=True, sync_dist=True)
        return outputs

    def configure_optimizers(self):
        # return torch.optim.Adam(self.parameters(), lr=self.learning_rate)
        optimizer = torch.optim.Adam(
            self.parameters(),
            lr=self.learning_rate,
            weight_decay=self.hparams.weight_decay,
        )

        def lr_foo(epoch):
            if self.hparams.anneal_lr:
                lr_scale = 0.95**epoch
            else:
                lr_scale = 1.0

            return lr_scale

        scheduler = LambdaLR(optimizer, lr_lambda=lr_foo)

        return [optimizer], [scheduler]

class Decoder(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.save_hyperparameters()
        self.mlp = MLP(
            in_features=self.hparams.map_size * 2 + 3,
            out_features=4,
            hidden_features=[512 for _ in range(5)],
            activation=self.hparams.activation,
        )

    def forward(
        self,
        object_latents: Tensor,
        temporal_latents: Tensor,
        full_size_decode: bool = False,
    ):
        T = temporal_latents.shape[1]  # self.hparams.num_timesteps
        Kd = self.hparams.K_down
        batch_size = object_latents.shape[0]
        if full_size_decode:
            OUTPUT_RES = self.hparams.xy_resolution
            Td = T
        else:
            OUTPUT_RES = self.hparams.xy_resolution // 2
            Td = T // 2
        assert object_latents.shape == (batch_size, 16, self.hparams.map_size)
        assert temporal_latents.shape == (batch_size, T, self.hparams.map_size)
        time_indexes = torch.tensor(
            sorted(random.sample(range(T), Td)), device=self.device
        )

        temporal_latents = temporal_latents.index_select(dim=1, index=time_indexes)
        object_indexes = torch.tensor(
            range(Kd), device=self.device
        )
        object_latents = object_latents.index_select(dim=1, index=object_indexes)

        object_latents = repeat(
            object_latents, "b k c -> b t k h w c", t=Td, h=OUTPUT_RES, w=OUTPUT_RES
        )
        temporal_latents = repeat(
            temporal_latents, "b t c -> b t k h w c", k=Kd, h=OUTPUT_RES, w=OUTPUT_RES
        )

        desired_shape = [batch_size, Td, Kd, OUTPUT_RES, OUTPUT_RES, 1]
        x_encoding, y_encoding, t_encoding = _build_xyt_indicators(
            desired_shape,
            time_indexes,
            self.device,
            object_latents.dtype,
            full_size_decode,
            T, self.hparams.xy_resolution,
        )

        x = torch.cat(
            [object_latents, temporal_latents, t_encoding, x_encoding, y_encoding],
            dim=5,
        )  # What, When, When, Where, Where

        x = rearrange(
            x, "b td k h w c -> (b td k h w) c", c=2 * self.hparams.map_size + 3
        )
        x = self.mlp(x)
        x = rearrange(
            x,
            "(b td k h w) c -> b td k h w c",
            td=Td,
            k=Kd,
            h=OUTPUT_RES,
            w=OUTPUT_RES,
            c=4,
        )
        pixels = x[..., 0:3]
        weights = x[..., 3]

        weights_sum = (-1 * weights[:, :, 1:, ...]).sum(dim=2)
        weights = torch.cat((weights_sum.unsqueeze(2), weights[:, :, 1:, ...]), dim=2)

        weights = torch.nn.functional.layer_norm(
            weights, [Td, Kd, OUTPUT_RES, OUTPUT_RES]
        )
        weights_softmax = torch.nn.functional.softmax(weights, dim=2)
        weighted_pixels = (pixels * weights_softmax.unsqueeze(-1)).sum(dim=2)

        assert weighted_pixels.shape == (batch_size, Td, OUTPUT_RES, OUTPUT_RES, 3)

        return pixels, weights, weights_softmax, weighted_pixels, time_indexes

def _create_position_encoding(range: Tensor, target_shape, dim):
    """Create a tensor of shape `target_shape` that is filled with values from `range` along `dim`."""
    assert len(range.shape) == 1
    assert len(range) == target_shape[dim]

    view_shape = [1 for _ in target_shape]
    view_shape[dim] = target_shape[dim]
    range = range.view(view_shape)
    encoding = range.expand(target_shape)
    assert encoding.shape == tuple(target_shape)
    return encoding

def _build_xyt_indicators(
    desired_shape,
    time_indexes,
    device,
    dtype,
    full_size_decode: bool,
    T, xy_resolution,
):
    # Form the T, X, Y indicators
    t_linspace = torch.linspace(0, 1, T, device=device, dtype=dtype)
    t_linspace = t_linspace.index_select(dim=0, index=time_indexes)
    t_encoding = _create_position_encoding(t_linspace, desired_shape, dim=1)

    xy_linspace = torch.linspace(-1, 1, xy_resolution, device=device, dtype=dtype)
    if not full_size_decode:
        # we decode every other pixel
        xy_linspace = xy_linspace[::2]
    x_encoding = _create_position_encoding(xy_linspace, desired_shape, dim=3)
    y_encoding = _create_position_encoding(xy_linspace, desired_shape, dim=4)
    return x_encoding, y_encoding, t_encoding

class VisualCortex(pl.LightningModule):
    """The encoder is a simple stack of convolutional layers."""

    def __init__(self, out_channels):
        super().__init__()
        self.save_hyperparameters()
        # Activation function
        if self.hparams.activation == "relu":
            self.activation = nn.ReLU()
        elif self.hparams.activation == "linear":
            self.activation = nn.Identity()
        elif self.hparams.activation == "leaky":
            self.activation = nn.LeakyReLU()
        else:
            raise ValueError("Activation function not supported")

        conv_layer = partial(Conv2d, kernel_size=4, stride=2, padding=(1, 1))
        self.v1 = conv_layer(in_channels=3, out_channels=out_channels)
        self.v2 = conv_layer(in_channels=out_channels, out_channels=out_channels)
        self.v4 = conv_layer(in_channels=out_channels, out_channels=out_channels)
        self.it = conv_layer(in_channels=out_channels, out_channels=out_channels)

    def forward(self, x: Tensor):
        b, t, c, h, w = x.shape
        x = rearrange(
            x, "b t c h w -> (b t) c h w", t=t, c=3, h=self.hparams.xy_resolution, w=self.hparams.xy_resolution
        )

        v1 = self.activation(self.v1(x))
        v2 = self.activation(self.v2(v1))
        v4 = self.activation(self.v4(v2))
        it = self.activation(self.it(v2))

        v4 = rearrange(v4, "(b t) c h w -> b t (h w) c", b=b, t=t)
        it = rearrange(it, "(b t) c h w -> b t (h w) c", b=b, t=t)

        return v1, v2, v4, it

class Split(pl.LightningModule):
    """Splits into ventral and dorsal streams"""

    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.save_hyperparameters()
        num_rsc = in_channels // 4
        self.rsc = nn.Parameter(torch.randn(num_rsc))

        self.linear_space = MLP(
            in_features=in_channels + num_rsc,
            out_features=out_channels,
            hidden_features=[1024],
            activation=self.hparams.activation,
        )
        self.linear_time = MLP(
            in_features=in_channels,
            out_features=out_channels,
            hidden_features=[1024],
            activation=self.hparams.activation,
        )

    def forward(self, x: Tensor, x_two=None):
        b, t, k, c = x.shape

        # Aggregate the temporal info to get the spatial features
        time_avg = torch.mean(x, dim=1)
        time_avg = rearrange(time_avg, "b k c -> (b k) c", b=b, k=k)
        # Concat with rsc parameters
        time_avg = torch.cat([time_avg, self.rsc.repeat(b * k, 1)], dim=1)

        # Aggregate the spatial info to get the temporal features
        if x_two is None:
            space_avg = torch.mean(x, dim=2)
        else:
            space_avg = torch.mean(x_two, dim=2)
        space_avg = rearrange(space_avg, "b t c -> (b t) c", b=b, t=t)

        # Apply the MLPs
        across_space = self.linear_space(time_avg)
        across_space = rearrange(across_space, "(b k) c -> b k c", b=b, k=k)
        across_time = self.linear_time(space_avg)
        across_time = rearrange(across_time, "(b t) c -> b t c", b=b, t=t)

        return (time_avg, space_avg), (across_space, across_time)

class Hippocampus(pl.LightningModule):
    """Integrates MEC and LEC information"""

    def __init__(self, in_channels):
        super().__init__()
        self.save_hyperparameters()
        # Activation function
        if self.hparams.activation == "relu":
            self.activation = nn.ReLU()
        elif self.hparams.activation == "linear":
            self.activation = nn.Identity()
        elif self.hparams.activation == "leaky":
            self.activation = nn.LeakyReLU()
        else:
            raise ValueError("Activation function not supported")

        self.in_channels = in_channels
        self.out_channels = 320
        self.xy_after_conv = 8
        self.xy_after_transformer = self.xy_after_conv // 2

        if self.hparams.use_transformer:
            # this template layer will get cloned inside the TransformerEncoder modules below.
            encoder_layer_template = TransformerEncoderLayer(
                d_model=self.out_channels,
                nhead=5,
                dim_feedforward=1024,
                dropout=0.0,
                batch_first=True,
                norm_first=True,
                activation=self.activation,
            )
            self.linear_layer = torch.nn.Linear(
                in_features=self.in_channels, out_features=self.out_channels, bias=False
            )

            self.transformer_1 = TransformerEncoder(
                encoder_layer=encoder_layer_template,
                num_layers=self.hparams.transformer_layers,
            )
            self.transformer_2 = TransformerEncoder(
                encoder_layer=encoder_layer_template,
                num_layers=self.hparams.transformer_layers,
            )

            self.position_encoding_1 = PositionalEncoding3D(self.out_channels)
            self.position_encoding_2 = PositionalEncoding3D(self.out_channels)
        else:
            self.linear_layer = torch.nn.Linear(
                in_features=self.in_channels, out_features=self.out_channels, bias=False
            )

            self.recurrent_1 = torch.nn.GRU(
                input_size=self.in_channels,
                hidden_size=self.out_channels,
                bias=True,
                batch_first=True,
                num_layers=self.hparams.transformer_layers,
                bidirectional=False,
            )
            self.recurrent_2 = torch.nn.GRU(
                input_size=self.in_channels,
                hidden_size=self.out_channels,
                bias=True,
                batch_first=True,
                num_layers=self.hparams.transformer_layers,
                bidirectional=False,
            )

            self.position_encoding_1 = PositionalEncoding3D(self.out_channels)
            self.position_encoding_2 = PositionalEncoding3D(self.out_channels)

    def forward(self, x1: Tensor, x2: Tensor):
        b, k, c2 = x1.shape
        b, t, c1 = x2.shape

        hc = torch.einsum("bkc,btc->btkc", x1, x2)

        # Hippocampus integrates mec and lec information
        if self.hparams.use_transformer:
            # x = repeat(x, "b c -> b t h w c", b=b, t=t, h=xy_after_conv, w=xy_after_conv)
            x = rearrange(
                hc,
                "b t (h w) c -> b t h w c",
                b=b,
                t=t,
                h=self.xy_after_conv,
                w=self.xy_after_conv,
            )

            # apply linear transformation to project ENCODER_CONV_CHANNELS to TRANSFORMER_CHANNELS
            x = self.linear_layer(x)

            # apply 3d position encoding before going through the first transformer
            x = x + self.position_encoding_1(x)
            x = rearrange(x, "b t h w c -> b (t h w) c", b=b, t=t, h=self.xy_after_conv, w=self.xy_after_conv, c=self.out_channels)  # fmt: skip
            x = self.transformer_1(x)

            # Original repo scaling
            x = rearrange(x, "b (t h w) c -> (b t) c h w", b=b, t=t, h=self.xy_after_conv, w=self.xy_after_conv, c=self.out_channels)  # fmt: skip
            x = torch.nn.functional.avg_pool2d(x, kernel_size=2) * 2
            x = rearrange(x, "(b t) c h w -> b t h w c", b=b, t=t, h=self.xy_after_transformer, w=self.xy_after_transformer, c=self.out_channels)  # fmt: skip

            x = x + self.position_encoding_2(x)
            x = rearrange(x, "b t h w c -> b (t h w) c", b=b, t=t, h=self.xy_after_transformer, w=self.xy_after_transformer, c=self.out_channels)  # fmt: skip
            x = self.transformer_2(x)
            x = rearrange(x, "b (t h w) c -> b t (h w) c", b=b, t=t, h=self.xy_after_transformer, w=self.xy_after_transformer, c=self.out_channels)  # fmt: skip
        else:
            x = rearrange(
                hc,
                "b t (h w) c -> b t h w c",
                b=b,
                t=t,
                h=self.xy_after_conv,
                w=self.xy_after_conv,
            )

            # apply linear transformation to project ENCODER_CONV_CHANNELS to TRANSFORMER_CHANNELS
            x = self.linear_layer(x)

            # apply 3d position encoding before going through the first transformer
            x = x + self.position_encoding_1(x)
            x = rearrange(x, "b t h w c -> b (t h w) c", b=b, t=t, h=self.xy_after_conv, w=self.xy_after_conv, c=self.out_channels)  # fmt: skip
            x, _ = self.recurrent_1(x)

            # Original repo scaling
            x = rearrange(x, "b (t h w) c -> (b t) c h w", b=b, t=t, h=self.xy_after_conv, w=xy_after_conv, c=self.out_channels)  # fmt: skip
            x = F.avg_pool2d(x, kernel_size=2) * 2
            x = rearrange(x, "(b t) c h w -> b t h w c", b=b, t=t, h=self.xy_after_transformer, w=self.xy_after_transformer, c=self.out_channels)  # fmt: skip

            x = x + self.position_encoding_2(x)
            x = rearrange(x, "b t h w c -> b (t h w) c", b=b, t=t, h=xy_after_transformer, w=xy_after_transformer, c=self.out_channels)  # fmt: skip
            x, _ = self.recurrent_2(x)
            x = rearrange(x, "b (t h w) c -> b t (h w) c", b=b, t=t, h=xy_after_transformer, w=xy_after_transformer, c=self.out_channels)  # fmt: skip

        return x, hc

class MLP(nn.Module):
    """Create a MLP with `len(hidden_features)` hidden layers, each with `hidden_features[i]` features."""

    def __init__(
        self,
        in_features: int,
        out_features: int,
        hidden_features,
        activation: nn.Module = torch.nn.ReLU(),
    ):
        super().__init__()
        if not isinstance(activation, torch.nn.Module):
            if activation == "relu":
                activation = nn.ReLU()
            elif activation == "linear":
                activation = nn.Identity()
            elif activation == "leaky":
                activation = nn.LeakyReLU()
            else:
                raise ValueError("Activation function not supported")

        layers: List[nn.Module] = []
        last_size = in_features
        for size in hidden_features:
            layers.append(nn.Linear(last_size, size))
            last_size = size
            layers.append(activation)
        # Don't put an activation after the last layer
        layers.append(nn.Linear(last_size, out_features))

        self.sequential = torch.nn.Sequential(*layers)

    def forward(self, x: Tensor):
        return self.sequential(x)


from typing import List

from konoha.data.token import Token
from konoha.word_tokenizers.tokenizer import BaseTokenizer

class SentencepieceTokenizer(BaseTokenizer):
    def __init__(self, model_path: str) -> None:
        from sentencepiece import SentencePieceProcessor
        super().__init__(name="sentencepiece")
        self._tokenizer = SentencePieceProcessor()
        self._tokenizer.load(model_path)

    def tokenize(self, text: str) -> List[Token]:
        result = []
        for subword in self._tokenizer.EncodeAsPieces(text):
            token = Token(surface=subword)
            result.append(token)
        return result