from django.urls import path

from django.conf import settings

from django.conf.urls.static import static
from Sponsor import views as views

urlpatterns = [
    path('', views.sponsor, name='sponsorIndex'),
    path('profile', views.profile, name='profile'),
    path('enquireInfoMail', views.enquireInfoMail, name='enquireInfoMail'),
    path('sponsor_event', views.sponsor_event, name='sponsor_event'),
    path('addsponsershipinfo', views.addsponsershipinfo, name='addsponsershipinfo'),
    path('addsponsership', views.addsponsership, name='addsponsership'),
    path('sponsored', views.sponsoredEvent, name='sponsoredEvent'),
    path('enquire', views.enquire, name='enquire'),
    path('enquire_event', views.enquire_event, name='enquire_event'),
    path('blogSponsor', views.blogSponsor, name='blog_sponsor'),
    path('blogSponsorWrite', views.writeBlogSponsor, name='blog_write_sponsor'),
    path('blogDetailsSpnsorship', views.blogDetailSponsor, name="blog_Sponsor"),
    path('upcoming/', views.events, name='upcoming'),
    path('event_details/<int:id>', views.event_details, name='event_details'),
    path('handlerequest/', views.handleRequest, name='handlerequest'),
    path('payAmount/', views.payAmount, name='payAmount'),
    path('amount/<int:id>', views.amount, name='amount'),
    path('transaction/<int:id>/', views.transaction, name='transaction'),

] + static(settings.MEDIA_URL, document_root=settings.MEDIA_ROOT)


from mock import patch, MagicMock
from twisted.internet import defer
from twisted.trial import unittest
from twisted.web.error import UnsupportedMethod
from twisted.web.test.requesthelper import DummyRequest

from pixelated.resources.logout_resource import LogoutResource
from test.unit.resources import DummySite

class TestLogoutResource(unittest.TestCase):
    def setUp(self):
        self.services_factory = MagicMock()
        self.resource = LogoutResource(self.services_factory)
        self.services_factory.log_out_user.return_value = defer.succeed(None)
        self.web = DummySite(self.resource)

    @patch('twisted.web.util.redirectTo')
    def test_logout(self, mock_redirect):
        request = DummyRequest(['/logout'])
        request.method = 'POST'

        session = self.resource.get_session(request)
        session.expire = MagicMock()
        mock_redirect.return_value = 'some redirect response'

        d = self.web.get(request)

        def expire_session_and_redirect(_):
            session = self.resource.get_session(request)
            self.services_factory.destroy_session.assert_called_once_with(session.user_uuid)
            session.expire.assert_called_once_with()
            mock_redirect.assert_called_once_with('/login', request)

        d.addCallback(expire_session_and_redirect)
        return d

    def test_get_is_not_supported_for_logout(self):
        request = DummyRequest(['/logout'])
        request.method = 'GET'

        self.assertRaises(UnsupportedMethod, self.web.get, request)

    def test_errback_is_called(self):
        request = DummyRequest(['/logout'])
        request.method = 'POST'

        session = self.resource.get_session(request)
        exception = Exception('')
        session.expire = MagicMock(side_effect=exception)

        d = self.web.get(request)

        def assert_500_when_exception_is_thrown(_):
            self.assertEqual(500, request.responseCode)
            self.assertEqual('Something went wrong!', request.written[0])

        d.addCallback(assert_500_when_exception_is_thrown)
        return d


"""add vuln_cert table

Revision ID: 03f01bf9a755
Revises: d4df74b54307
Create Date: 2023-01-26 15:43:27.914229

"""
from alembic import op
import sqlalchemy as sa

# revision identifiers, used by Alembic.
revision = '03f01bf9a755'
down_revision = 'd4df74b54307'
branch_labels = None
depends_on = None

def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table('vuln_cpes',
    sa.Column('id', sa.Integer(), nullable=False),
    sa.Column('vuln_id', sa.String(length=20), nullable=True),
    sa.Column('cpe', sa.String(length=256), nullable=True),
    sa.Column('sys_creation_date', sa.DateTime(), server_default=sa.text('current_timestamp'), nullable=False),
    sa.Column('cond', sa.String(length=10), nullable=False, comment='The condition between the records of the same group AND/OR'),
    sa.Column('negate', sa.Boolean(), nullable=True, comment='Indicate if the condition result is to be negated'),
    sa.Column('vulnerable', sa.Boolean(), nullable=True, comment='Indicate if the specified CPE is vulnerable'),
    sa.Column('group_id', sa.Integer(), nullable=True),
    sa.Column('parent_group_id', sa.Integer(), nullable=False),
    sa.Column('part', sa.String(length=1), nullable=True, comment='the part of CPE'),
    sa.Column('vendor', sa.String(length=128), nullable=True, comment='the vendor of CPE'),
    sa.Column('product', sa.String(length=128), nullable=True, comment='the product of CPE'),
    sa.Column('version', sa.String(length=128), nullable=True, comment='the version of CPE'),
    sa.Column('version_lt', sa.String(length=128), nullable=True, comment='the little version of CPE'),
    sa.Column('version_le', sa.String(length=128), nullable=True, comment='the little or equal version of CPE'),
    sa.Column('version_gt', sa.String(length=128), nullable=True, comment='the greater version of CPE'),
    sa.Column('version_ge', sa.String(length=128), nullable=True, comment='the greater or equal version of CPE'),
    sa.Column('update', sa.String(length=128), nullable=True, comment='the update of CPE'),
    sa.Column('edition', sa.String(length=128), nullable=True, comment='the edition of CPE'),
    sa.Column('language', sa.String(length=128), nullable=True, comment='the language of CPE'),
    sa.Column('sw_edition', sa.String(length=128), nullable=True, comment='the sw_edition of CPE'),
    sa.Column('target_sw', sa.String(length=128), nullable=True, comment='the target_sw of CPE'),
    sa.Column('target_hw', sa.String(length=128), nullable=True, comment='the target_hw of CPE'),
    sa.PrimaryKeyConstraint('id'),
    comment='Table that contains the list of Vulnerabilities'
    )
    op.create_index(op.f('ix_vuln_cpes_vuln_id'), 'vuln_cpes', ['vuln_id'], unique=False)
    op.create_index('vuln_cpe_idx1', 'vuln_cpes', ['product', 'version', 'vendor'], unique=False)
    op.create_unique_constraint('vuln_uix_1', 'vuln', ['vuln_id'])
    op.create_table_comment(
        'vuln',
        'Table that contains the list of Vulnerabilities',
        existing_comment=None,
        schema=None
    )
    # ### end Alembic commands ###

def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_table_comment(
        'vuln',
        existing_comment='Table that contains the list of Vulnerabilities',
        schema=None
    )
    op.drop_constraint('vuln_uix_1', 'vuln', type_='unique')
    op.drop_index('vuln_cpe_idx1', table_name='vuln_cpes')
    op.drop_index(op.f('ix_vuln_cpes_vuln_id'), table_name='vuln_cpes')
    op.drop_table('vuln_cpes')
    # ### end Alembic commands ###


#!/usr/bin/env python3
import requests
import json
from pprint import pprint

# url that returns the json
URL= "http://127.0.0.1:2224/json"
resp= requests.get(URL).json()

print("Original response")

pprint(resp)

# new object to be added to the json/dictionary
newcharacter = {
             "name" : "Jon Snow",
             "house" : "House Stark",
             "age" : 21,
             }

# posting to our endpoint
resp_post = requests.post("http://127.0.0.1:2224/newcharacter", json=json.dumps(newcharacter))

print('------------------------------------------------')
print("New response after the new object is posted")

# pretty print the response for a cleaner and more readeable output
pprint(resp_post.json())


from envs import Arm1DEnv,ExerciseEnv
from rl_agents import SAC

train = False

if train:
# create environemnt
    env = Arm1DEnv(sim_time=3,
                    visualize=False,
                    show_act_plot=False,
                    fixed_init=False,
                    fixed_target=False,
                    integrator_accuracy=1e-5,
                    step_size=1e-2,
                    with_fes=True)

    # create agent
    agent = SAC(env=env,
                mem_size=10000,
                batch_size=500,
                gamma=0.99,
                alpha=1.0,
                dir_name='Arm1DEnv_FES_brach_freq_1_nn_64_32_16_v3\\sac',
                save_rate=50,
                print_rate=1,
                start_lr = 5e-4,
                load_model=True,
                hidden_layers=(64,32,16))

    # train agent
    agent.learn(n_epochs=1000000, verbose=True, pulse_frequency_steps = 1, plot_tensorboard =True)
else:
    env = ExerciseEnv(sim_time=12,
                    visualize=True,
                    show_act_plot=False,
                    fixed_init=True,
                    #fixed_target=False,
                    integrator_accuracy=1e-5,
                    step_size=1e-2,
                    with_fes=True)

    # create agent
    agent = SAC(env=env,
                mem_size=10000,
                batch_size=500,
                gamma=0.99,
                alpha=1.0,
                dir_name='Arm1DEnv_FES_brach_freq_1_nn_64_32_16_v3\\sac',
                save_rate=50,
                print_rate=1,
                start_lr = 1e-4,
                load_model=True,
                hidden_layers=(64,32,16))

    agent.test(n_attemps=100, verbose=True, pulse_frequency_steps = 1)


"""hospital_mangement URL Configuration

The `urlpatterns` list routes URLs to views. For more information please see:
    https://docs.djangoproject.com/en/3.2/topics/http/urls/
Examples:
Function views
    1. Add an import:  from my_app import views
    2. Add a URL to urlpatterns:  path('', views.home, name='home')
Class-based views
    1. Add an import:  from other_app.views import Home
    2. Add a URL to urlpatterns:  path('', Home.as_view(), name='home')
Including another URLconf
    1. Import the include() function: from django.urls import include, path
    2. Add a URL to urlpatterns:  path('blog/', include('blog.urls'))
"""
from django.contrib import admin
from django.urls import path
from hospital import views
# from .views import Patient
urlpatterns = [
    path('admin/', admin.site.urls),
    path('accounts/login/', views.hospital_login),
    path('base/', views.BASE, name="BASE"),

    # --------- main page--------------------------
    # path('welcome/', views.Dashboard, name="Dashboard"),

    # ----------add patients ------------------------
    path('Add-patient/', views.home , name="home_page"),

    # ----------- Show all Patients-------------------
    path('show-paitanse/', views.show_patients, name="show_patients"),
    path('csvcreate-paitanse/', views.P_details.csv_create, name="csv_create"),
    path('excelcreate-paitanse/', views.P_details.excel_create, name="excel_create"),

    #------------ Edit patients --------------------
    path('edit-paitanse/<int:pk>/', views.edit_patients, name="edit_patients"),

    # ------------ delete patients---------------------------
    path('delete-paitanse/<int:pk>/', views.delete_patients, name="delete_patients"),

    #------------------add doctors --------------------------- show_patient
    path('doctor/', views.doctors, name="Add_Doctor"),

#----------------------- show doctors -------------------------------------
    path('all-doctor/', views.doctor_details, name="all_doctor"),

    path('edit-doctors/<int:pk>/', views.edit_doctor, name="edit_doctor"),
    path('delete-doctors/<int:pk>/', views.delete_doctors, name="delete_doctors"),

    path('csvcreate-doctor/', views.Doctor.csv_create, name="csv_create_doctor"),
    path('excelcreate-doctor/', views.Doctor.excel_create, name="excel_create_doctor"),

#-----------------------payment-----------------------------
    path('Payment/', views.Payment, name="Payment"),

#------------------------show payment---------------------------------------
    path('show-payment/', views.show_payment, name="show_payment"),
    path('edit-payment/<int:pk>/', views.edit_payment, name="edit_payment"),
    path('delete-payment/<int:pk>/', views.delete_payment, name="delete_payment"),
    path('csvcreate-payment/', views.Transaction.csv_create, name="csv_create_payment"),
    path('excelcreate-payment/', views.Transaction.excel_create, name="excel_create_payment"),
    path("binomial_graph/", views.binomial_graph, name="binomial_graph"),
    path('payment-pdf/', views.getid, name = 'getid'),

#-------------------------Room--------------------------------------------------
    path('room/', views.room, name="room"),
#-------------------------show room----------------------------------------------------------- Room
    path('show-room/', views.show_room, name="show_room"),
    path('edit-room/<int:pk>/', views.edit_room, name="edit_room"),
    path('delete-room/<int:pk>/', views.delete_room, name="delete_room"),
    path('csvcreate-room/', views.PRoom.csv_create, name="csv_create_Room"),
    path('excelcreate-room/', views.PRoom.excel_create, name="excel_create_Room"),

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

    path('Appointment/', views.appointment, name="Appointment"),
    path('show-appointment/', views.show_appointment, name="show_appointment"),
    path('edit-appointment/<int:pk>/', views.edit_appointment, name="edit_appointment"),
    path('delete-appointment/<int:pk>/', views.delete_appointment, name="delete_appointment"),

    path('csvcreate-apt/', views.PApointment.csv_create, name="csv_create_apointment"),
    path('excelcreate-apt/', views.PApointment.excel_create, name="excel_create_apointment"),

    path('h_form/', views.hospital_form , name = "user_form"),
    path('h_login/', views.hospital_login , name = "user_login"),
    path('logout/',views.logout_user , name = "user_logout"),

]


import csv
import numpy as np
np.random.seed(0)
import pickle,os
from __init__ import ROOT_PATH,datanames
from sklearn.preprocessing import StandardScaler
import collections

def read_ionosphere():
    file = ROOT_PATH+'/data/ionosphere.data'
    def str2int(s):
        return 1 if s is 'g' else -1
    # print('reading Ionosphere')

    with open(file,'r') as rf:
        reader = csv.reader(rf)
        lines = list(reader)
        n = len(lines[0])
        z = [[float(l[i]) if i < n-1 else str2int(l[i]) for i in range(n) if i != 1] for l in lines]
    return np.array(z)

def read_breast_cancer():
    file = ROOT_PATH+'/data/breast-cancer-wisconsin.data'
    def str2int(s):
        return 1 if s is '4' else -1

    # print('reading Cancer')
    with open(file,'r') as rf:
        reader = csv.reader(rf)
        lines = list(reader)
        n = len(lines[0])
        z = [[float(l[i]) if i < n-1 else str2int(l[i]) for i in range(1,n)] for l in lines if '?' not in l]
    return np.array(z)

def read_glass():
    file = ROOT_PATH+'/data/glass.data'

    # print('reading glass')
    with open(file,'r') as rf:
        reader = csv.reader(rf)
        lines = list(reader)
        n = len(lines[0])
        z = [[float(l[i]) if i <n-1 else float(l[i])>4 for i in range(n)] for l in lines if '?' not in l]
    return np.array(z)

def read_sonar():
    file = ROOT_PATH+'/data/sonar.all-data'
    def str2int(s):
        return 1 if s is 'R' else -1
    # print('reading Sonar')
    with open(file,'r') as rf:
        reader = csv.reader(rf)
        lines = list(reader)
        n = len(lines[0])
        z = [[float(l[i]) if i < n-1 else str2int(l[i])  for i in range(n)] for l in lines]
    n = len(z)
    m = len(z[0])
    return np.array(z)

def read_crabs():

    file = ROOT_PATH+'/data/crabs.dat'

    str2int = {'M':1,'F':-1,'B':1,'O':-1}
    # print('reading crabs')
    with open(file,'r') as rf:
        reader = csv.reader(rf)
        lines = list(reader)
        lines = [l[0].split() for l in lines[1:]]
        n = len(lines[0])
        z = [[str2int[l[i]] if i <= 1 else float(l[i])  for i in range(n)] for l in lines]
    m = len(z[0])
    z = np.array(z)
    x = np.array([z[:,i] for i in range(m) if i!= 1]).T
    y = z[:,1]
    z = np.hstack((x,y.reshape((-1,1))))
    return z

def read_pima():
    str2int = {'Yes':1,'No':-1}

    def read_and_remove_na(fileid):
        file = ROOT_PATH+'/data/{}'.format(fileid)
        # print('reading {}'.format(fileid))
        with open(file, 'r') as rf:
            reader = csv.reader(rf)
            lines = list(reader)
            lines = [l[0].split() for l in lines[1:]]
            lines = [l for l in lines if 'NA' not in l]
            n = len(lines[0])
            z = [[str2int[l[i]] if i == n - 1 else float(l[i]) for i in range(n)] for l in lines]
        return z

    z_tr =  read_and_remove_na('pima.tr')
    z_tr2 = read_and_remove_na('pima.tr2')
    z_te = read_and_remove_na('pima.te')

    z = np.vstack((z_tr,z_tr2))
    z = np.vstack((z,z_te))
    return z

def read_wine(l1,l2):
    file = ROOT_PATH + '/data/wine.data'
    # print('reading wine')
    with open(file, 'r') as rf:
        reader = csv.reader(rf)
        lines = list(reader)
        lines = np.array(lines,dtype=float)
        n = len(lines[0])
        col_ids = np.array([i+1 if i != n-1 else 0 for i in range(n)])

        lines = lines[:,col_ids]

        label1, label2 = l1,l2
        lines = np.array([l for l in lines if l[-1] == label1 or l[-1] == label2])
        lines[:, -1] -= (label1 + label2) / 2
        lines[:, -1] /= abs(label1 - label2) / 2
        return lines

def split_data(data,dataname,seed,normalize=False):
    np.random.seed(seed)
    np.random.shuffle(data)
    n = data.shape[0]
    nfolds = 10
    l = int(n / nfolds)

    def loop(i, data, l):
        # split data
        test = data[i * l:i * l + l]
        train = np.vstack((data[:i * l], data[i * l + l:]))
        x_train = train[:, :-1]
        y_train = train[:, -1]
        x_test = test[:, :-1]
        y_test = test[:, -1]
        if normalize:
            scaler = StandardScaler()
            scaler.fit(x_train)
            x_train = scaler.transform(x_train)
            x_test = scaler.transform(x_test)
        dic = {'x_train':x_train, 'y_train':y_train, 'x_test':x_test,'y_test':y_test}
        save_obj(dic,'{}_{}'.format(dataname,i))
    for i in range(nfolds):
        loop(i,data,l)

def save_obj(obj, name):
    with open(ROOT_PATH+'/data/split_data/'+ name + '.pkl', 'wb') as f:
        pickle.dump(obj, f, pickle.HIGHEST_PROTOCOL)
        # print(f)

def generate_data():
    for i in range(100):
        seed = i * 10
        split_data(read_wine(1,3), 'wine13_{}'.format(i), seed,True)
        split_data(read_wine(2, 3), 'wine23_{}'.format(i), seed,True)
        split_data(read_wine(1,2), 'wine12_{}'.format(i), seed,True)

    # download lookup table

if __name__ == '__main__':
    generate_data()


from django import forms
from django.db.models.expressions import RawSQL
from django.utils import timezone

class QualifyingType:
    ORDER_BY = ["-qualifying_score"]

    class Form(forms.Form):
        def __init__(self, qualifying_type, *args, **kwargs):
            submitting = kwargs.pop("submitting", True)
            super().__init__(*args, **kwargs)
            self.qualifying_type = qualifying_type

            if not self.qualifying_type.qualifier.event.vod_required:
                self.fields["vod"].help_text = "(optional)"
                self.fields["vod"].required = False

            # If we're not submitting, ignore vod and details fields
            if not submitting:
                del self.fields["vod"]
                del self.fields["details"]

        vod = forms.URLField(label="VOD")
        details = forms.CharField(widget=forms.Textarea, required=False, label="Details",
                                  help_text="(optional)")

        # Fields to display when submitting a qualifier
        def submit_fields(self):
            for field in ([field for field, _ in self.qualifying_type.EXTRA_FIELDS] +
                          ["vod", "details"]):
                yield self[field]

        # Fields to display when editing a qualifier
        def edit_fields(self):
            for field, _ in self.qualifying_type.EXTRA_FIELDS:
                yield self[field]

        def save(self):
            self.qualifying_type.save_form(self.cleaned_data)

    def __init__(self, qualifier):
        self.qualifier = qualifier
        self.vod = qualifier.vod
        self.details = qualifier.details
        self.load_data(self.qualifier.qualifying_data)

    def display_values(self):
        return ([(label, self.format_value(field, getattr(self, field))) for field, label in self.EXTRA_FIELDS] +
                [("Total score", self.format_score())])

    def format_value(self, field, value):
        return value

    def format_score(self):
        return "{:,}".format(self.qualifier.qualifying_score)

    def form(self, *args):
        submitting = not self.qualifier.submitted
        initial = {
            "vod": self.vod,
            "details": self.details,
        }
        for field, _ in self.EXTRA_FIELDS:
            initial[field] = getattr(self, field)

        return self.Form(self, *args, initial=initial, submitting=submitting)

    def save(self):
        self.qualifier.vod = self.vod
        self.qualifier.details = self.details
        self.qualifier.qualifying_score = self.qualifying_score()
        self.qualifier.qualifying_data = self.qualifying_data()
        self.qualifier.save()

    def save_form(self, cleaned_data):
        submit = not self.qualifier.submitted
        if submit:
            self.qualifier.submitted = True
            self.qualifier.submitted_at = timezone.now()

        for attr, value in cleaned_data.items():
            setattr(self, attr, value)

        self.save()

        if submit:
            self.qualifier.report_submitted()

class HighestScore(QualifyingType):
    NAME = "Highest Score"
    EXTRA_FIELDS = [
        ("score", "Score"),
    ]

    class Form(QualifyingType.Form):
        score = forms.IntegerField(min_value=0, label="Score")

    def load_data(self, data):
        self.score = data[0] if data else None

    def qualifying_score(self):
        return self.score

    def qualifying_data(self):
        return [self.score]

class Highest2Scores(QualifyingType):
    NAME = "Highest 2 Scores"
    EXTRA_FIELDS = [
        ("score1", "Score 1"),
        ("score2", "Score 2"),
    ]

    class Form(QualifyingType.Form):
        score1 = forms.IntegerField(min_value=0, label="Score 1")
        score2 = forms.IntegerField(min_value=0, label="Score 2")

    def load_data(self, data):
        self.score1 = data[0] if data else None
        self.score2 = data[1] if data else None

    def qualifying_score(self):
        return (self.score1 + self.score2) // 2

    def qualifying_data(self):
        return list(reversed(sorted([self.score1, self.score2])))

class Highest3Scores(QualifyingType):
    NAME = "Highest 3 Scores"
    EXTRA_FIELDS = [
        ("score1", "Score 1"),
        ("score2", "Score 2"),
        ("score3", "Score 3"),
    ]

    class Form(QualifyingType.Form):
        score1 = forms.IntegerField(min_value=0, label="Score 1")
        score2 = forms.IntegerField(min_value=0, label="Score 2")
        score3 = forms.IntegerField(min_value=0, label="Score 3")

    def load_data(self, data):
        self.score1 = data[0] if data else None
        self.score2 = data[1] if data else None
        self.score3 = data[2] if data else None

    def qualifying_score(self):
        return (self.score1 + self.score2 + self.score3) // 3

    def qualifying_data(self):
        return list(reversed(sorted([self.score1, self.score2, self.score3])))

class MostMaxouts(QualifyingType):
    NAME = "Most Maxouts"
    EXTRA_FIELDS = [
        ("maxouts", "Maxout Count"),
        ("kicker", "Kicker"),
    ]
    ORDER_BY = [
        RawSQL("(qualifying_data::json->>'maxouts')::integer", ()).desc(),
        RawSQL("(qualifying_data::json->>'kicker')::integer", ()).desc()
    ]

    class Form(QualifyingType.Form):
        maxouts = forms.IntegerField(
            min_value=0, label="Maxout Count",
            help_text=("Total number of maxouts (scores over 999,999) you got. If you did not "
                       "maxout, enter 0.")
        )
        kicker = forms.IntegerField(
            min_value=0, label="Kicker",
            help_text=("Your highest non-maxout score. If you did not maxout, this is your highest "
                       "score.")
        )

    def load_data(self, data):
        self.maxouts = data["maxouts"] if data else None
        self.kicker = data["kicker"] if data else None

    def qualifying_score(self):
        return self.maxouts

    def format_score(self):
        return "{}; {:,}".format(self.maxouts, self.kicker)

    def qualifying_data(self):
        return {
            "maxouts": self.maxouts,
            "kicker": self.kicker,
        }

def format_time(time):
    minutes = int(time // 60)
    seconds = time % 60
    return "{}:{:06.3f}".format(minutes, seconds)

class LowestTime(QualifyingType):
    NAME = "Lowest Time"
    EXTRA_FIELDS = [
        ("time", "Time"),
    ]
    ORDER_BY = ["qualifying_score"] # increasing; lowest time comes first

    class Form(QualifyingType.Form):
        time = forms.DurationField(label="Time", help_text="(e.g. 12:34.567)")

        def __init__(self, *args, **kwargs):
            if kwargs["initial"]["time"]:
                kwargs["initial"]["time"] = format_time(kwargs["initial"]["time"])
            super().__init__(*args, **kwargs)

        def clean_time(self):
            time = self.cleaned_data["time"]
            return time.total_seconds()

    def load_data(self, data):
        self.time = data[0] if data else None

    def qualifying_score(self):
        return round(self.time * 1000)

    def format_value(self, field, value):
        if field == "time":
            return format_time(value)
        else:
            return value

    def format_score(self):
        return format_time(self.time)

    def qualifying_data(self):
        return [self.time]

QUALIFYING_TYPES = {
    1: HighestScore,
    2: Highest2Scores,
    3: Highest3Scores,
    4: MostMaxouts,
    5: LowestTime,
}
CHOICES = [(n, qualifying_type.NAME) for n, qualifying_type in QUALIFYING_TYPES.items()]


import subprocess

from .utils import CommandNotFoundError, T, shjoin

class Local:
    def display(self, args):
        print(T.bold_green(f"(local) $ ", shjoin(args)))

    def silent_get(self, *args, **kwargs):
        return subprocess.check_output(
            args,
            universal_newlines=True,
            **kwargs,
        )

    def get(self, *args, **kwargs):
        self.display(args)
        return subprocess.check_output(
            args,
            universal_newlines=True,
            **kwargs,
        )

    def run(self, *args, **kwargs):
        self.display(args)
        try:
            return subprocess.run(
                args,
                universal_newlines=True,
                **kwargs,
            )
        except FileNotFoundError as e:
            if e.filename == args[0]:
                raise CommandNotFoundError(e.filename)
            else:
                raise

    def popen(self, *args, **kwargs):
        self.display(args)
        return subprocess.Popen(
            args,
            universal_newlines=True,
            **kwargs,
        )

    def check_passwordless(self, host):
        results = self.run(
            "ssh",
            "-oPreferredAuthentications=publickey",
            host,
            "echo OK",
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )
        if results.returncode != 0:
            if "Permission denied" in results.stderr:
                return False
            else:
                print(results.stdout)
                print(results.stderr)
                exit(f"Failed to connect to {host}, could not understand error")
        else:
            print("# OK")
            return True


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

"""
This file is to create patterned image. The input is needed from coord_generator.py .

Data needed:
    prefix = 'm%02df%03d' % (idx_set[idx][0], idx_set[idx][1])
    1. xy_pro1_cv.npy
    2. mask_cam
    3. pattern.png

Main work:
    1. Set jump value for image pair calculation.
        The jump value is used for frame calculation. Image pair will be calculated with i and i+jump frame.
        Default: 1
    2. Load needed information. Including:
        xy_pro1_cv.npy
        mask_cam.png
        pattern.png
    3. Draw pattern by using grid_sample.
        (https://pytorch.org/docs/stable/nn.html?highlight=grid_sample#torch.nn.functional.grid_sample)
        sample pattern.png, by using xy_pro1_cv.npy.
        Get: image_cam.png
    4. Save information
        <image_cam>_1.png
        <image_cam>_2.png

"""

import csv
import sys
import os
import torch
from matplotlib import pyplot as plt
import numpy as np

def load_as_torch(name, suffix, main_path, prefix, channel=1):
    """Load file as torch tensor.

    Possible data:
        depth_cam.npy   [H, W], float
        mask_cam.png    [H, W], byte
        xy_pro1_cv.npy  [H, W, 2], float
        xy_cam_p1v.npy  [H, W, 2], float
        mask_pro.png    [H, W], byte

    Args:
        :param name: the data name.
        :param suffix: '.png' or '.npy'.
        :param main_path: the file path with all the m00f000 data.
        :param prefix: m00f000.

    Return:
        :return: Loaded data or None(if suffix is not available).
                Data is always 4D [N, C, H, W].

    Raises:
        None.
    """
    full_name = '%s/%s_%s%s' % (main_path, prefix, name, suffix)
    if suffix == '.png':
        item = torch.from_numpy(plt.imread(full_name))
        if channel == 1:
            item = item[:, :, 0].unsqueeze(0)
        else:
            item = item.permute(2, 0, 1)
        item = item.unsqueeze(0)
        item = item.byte()
        return item.cuda()
    elif suffix == '.npy':
        item = torch.from_numpy(np.load(full_name).astype(float))
        if len(item.shape) == 2:
            item = item.unsqueeze(0)
        else:
            item = item.permute(2, 0, 1)
        item[torch.isnan(item)] = 0
        item = item.unsqueeze(0)
        return item.cuda()
    else:
        return None

def save_from_torch(mat, name, suffix, out_path, prefix):
    """Save file from torch tensor to npy or png file.

    detach, cpu, numpy. use plt.imsave and np.save function.

    :param name:
    :param suffix:
    :param out_path:
    :param prefix:
    :return:
    """
    full_name = '%s/%s_%s%s' % (out_path, prefix, name, suffix)
    if suffix == '.png':
        # mat_np = mat.detach().cpu().squeeze().numpy()
        mat_np = mat.detach().cpu().squeeze().permute(1, 2, 0).numpy()
        plt.imsave(full_name, mat_np)
    elif suffix == '.npy':
        mat = mat.detach().cpu().squeeze()
        if len(mat.shape) == 3:
            mat = mat.permute(1, 2, 0)
        mat_np = mat.numpy().astype(np.float32)
        np.save(full_name, mat_np)

def main():

    # Some parameters (You should edit here every time):
    main_path = '/media/qiao/数据文档/SLDataSet/Thing10K/test_dataset'
    out_path = '/media/qiao/数据文档/SLDataSet/Thing10K/4pix_dataset/test'
    pattern_path = '/media/qiao/数据文档/SLDataSet/Thing10K/'
    model_num = 2
    frame_num = 100
    jump = 1
    image_shape = (1024, 1280)

    # Main Loop
    if not os.path.exists(out_path):
        os.mkdir(out_path)
    # x_cam_grid = torch.arange(0, image_shape[1]).reshape(1, -1).repeat(image_shape[0], 1)
    # y_cam_grid = torch.arange(0, image_shape[0]).reshape(-1, 1).repeat(1, image_shape[1])
    # xy_cam_grid = torch.stack((x_cam_grid, y_cam_grid), dim=0).reshape(1, 2, image_shape[0], image_shape[1]).float()
    # xy_cam_grid = xy_cam_grid.cuda()
    for m_idx in range(1, model_num + 1):
        # Step 1: Load all information needed.
        print("Loading No.%02d model..." % m_idx, end='', flush=True)
        mask_cam = []
        xy_pro1_cv = []
        pattern = load_as_torch('pattern', '.png', pattern_path, '4pix', channel=3)
        for f_idx in range(0, frame_num):
            prefix = 'm%02df%03d' % (m_idx, f_idx)
            mask_cam.append(load_as_torch('mask_cam', '.png', main_path, prefix))
            xy_pro1_cv.append(load_as_torch('xy_pro1_cv', '.npy', main_path, prefix))
        print("Finished.")

        # Step 2: Process every frame, calculate image_cam.
        print("Calculate No.%02d image_cam..." % m_idx, end='', flush=True)
        image_cam = []
        for f_idx in range(0, frame_num):
            i = f_idx
            # xy_pro1_cv to [-1, 1]
            xy_pro1_cv[i][:, 0, :, :] = xy_pro1_cv[i][:, 0, :, :] / (image_shape[1] - 1) * 2.0 - 1.0
            xy_pro1_cv[i][:, 1, :, :] = xy_pro1_cv[i][:, 1, :, :] / (image_shape[0] - 1) * 2.0 - 1.0
            xy_pro1_cv[i][mask_cam[i].repeat(1, 2, 1, 1) == 0] = -1
            xy_pro1_cv[i] = xy_pro1_cv[i].permute(0, 2, 3, 1)  # [N(1), H, W, 2]
            # Sample from xy_cam_p1v
            image = torch.nn.functional.grid_sample(input=pattern.float(), grid=xy_pro1_cv[i].float())
            image[mask_cam[i].repeat(1, 3, 1, 1) == 0] = 0
            image_cam.append(image)
        print("Finished.")

        # Step 3: Save image as image pairs
        print("Save No.%02d image pairs..." % m_idx, end='', flush=True)
        for f_idx in range(0, frame_num):
            prefix = 'm%02df%03d' % (m_idx, f_idx)
            i = f_idx
            j = i + jump
            if j >= frame_num:
                j -= frame_num
            save_from_torch(image_cam[i], 'img_cam_img1', '.png', out_path, prefix)
            save_from_torch(image_cam[j], 'img_cam_img2', '.png', out_path, prefix)
        print("Finished.")

if __name__ == '__main__':
    assert len(sys.argv) >= 1

    main()


# pylint: disable=C0413, W0212
import os
import sys
import unittest
import numpy as np
from pathlib import Path

sys.path.insert(0, os.getcwd())
from legacy_data_pipeline.pan_motor_timebase_aligner import PanMotorTimeBaseAligner
from legacy_data_pipeline.legacy_data_file_system_helper import LegacyDataFileSystemHelper as LDFH
from video_and_photo_tools.image_from_video import ImageFromVideo

class TestPanMoterTimebaseAligner(unittest.TestCase):

    HEAD_PATH = Path('/Volumes/WD')
    TEST_SESSION_DIR_NAME = 'Aug_17_Palo_Alto_High_2nd_time_B80_ottofillmore'

    def setUp(self):
        self.ldfh = LDFH(self.__class__.HEAD_PATH)
        self.pmta = PanMotorTimeBaseAligner(self.__class__.TEST_SESSION_DIR_NAME,
                                            num_points=10)

    def dont_test_analyze_motor_data(self):
        data = self.pmta.get_motor_data()
        differences = np.ediff1d(data)
        sorted_differences = np.sort(np.abs(differences))[::-1]

    def test_visualize_peaks(self):
        self.pmta.visualize_motor_local_maxima_and_minima()

    def dont_test_show_scrubber_for_maxima(self):
        self.pmta.show_scrub_picker_for_maxima()

    def dont_test_show_scrubber_for_minima(self):
        self.pmta.show_scrub_picker_for_minima()

    def dont_test_get_images_around_time(self):
        times = self.pmta._get_time_at_maxima()
        images_around_times = [self.pmta.get_images_around_time(time)
                               for time in times]
        for images in images_around_times:
            for image in images:
                self.assertIsInstance(image, ImageFromVideo)

    def test_run(self):
        self.pmta.show_scrub_picker_for_min_and_max()
        self.pmta.save_results()

if __name__ == '__main__':
    unittest.main()


import pandas as pd
from data.process_data.process_data import process_data

DUNNHUMBY_PATH = 'data/dunnhumby/dunnhumby.csv'
DUNNHUMBY_OUTPUT = 'dunnhumby_all'
DUNNHUMBY_USER_COL = 'household_key'
DUNNHUMBY_TIME_COL = 'DAY'
DUNNHUMBY_ITEM_COL = 'PRODUCT_ID'
DUNNHUMBY_MIN_ITEM_CNT = 10
DUNNHUMBY_MIN_USER_ITEM_CNT = 10
DUNNHUMBY_MAX_USER_ITEM_CNT = 10000
DUNNHUMBY_MIN_USER_SET_CNT = 4
DUNNHUMBY_MAX_USER_SET_CNT = 1000

def main():
    df = pd.read_csv(DUNNHUMBY_PATH)
    print(df.head())
    df = df[(df.DAY >= 1) & (df.DAY <= 300)]
    process_data(df=df,
                 output=DUNNHUMBY_OUTPUT,
                 user_col=DUNNHUMBY_USER_COL,
                 time_col=DUNNHUMBY_TIME_COL,
                 item_col=DUNNHUMBY_ITEM_COL,
                 min_item_cnt=DUNNHUMBY_MIN_ITEM_CNT,
                 min_user_item_cnt=DUNNHUMBY_MIN_USER_ITEM_CNT,
                 max_user_item_cnt=DUNNHUMBY_MAX_USER_ITEM_CNT,
                 min_user_set_cnt=DUNNHUMBY_MIN_USER_SET_CNT,
                 max_user_set_cnt=DUNNHUMBY_MAX_USER_SET_CNT)

if __name__ == '__main__':
    main()


import streamlit as st
import os
import glob
from dotenv import load_dotenv
from PyPDF2 import PdfReader
from langchain.text_splitter import CharacterTextSplitter
from langchain.embeddings import OpenAIEmbeddings
from langchain.vectorstores import FAISS
from langchain.chat_models import ChatOpenAI
from langchain.memory import ConversationBufferMemory
from langchain.chains import ConversationalRetrievalChain
from htmlTemplates import css, bot_template, user_template

def write_to_file(question, response, file_path="chat_history.txt"):
    with open(file_path, "a") as file:
        file.write(f"Question: {question}\nResponse: {response}\n\n")

def get_pdf_text(pdf_docs):
    text = ""
    for pdf in pdf_docs:
        pdf_reader = PdfReader(pdf)
        for page in pdf_reader.pages:
            text += page.extract_text()
    return text

def get_text_chunks(text):
    text_splitter = CharacterTextSplitter(
        separator="\n",
        chunk_size=1000,
        chunk_overlap=200,
        length_function=len
    )
    chunks = text_splitter.split_text(text)
    return chunks

def get_vectorstore(text_chunks):
    embeddings = OpenAIEmbeddings()
    vectorstore = FAISS.from_texts(texts=text_chunks, embedding=embeddings)
    return vectorstore

def get_conversation_chain(vectorstore):
    llm = ChatOpenAI()
    memory = ConversationBufferMemory(
        memory_key='chat_history', return_messages=True)
    conversation_chain = ConversationalRetrievalChain.from_llm(
        llm=llm,
        retriever=vectorstore.as_retriever(),
        memory=memory
    )
    return conversation_chain

def handle_userinput(user_question):
    response = st.session_state.conversation({'question': user_question})
    st.session_state.chat_history = response['chat_history']

    for i, message in enumerate(st.session_state.chat_history):
        if i % 2 == 0:
            user_message = message.content
            st.write(user_template.replace("{{MSG}}", user_message), unsafe_allow_html=True)
        else:
            bot_message = message.content
            st.write(bot_template.replace("{{MSG}}", bot_message), unsafe_allow_html=True)
            write_to_file(user_message, bot_message)  # Write to file after each exchange

def main():
    load_dotenv()
    st.set_page_config(page_title="Macroeconomic Researcher and Large Language Chat GPT -by TechTitans",
                       page_icon="")
    st.write(css, unsafe_allow_html=True)

    if "conversation" not in st.session_state:
        st.session_state.conversation = None
    if "chat_history" not in st.session_state:
        st.session_state.chat_history = None

    st.header("Macroeconomic Researcher and Large Language Chat GPT -by TechTitans")
    user_question = st.text_input("Ask your questions :")
    if user_question:
        handle_userinput(user_question)

    local_directory = "PDFS"  # Replace with your directory path

    pdf_files = glob.glob(os.path.join(local_directory, "*.pdf"))
    if pdf_files:
        with st.spinner("Processing your question..."):
            pdf_docs = [open(file, 'rb') for file in pdf_files]
            # get pdf text
            raw_text = get_pdf_text(pdf_docs)

            # Close the opened PDF files
            for pdf in pdf_docs:
                pdf.close()

            # get the text chunks
            text_chunks = get_text_chunks(raw_text)

            # create vector store
            vectorstore = get_vectorstore(text_chunks)

            # create conversation chain
            st.session_state.conversation = get_conversation_chain(vectorstore)

if __name__ == '__main__':
    main()


"""
This module provides `TIM`, PyTermGUI's Terminal Inline Markup language. It is a simple,
performant and easy to read way to style, colorize & modify text.

Basic rundown
-------------

TIM is included with the purpose of making styling easier to read and manage.

Its syntax is based on square brackets, within which tags are strictly separated by one
space character. Tags can stand for colors (xterm-256, RGB or HEX, both background &
foreground), styles, unsetters and macros.

The 16 simple colors of the terminal exist as named tags that refer to their numerical
value.

Here is a simple example of the syntax, using the `pytermgui.pretty` submodule to
syntax-highlight it inside the REPL:

```python3
>>> from pytermgui import pretty
>>> '[141 @61 bold] Hello [!upper inverse] There '
```

General syntax
--------------

Background colors are always denoted by a leading `@` character in front of the color
tag. Styles are just the name of the style and macros have an exclamation mark in front
of them. Additionally, unsetters use a leading slash (`/`) for their syntax. Color
tokens have special unsetters: they use `/fg` to cancel foreground colors, and `/bg` to
do so with backgrounds.

### Macros:

Macros are any type of callable that take at least *args; this is the value of the plain
text enclosed by the tag group within which the given macro resides. Additionally,
macros can be given any number of positional arguments from within markup, using the
syntax:

```
[!macro(arg1:arg2:arg3)]Text that the macro applies to.[/!macro]plain text, no macro
```

This syntax gets parsed as follows:

```python3
macro("Text that the macro applies to.", "arg1", "arg2", "arg3")
```

`macro` here is whatever the name `macro` was defined as prior.

### Colors:

Colors can be of three general types: xterm-256, RGB and HEX.

`xterm-256` stands for one of the 256 xterm colors. You can use `ptg -c` to see the all
of the available colors. Its syntax is just the 0-base index of the color, like `[141]`

`RGB` colors are pretty self explanatory. Their syntax is follows the format
`RED;GREEN;BLUE`, such as `[111;222;333]`.

`HEX` colors are basically just RGB with extra steps. Their syntax is `#RRGGBB`, such as
`[#FA72BF]`. This code then gets converted to a tuple of RGB colors under the hood, so
from then on RGB and HEX colors are treated the same, and emit the same tokens.

As mentioned above, all colors can be made to act on the background instead by
prepending the color tag with `@`, such as `@141`, `@111;222;333` or `@#FA72BF`. To
clear these effects, use `/fg` for foreground and `/bg` for background colors.

`MarkupLanguage` and instancing
-------------------------------

All markup behaviour is done by an instance of the `MarkupLanguage` class. This is done
partially for organization reasons, but also to allow a sort of sandboxing of custom
definitions and settings.

PyTermGUI provides the `tim` name as the global markup language instance. For historical
reasons, the same instance is available as `markup`. This should be used pretty much all
of the time, and custom instances should only ever come about when some
security-sensitive macro definitions are needed, as `markup` is used by every widget,
including user-input ones such as `InputField`.

For the rest of this page, `MarkupLanguage` will refer to whichever instance you are
using.

TL;DR : Use `tim` always, unless a security concern blocks you from doing so.

Caching
-------

By default, all markup parse results are cached and returned when the same input is
given. To disable this behaviour, set your markup instance (usually `markup`)'s
`should_cache` field to False.

Customization
-------------

There are a couple of ways to customize how markup is parsed. Custom tags can be created
by calling `MarkupLanguage.alias`. For defining custom macros, you can use
`MarkupLanguage.define`. For more information, see each method's documentation.
"""
# pylint: disable=too-many-lines

from __future__ import annotations

from argparse import ArgumentParser
from contextlib import suppress
from dataclasses import dataclass
from enum import Enum
from enum import auto as _auto
from functools import cached_property
from random import shuffle
from typing import Callable, Iterator, List, Tuple

from .colors import Color, StandardColor, str_to_color
from .exceptions import AnsiSyntaxError, ColorSyntaxError, MarkupSyntaxError
from .regex import RE_ANSI, RE_LINK, RE_MACRO, RE_MARKUP
from .terminal import get_terminal

__all__ = [
    "StyledText",
    "MacroCallable",
    "MacroCall",
    "MarkupLanguage",
    "markup",
    "tim",
]

MacroCallable = Callable[..., str]
MacroCall = Tuple[MacroCallable, List[str]]

STYLE_MAP = {
    "bold": "1",
    "dim": "2",
    "italic": "3",
    "underline": "4",
    "blink": "5",
    "blink2": "6",
    "inverse": "7",
    "invisible": "8",
    "strikethrough": "9",
    "overline": "53",
}

UNSETTER_MAP: dict[str, str | None] = {
    "/": "0",
    "/bold": "22",
    "/dim": "22",
    "/italic": "23",
    "/underline": "24",
    "/blink": "25",
    "/blink2": "26",
    "/inverse": "27",
    "/invisible": "28",
    "/strikethrough": "29",
    "/fg": "39",
    "/bg": "49",
    "/overline": "54",
}

def macro_align(width: str, alignment: str, content: str) -> str:
    """Aligns given text using fstrings.

    Args:
        width: The width to align to.
        alignment: One of "left", "center", "right".
        content: The content to align; implicit argument.
    """

    aligner = "<" if alignment == "left" else (">" if alignment == "right" else "^")
    return f"{content:{aligner}{width}}"

def macro_expand(lang: MarkupLanguage, tag: str) -> str:
    """Expands a tag alias."""

    if not tag in lang.user_tags:
        return tag

    return lang.get_markup(f"\x1b[{lang.user_tags[tag]}m ")[:-1]

def macro_strip_fg(item: str) -> str:
    """Strips foreground color from item"""

    return markup.parse(f"[/fg]{item}")

def macro_strip_bg(item: str) -> str:
    """Strips foreground color from item"""

    return markup.parse(f"[/bg]{item}")

def macro_shuffle(item: str) -> str:
    """Shuffles a string using shuffle.shuffle on its list cast."""

    shuffled = list(item)
    shuffle(shuffled)

    return "".join(shuffled)

def macro_link(*args) -> str:
    """Creates a clickable hyperlink.

    Note:
        Since this is a pretty new feature for terminals, its support is limited.
    """

    *uri_parts, label = args
    uri = ":".join(uri_parts)

    return f"\x1b]8;;{uri}\x1b\\{label}\x1b]8;;\x1b\\"

def _apply_colors(colors: list[str] | list[int], item: str) -> str:
    """Applies the given list of colors to the item, spread out evenly."""

    blocksize = max(round(len(item) / len(colors)), 1)

    out = ""
    current_block = 0
    for i, char in enumerate(item):
        if i % blocksize == 0 and current_block < len(colors):
            out += f"[{colors[current_block]}]"
            current_block += 1

        out += char

    return markup.parse(out)

def macro_rainbow(item: str) -> str:
    """Creates rainbow-colored text."""

    colors = ["red", "208", "yellow", "green", "bright-blue", "blue", "93"]

    return _apply_colors(colors, item)

def macro_gradient(base_str: str, item: str) -> str:
    """Creates an xterm-256 gradient from a base color.

    This exploits the way the colors are arranged in the xterm color table; every
    36th color is the next item of a single gradient.

    The start of this given gradient is calculated by decreasing the given base by 36 on
    every iteration as long as the point is a valid gradient start.

    After that, the 6 colors of this gradient are calculated and applied.
    """

    if not base_str.isdigit():
        raise ValueError(f"Gradient base has to be a digit, got {base_str}.")

    base = int(base_str)
    if base < 16 or base > 231:
        raise ValueError("Gradient base must be between 16 and 232")

    while base > 52:
        base -= 36

    colors = []
    for i in range(6):
        colors.append(base + 36 * i)

    return _apply_colors(colors, item)

class TokenType(Enum):
    """An Enum to store various token types."""

    LINK = _auto()
    """A terminal hyperlink."""

    PLAIN = _auto()
    """Plain text, nothing interesting."""

    COLOR = _auto()
    """A color token. Has a `pytermgui.colors.Color` instance as its data."""

    STYLE = _auto()
    """A builtin terminal style, such as `bold` or `italic`."""

    MACRO = _auto()
    """A PTG markup macro. The macro itself is stored inside `self.data`."""

    ESCAPED = _auto()
    """An escaped token."""

    UNSETTER = _auto()
    """A token that unsets some other attribute."""

    POSITION = _auto()
    """A token representing a positioning string. `self.data` follows the format `x,y`."""

@dataclass
class Token:
    """A class holding information on a singular markup or ANSI style unit.

    Attributes:
    """

    ttype: TokenType
    """The type of this token."""

    data: str | MacroCall | Color | None
    """The data contained within this token. This changes based on the `ttype` attr."""

    name: str = "<unnamed-token>"
    """An optional display name of the token. Defaults to `data` when not given."""

    def __post_init__(self) -> None:
        """Sets `name` to `data` if not provided."""

        if self.name == "<unnamed-token>":
            if isinstance(self.data, str):
                self.name = self.data

            elif isinstance(self.data, Color):
                self.name = self.data.name

            else:
                raise TypeError

        # Create LINK from a plain token
        if self.ttype is TokenType.PLAIN:
            assert isinstance(self.data, str)

            link_match = RE_LINK.match(self.data)

            if link_match is not None:
                self.data, self.name = link_match.groups()
                self.ttype = TokenType.LINK

        if self.ttype is TokenType.ESCAPED:
            assert isinstance(self.data, str)

            self.name = self.data[1:]

    def __eq__(self, other: object) -> bool:
        """Checks equality with `other`."""

        if other is None:
            return False

        if not isinstance(other, type(self)):
            return False

        return other.data == self.data and other.ttype is self.ttype

    @cached_property
    def sequence(self) -> str | None:
        """Returns the ANSI sequence this token represents."""

        if self.data is None:
            return None

        if self.ttype in [TokenType.PLAIN, TokenType.MACRO, TokenType.ESCAPED]:
            return None

        if self.ttype is TokenType.LINK:
            return macro_link(self.data, self.name)

        if self.ttype is TokenType.POSITION:
            assert isinstance(self.data, str)
            position = self.data.split(",")
            return f"\x1b[{position[1]};{position[0]}H"

        # Colors and styles
        data = self.data

        if self.ttype in [TokenType.STYLE, TokenType.UNSETTER]:
            return f"\033[{data}m"

        assert isinstance(data, Color)
        return data.sequence

class StyledText(str):
    """A styled text object.

    The purpose of this class is to implement some things regular `str`
    breaks at when encountering ANSI sequences.

    Instances of this class are usually spat out by `MarkupLanguage.parse`,
    but may be manually constructed if the need arises. Everything works even
    if there is no ANSI tomfoolery going on.
    """

    value: str
    """The underlying, ANSI-inclusive string value."""

    _plain: str | None = None
    _tokens: list[Token] | None = None

    def __new__(cls, value: str = ""):
        """Creates a StyledText, gets markup tags."""

        obj = super().__new__(cls, value)
        obj.value = value

        return obj

    def _generate_tokens(self) -> None:
        """Generates self._tokens & self._plain."""

        self._tokens = list(tim.tokenize_ansi(self.value))

        self._plain = ""
        for token in self._tokens:
            if token.ttype is not TokenType.PLAIN:
                continue

            assert isinstance(token.data, str)
            self._plain += token.data

    @property
    def tokens(self) -> list[Token]:
        """Returns all markup tokens of this object.

        Generated on-demand, at the first call to this or the self.plain
        property.
        """

        if self._tokens is not None:
            return self._tokens

        self._generate_tokens()
        assert self._tokens is not None
        return self._tokens

    @property
    def plain(self) -> str:
        """Returns the value of this object, with no ANSI sequences.

        Generated on-demand, at the first call to this or the self.tokens
        property.
        """

        if self._plain is not None:
            return self._plain

        self._generate_tokens()
        assert self._plain is not None
        return self._plain

    def plain_index(self, index: int | None) -> int | None:
        """Finds given index inside plain text."""

        if index is None:
            return None

        styled_chars = 0
        plain_chars = 0
        negative_index = False

        tokens = self.tokens.copy()
        if index < 0:
            tokens.reverse()
            index = abs(index)
            negative_index = True

        for token in tokens:
            if token.data is None:
                continue

            if token.ttype is not TokenType.PLAIN:
                assert token.sequence is not None
                styled_chars += len(token.sequence)
                continue

            assert isinstance(token.data, str)
            for _ in range(len(token.data)):
                if plain_chars == index:
                    if negative_index:
                        return -1 * (plain_chars + styled_chars)

                    return styled_chars + plain_chars

                plain_chars += 1

        return None

    def __len__(self) -> int:
        """Gets "real" length of object."""

        return len(self.plain)

    def __getitem__(self, subscript: int | slice) -> str:
        """Gets an item, adjusted for non-plain text.

        Args:
            subscript: The integer or slice to find.

        Returns:
            The elements described by the subscript.

        Raises:
            IndexError: The given index is out of range.
        """

        if isinstance(subscript, int):
            plain_index = self.plain_index(subscript)
            if plain_index is None:
                raise IndexError("StyledText index out of range")

            return self.value[plain_index]

        return self.value[
            slice(
                self.plain_index(subscript.start),
                self.plain_index(subscript.stop),
                subscript.step,
            )
        ]

class MarkupLanguage:
    """A class representing an instance of a Markup Language.

    This class is used for all markup/ANSI parsing, tokenizing and usage.

    ```python3
    from pytermgui import tim

    tim.alias("my-tag", "@152 72 bold")
    tim.print("This is [my-tag]my-tag[/]!")
    ```

    """

    raise_unknown_markup: bool = False
    """Raise `pytermgui.exceptions.MarkupSyntaxError` when encountering unknown markup tags."""

    def __init__(self, default_macros: bool = True) -> None:
        """Initializes a MarkupLanguage.

        Args:
            default_macros: If not set, the builtin macros are not defined.
        """

        self.tags: dict[str, str] = STYLE_MAP.copy()
        self._cache: dict[str, StyledText] = {}
        self.macros: dict[str, MacroCallable] = {}
        self.user_tags: dict[str, str] = {}
        self.unsetters: dict[str, str | None] = UNSETTER_MAP.copy()

        self.should_cache: bool = True

        if default_macros:
            self.define("!link", macro_link)
            self.define("!align", macro_align)
            self.define("!markup", self.get_markup)
            self.define("!shuffle", macro_shuffle)
            self.define("!strip_bg", macro_strip_bg)
            self.define("!strip_fg", macro_strip_fg)
            self.define("!rainbow", macro_rainbow)
            self.define("!gradient", macro_gradient)
            self.define("!upper", lambda item: str(item.upper()))
            self.define("!lower", lambda item: str(item.lower()))
            self.define("!title", lambda item: str(item.title()))
            self.define("!capitalize", lambda item: str(item.capitalize()))
            self.define("!expand", lambda tag: macro_expand(self, tag))
            self.define("!debug", lambda *args: ",".join(ascii(arg) for arg in args))

        self.alias("code", "dim @black")
        self.alias("code.str", "142")
        self.alias("code.multiline_str", "code.str")
        self.alias("code.none", "167")
        self.alias("code.global", "214")
        self.alias("code.number", "175")
        self.alias("code.keyword", "203")
        self.alias("code.identifier", "109")
        self.alias("code.name", "code.global")
        self.alias("code.comment", "240 italic")
        self.alias("code.builtin", "code.global")
        self.alias("code.file", "code.identifier")
        self.alias("code.symbol", "code.identifier")

    def _get_color_token(self, tag: str) -> Token | None:
        """Tries to get a color token from the given tag.

        Args:
            tag: The tag to parse.

        Returns:
            A color token if the given tag could be parsed into one, else None.
        """

        try:
            color = str_to_color(tag, use_cache=self.should_cache)

        except ColorSyntaxError:
            return None

        return Token(name=color.value, ttype=TokenType.COLOR, data=color)

    def _get_style_token(self, tag: str) -> Token | None:
        """Tries to get a style (including unsetter) token from tags, user tags and unsetters.

        Args:
            tag: The tag to parse.

        Returns:
            A `Token` if one could be created, None otherwise.
        """

        if tag in self.unsetters:
            return Token(name=tag, ttype=TokenType.UNSETTER, data=self.unsetters[tag])

        if tag in self.user_tags:
            return Token(name=tag, ttype=TokenType.STYLE, data=self.user_tags[tag])

        if tag in self.tags:
            return Token(name=tag, ttype=TokenType.STYLE, data=self.tags[tag])

        return None

    def print(self, *args, **kwargs) -> None:
        """Parse all arguments and pass them through to print, along with kwargs."""

        parsed = []
        for arg in args:
            parsed.append(self.parse(str(arg)))

        get_terminal().print(*parsed, **kwargs)

    def tokenize_markup(self, markup_text: str) -> Iterator[Token]:
        """Converts the given markup string into an iterator of `Token`.

        Args:
            markup_text: The text to look at.

        Returns:
            An iterator of tokens. The reason this is an iterator is to possibly save
            on memory.
        """

        end = 0
        start = 0
        cursor = 0
        for match in RE_MARKUP.finditer(markup_text):
            full, escapes, tag_text = match.groups()
            start, end = match.span()

            # Add plain text between last and current match
            if start > cursor:
                yield Token(ttype=TokenType.PLAIN, data=markup_text[cursor:start])

            if not escapes == "" and len(escapes) % 2 == 1:
                cursor = end
                yield Token(ttype=TokenType.ESCAPED, data=full[len(escapes) - 1 :])
                continue

            for tag in tag_text.split():
                token = self._get_style_token(tag)
                if token is not None:
                    yield token
                    continue

                # Try to find a color token
                token = self._get_color_token(tag)
                if token is not None:
                    yield token
                    continue

                macro_match = RE_MACRO.match(tag)
                if macro_match is not None:
                    name, args = macro_match.groups()
                    macro_args = () if args is None else args.split(":")

                    if not name in self.macros:
                        raise MarkupSyntaxError(
                            tag=tag,
                            cause="is not a defined macro",
                            context=markup_text,
                        )

                    yield Token(
                        name=tag,
                        ttype=TokenType.MACRO,
                        data=(self.macros[name], macro_args),
                    )
                    continue

                if self.raise_unknown_markup:
                    raise MarkupSyntaxError(
                        tag=tag, cause="not defined", context=markup_text
                    )

            cursor = end

        # Add remaining text as plain
        if len(markup_text) > cursor:
            yield Token(ttype=TokenType.PLAIN, data=markup_text[cursor:])

    def tokenize_ansi(self, ansi: str) -> Iterator[Token]:
        """Converts the given ANSI string into an iterator of `Token`.

        Args:
            ansi: The text to look at.

        Returns:
            An iterator of tokens. The reason this is an iterator is to possibly save
            on memory.
        """

        def _is_in_tags(code: str, tags: dict[str, str]) -> str | None:
            """Determines whether a code is in the given dict of tags."""

            for name, current in tags.items():
                if current == code:
                    return name

            return None

        def _generate_color(
            parts: list[str], code: str
        ) -> tuple[str, TokenType, Color]:
            """Generates a color token."""

            data: Color
            if len(parts) == 1:
                data = StandardColor.from_ansi(code)
                name = data.name
                ttype = TokenType.COLOR

            else:
                data = str_to_color(code)
                name = data.name
                ttype = TokenType.COLOR

            return name, ttype, data

        end = 0
        start = 0
        cursor = 0

        # StyledText messes with indexing, so we need to cast it
        # back to str.
        if isinstance(ansi, StyledText):
            ansi = str(ansi)

        for match in RE_ANSI.finditer(ansi):
            code = match.groups()[0]
            start, end = match.span()

            if code is None:
                continue

            parts = code.split(";")

            if start > cursor:
                plain = ansi[cursor:start]

                yield Token(name=plain, ttype=TokenType.PLAIN, data=plain)

            name: str | None = code
            ttype = None
            data: str | Color = parts[0]

            # Styles & Unsetters
            if len(parts) == 1:
                # Covariancy is not an issue here, even though mypy seems to think so.
                name = _is_in_tags(parts[0], self.unsetters)  # type: ignore
                if name is not None:
                    ttype = TokenType.UNSETTER

                else:
                    name = _is_in_tags(parts[0], self.tags)
                    if name is not None:
                        ttype = TokenType.STYLE

            # Colors
            if ttype is None:
                with suppress(ColorSyntaxError):
                    name, ttype, data = _generate_color(parts, code)

            if name is None or ttype is None or data is None:
                if len(parts) != 2:
                    raise AnsiSyntaxError(
                        tag=parts[0], cause="not recognized", context=ansi
                    )

                name = "position"
                ttype = TokenType.POSITION
                data = ",".join(reversed(parts))

            yield Token(name=name, ttype=ttype, data=data)
            cursor = end

        if cursor < len(ansi):
            plain = ansi[cursor:]

            yield Token(ttype=TokenType.PLAIN, data=plain)

    def define(self, name: str, method: MacroCallable) -> None:
        """Defines a Macro tag that executes the given method.

        Args:
            name: The name the given method will be reachable by within markup.
                The given value gets "!" prepended if it isn't present already.
            method: The method this macro will execute.
        """

        if not name.startswith("!"):
            name = f"!{name}"

        self.macros[name] = method
        self.unsetters[f"/{name}"] = None

    def alias(self, name: str, value: str) -> None:
        """Aliases the given name to a value, and generates an unsetter for it.

        Note that it is not possible to alias macros.

        Args:
            name: The name of the new tag.
            value: The value the new tag will stand for.
        """

        def _get_unsetter(token: Token) -> str | None:
            """Get unsetter for a token"""

            if token.ttype is TokenType.PLAIN:
                return None

            if token.ttype is TokenType.UNSETTER:
                return self.unsetters[token.name]

            if token.ttype is TokenType.COLOR:
                assert isinstance(token.data, Color)

                if token.data.background:
                    return self.unsetters["/bg"]

                return self.unsetters["/fg"]

            name = f"/{token.name}"
            if not name in self.unsetters:
                raise KeyError(f"Could not find unsetter for token {token}.")

            return self.unsetters[name]

        if name.startswith("!"):
            raise ValueError('Only macro tags can always start with "!".')

        setter = ""
        unsetter = ""

        # Try to link to existing tag
        if value in self.user_tags:
            self.unsetters[f"/{name}"] = self.unsetters[f"/{value}"]
            self.user_tags[name] = self.user_tags[value]
            return

        for token in self.tokenize_markup(f"[{value}]"):
            if token.ttype is TokenType.PLAIN:
                continue

            assert token.sequence is not None
            setter += token.sequence

            t_unsetter = _get_unsetter(token)
            unsetter += f"\x1b[{t_unsetter}m"

        self.unsetters[f"/{name}"] = unsetter.lstrip("\x1b[").rstrip("m")
        self.user_tags[name] = setter.lstrip("\x1b[").rstrip("m")

        marked: list[str] = []
        for item in self._cache:
            if name in item:
                marked.append(item)

        for item in marked:
            del self._cache[item]

    # TODO: I cannot cut down the one-too-many branch that this has at the moment.
    #       We could look into it in the future, however.
    def parse(  # pylint: disable=too-many-branches
        self, markup_text: str
    ) -> StyledText:
        """Parses the given markup.

        Args:
            markup_text: The markup to parse.

        Returns:
            A `StyledText` instance of the result of parsing the input. This
            custom `str` class is used to allow accessing the plain value of
            the output, as well as to cleanly index within it. It is analogous
            to builtin `str`, only adds extra things on top.
        """

        applied_macros: list[tuple[str, MacroCall]] = []
        previous_token: Token | None = None
        previous_sequence = ""
        sequence = ""
        out = ""

        def _apply_macros(text: str) -> str:
            """Apply current macros to text"""

            for _, (method, args) in applied_macros:
                text = method(*args, text)

            return text

        def _is_same_colorgroup(previous: Token, new: Token) -> bool:
            if not isinstance(new.data, Color) or not isinstance(previous.data, Color):
                return False

            return (
                type(previous) is type(new)
                and previous.data.background == new.data.background
            )

        if (
            self.should_cache
            and markup_text in self._cache
            and len(RE_MACRO.findall(markup_text)) == 0
        ):
            return self._cache[markup_text]

        token: Token
        for token in self.tokenize_markup(markup_text):
            if sequence != "" and previous_token == token:
                continue

            # Optimize out previously added color tokens, as only the most
            # recent would be visible anyways.
            if (
                token.sequence is not None
                and previous_token is not None
                and _is_same_colorgroup(previous_token, token)
            ):
                sequence = token.sequence
                continue

            if token.ttype == TokenType.UNSETTER and token.data == "0":
                out += "\033[0m"
                sequence = ""
                applied_macros = []
                continue

            previous_token = token

            # Macro unsetters are stored with None as their data
            if token.data is None and token.ttype is TokenType.UNSETTER:
                for item, data in applied_macros.copy():
                    macro_match = RE_MACRO.match(item)
                    assert macro_match is not None

                    macro_name = macro_match.groups()[0]

                    if f"/{macro_name}" == token.name:
                        applied_macros.remove((item, data))

                continue

            if token.ttype is TokenType.MACRO:
                assert isinstance(token.data, tuple)

                applied_macros.append((token.name, token.data))
                continue

            if token.sequence is None:
                applied = sequence

                if not out.endswith("\x1b[0m"):
                    for item in previous_sequence.split("\x1b"):
                        if item == "" or item[1:-1] in self.unsetters.values():
                            continue

                        item = f"\x1b{item}"
                        applied = applied.replace(item, "")

                out += applied + _apply_macros(token.name)
                previous_sequence = sequence
                sequence = ""
                continue

            sequence += token.sequence

        if sequence + previous_sequence != "":
            out += "\x1b[0m"

        out = StyledText(out)
        self._cache[markup_text] = out
        return out

    def get_markup(self, ansi: str) -> str:
        """Generates markup from ANSI text.

        Args:
            ansi: The text to get markup from.

        Returns:
            A markup string that can be parsed to get (visually) the same
            result. Note that this conversion is lossy in a way: there are some
            details (like macros) that cannot be preserved in an ANSI->Markup->ANSI
            conversion.
        """

        current_tags: list[str] = []
        out = ""
        for token in self.tokenize_ansi(ansi):
            if token.ttype is TokenType.PLAIN:
                if len(current_tags) != 0:
                    out += "[" + " ".join(current_tags) + "]"

                assert isinstance(token.data, str)
                out += token.data
                current_tags = []
                continue

            if token.ttype is TokenType.ESCAPED:
                assert isinstance(token.data, str)

                current_tags.append(token.data)
                continue

            current_tags.append(token.name)

        return out

    def prettify_ansi(self, text: str) -> str:
        """Returns a prettified (syntax-highlighted) ANSI str.

        This is useful to quickly "inspect" a given ANSI string. However,
        for most real uses `MarkupLanguage.prettify_markup` would be
        preferable, given an argument of `MarkupLanguage.get_markup(text)`,
        as it is much more verbose.

        Args:
            text: The ANSI-text to prettify.

        Returns:
            The prettified ANSI text. This text's styles remain valid,
            so copy-pasting the argument into a command (like printf)
            that can show styled text will work the same way.
        """

        out = ""
        sequences = ""
        for token in self.tokenize_ansi(text):
            if token.ttype is TokenType.PLAIN:
                assert isinstance(token.data, str)
                out += token.data
                continue

            assert token.sequence is not None
            out += "\x1b[0m" + token.sequence + token.sequence.replace("\x1b", "\\x1b")
            sequences += token.sequence
            out += sequences

        return out

    def prettify_markup(self, text: str) -> str:
        """Returns a prettified (syntax-highlighted) markup str.

        Args:
            text: The markup-text to prettify.

        Returns:
            Prettified markup. This markup, excluding its styles,
            remains valid markup.
        """

        def _apply_macros(text: str) -> str:
            """Apply current macros to text"""

            for _, (method, args) in applied_macros:
                text = method(*args, text)

            return text

        def _pop_macro(name: str) -> None:
            """Pops a macro from applied_macros."""

            for i, (macro_name, _) in enumerate(applied_macros):
                if macro_name == name:
                    applied_macros.pop(i)
                    break

        def _finish(out: str, in_sequence: bool) -> str:
            """Adds ending cap to the given string."""

            if in_sequence:
                if not out.endswith("\x1b[0m"):
                    out += "\x1b[0m"

                return out + "]"

            return out + "[/]"

        styles: dict[TokenType, str] = {
            TokenType.MACRO: "210",
            TokenType.ESCAPED: "210 bold",
            TokenType.UNSETTER: "strikethrough",
        }

        applied_macros: list[tuple[str, MacroCall]] = []

        out = ""
        in_sequence = False
        current_styles: list[Token] = []

        for token in self.tokenize_markup(text):
            if token.ttype in [TokenType.PLAIN, TokenType.ESCAPED]:
                if in_sequence:
                    out += "]"

                in_sequence = False

                sequence = ""
                for style in current_styles:
                    if style.sequence is None:
                        continue

                    sequence += style.sequence

                out += f"{sequence}{_apply_macros(token.name)}\033[0m"
                continue

            out += " " if in_sequence else "["
            in_sequence = True

            if token.ttype is TokenType.UNSETTER:
                if token.name == "/":
                    applied_macros = []

                name = token.name[1:]

                if name in self.macros:
                    _pop_macro(name)

                current_styles.append(token)

                out += self.parse(
                    ("" if (name in self.tags) or (name in self.user_tags) else "")
                    + f"[{styles[TokenType.UNSETTER]}]/{name}"
                )
                continue

            if token.ttype is TokenType.MACRO:
                assert isinstance(token.data, tuple)

                name = token.name
                if "(" in name:
                    name = name[: token.name.index("(")]

                applied_macros.append((name, token.data))

                try:
                    out += token.data[0](*token.data[1], token.name)
                    continue

                except TypeError:  # Not enough arguments
                    pass

            if token.sequence is not None:
                current_styles.append(token)

            style_markup = styles.get(token.ttype) or token.name
            out += self.parse(f"[{style_markup}]{token.name}")

        return _finish(out, in_sequence)

    def get_styled_plains(self, text: str) -> Iterator[StyledText]:
        """Gets all plain tokens within text, with their respective styles applied.

        Args:
            text: The ANSI-sequence containing string to find plains from.

        Returns:
            An iterator of `StyledText` objects, each yielded when a new plain token is found,
            containing the styles that are relevant and active on the given plain.
        """

        def _apply_styles(styles: list[Token], text: str) -> str:
            """Applies given styles to text."""

            for token in styles:
                if token.ttype is TokenType.MACRO:
                    assert isinstance(token.data, tuple)
                    text = token.data[0](*token.data[1], text)
                    continue

                if token.sequence is None:
                    continue

                text = token.sequence + text

            return text

        def _pop_unsetter(token: Token, styles: list[Token]) -> list[Token]:
            """Removes an unsetter from the list, returns the new list."""

            if token.name == "/":
                return list(filter(lambda tkn: tkn.ttype is TokenType.POSITION, styles))

            target_name = token.name[1:]
            for style in styles:
                # bold & dim unsetters represent the same character, so we have
                # to treat them the same way.
                style_name = style.name

                if style.name == "dim":
                    style_name = "bold"

                if style_name == target_name:
                    styles.remove(style)

                elif (
                    style_name.startswith(target_name)
                    and style.ttype is TokenType.MACRO
                ):
                    styles.remove(style)

                elif style.ttype is TokenType.COLOR:
                    assert isinstance(style.data, Color)
                    if target_name == "fg" and not style.data.background:
                        styles.remove(style)

                    elif target_name == "bg" and style.data.background:
                        styles.remove(style)

            return styles

        def _pop_position(styles: list[Token]) -> list[Token]:
            for token in styles.copy():
                if token.ttype is TokenType.POSITION:
                    styles.remove(token)

            return styles

        styles: list[Token] = []
        for token in self.tokenize_ansi(text):
            if token.ttype is TokenType.COLOR:
                for i, style in enumerate(reversed(styles)):
                    if style.ttype is TokenType.COLOR:
                        assert isinstance(style.data, Color)
                        assert isinstance(token.data, Color)

                        if style.data.background != token.data.background:
                            continue

                        styles[len(styles) - i - 1] = token
                        break
                else:
                    styles.append(token)

                continue

            if token.ttype is TokenType.LINK:
                styles.append(token)
                yield StyledText(_apply_styles(styles, token.name))

            if token.ttype is TokenType.PLAIN:
                assert isinstance(token.data, str)
                yield StyledText(_apply_styles(styles, token.data))
                styles = _pop_position(styles)
                continue

            if token.ttype is TokenType.UNSETTER:
                styles = _pop_unsetter(token, styles)
                continue

            styles.append(token)

def main() -> None:
    """Main method"""

    parser = ArgumentParser()

    markup_group = parser.add_argument_group("Markup->ANSI")
    markup_group.add_argument(
        "-p", "--parse", metavar=("TXT"), help="parse a markup text"
    )
    markup_group.add_argument(
        "-e", "--escape", help="escape parsed markup", action="store_true"
    )
    # markup_group.add_argument(
    # "-o",
    # "--optimize",
    # help="set optimization level for markup parsing",
    # action="count",
    # default=0,
    # )

    markup_group.add_argument("--alias", action="append", help="alias src=dst")

    ansi_group = parser.add_argument_group("ANSI->Markup")
    ansi_group.add_argument(
        "-m", "--markup", metavar=("TXT"), help="get markup from ANSI text"
    )
    ansi_group.add_argument(
        "-s",
        "--show-inverse",
        action="store_true",
        help="show result of parsing result markup",
    )

    args = parser.parse_args()

    lang = MarkupLanguage()

    if args.markup:
        markup_text = lang.get_markup(args.markup)
        print(markup_text, end="")

        if args.show_inverse:
            print("->", lang.parse(markup_text))
        else:
            print()

    if args.parse:
        if args.alias:
            for alias in args.alias:
                src, dest = alias.split("=")
                lang.alias(src, dest)

        parsed = lang.parse(args.parse)

        if args.escape:
            print(ascii(parsed))
        else:
            print(parsed)

        return

tim = markup = MarkupLanguage()
"""The default TIM instances."""

if __name__ == "__main__":
    main()


import math
import numpy as np
from typing import Tuple

from replay_buffer import Block
from learner import calculate_mixed_td_errors

class LocalBuffer:
    """Store transitions of one episode"""

    def __init__(
        self,
        action_dim: int,
        config
    ):

        forward_steps = config["forward_steps"]
        burn_in_steps = config["burn_in_steps"]
        learning_steps = config["learning_steps"]
        gamma = config["gamma"]
        hidden_dim = config["hidden_dim"]
        block_length = config["block_length"]

        self.action_dim = action_dim
        self.gamma = gamma
        self.hidden_dim = hidden_dim
        self.forward_steps = forward_steps
        self.learning_steps = learning_steps
        self.burn_in_steps = burn_in_steps
        self.block_length = block_length
        self.curr_burn_in_steps = 0
        self.full_reward_buffer = []
        self.full_qval_buffer = []

    def __len__(self):
        return self.size

    def __getitem__(self, key):
        idx, buffer_type = key
        if buffer_type == "reward":
            return self.full_reward_buffer[idx]
        elif buffer_type == "qval":
            return self.full_qval_buffer[idx]

    def reset(self, init_obs: np.ndarray):
        self.obs_buffer = [init_obs]
        self.last_action_buffer = [
            np.array([1 if i == 0 else 0 for i in range(self.action_dim)], dtype=bool)
        ]
        self.last_reward_buffer = [0]
        self.hidden_buffer = [np.zeros((2, self.hidden_dim), dtype=np.float32)]
        self.action_buffer = []
        self.reward_buffer = []
        self.qval_buffer = []
        self.curr_burn_in_steps = 0
        self.size = 0
        self.sum_reward = 0
        self.done = False

    def reset_full_buffers(self):
        self.full_reward_buffer = []
        self.full_qval_buffer = []

    def add(
        self,
        action: int,
        reward: float,
        next_obs: np.ndarray,
        q_value: np.ndarray,
        hidden_state: np.ndarray,
    ):
        self.action_buffer.append(action)
        self.reward_buffer.append(reward)
        self.full_reward_buffer.append(reward)
        self.hidden_buffer.append(hidden_state)
        self.obs_buffer.append(next_obs)
        self.last_action_buffer.append(
            np.array(
                [1 if i == action else 0 for i in range(self.action_dim)], dtype=bool
            )
        )
        self.last_reward_buffer.append(reward)
        self.qval_buffer.append(q_value)
        self.full_qval_buffer.append(q_value)
        self.sum_reward += reward
        self.size += 1

    def finish(self, last_qval: np.ndarray = None) -> Tuple:
        assert self.size <= self.block_length
        # assert len(self.last_action_buffer) == self.curr_burn_in_steps + self.size + 1

        num_sequences = math.ceil(self.size / self.learning_steps)

        max_forward_steps = min(self.size, self.forward_steps)
        n_step_gamma = [self.gamma**self.forward_steps] * (
            self.size - max_forward_steps
        )

        # last_qval is none means episode done
        if last_qval is not None:
            self.qval_buffer.append(last_qval)
            n_step_gamma.extend(
                [self.gamma**i for i in reversed(range(1, max_forward_steps + 1))]
            )
        else:
            self.done = True
            self.qval_buffer.append(np.zeros_like(self.qval_buffer[0]))
            n_step_gamma.extend(
                [0 for _ in range(max_forward_steps)]
            )  # set gamma to 0 so don't need 'done'

        n_step_gamma = np.array(n_step_gamma, dtype=np.float32)

        obs = np.stack(self.obs_buffer)
        last_action = np.stack(self.last_action_buffer)
        last_reward = np.array(self.last_reward_buffer, dtype=np.float32)

        hiddens = np.stack(self.hidden_buffer[slice(0, self.size, self.learning_steps)])

        actions = np.array(self.action_buffer, dtype=np.uint8)

        qval_buffer = np.concatenate(self.qval_buffer)
        reward_buffer = self.reward_buffer + [0 for _ in range(self.forward_steps - 1)]
        n_step_reward = np.convolve(
            reward_buffer,
            [
                self.gamma ** (self.forward_steps - 1 - i)
                for i in range(self.forward_steps)
            ],
            "valid",
        ).astype(np.float32)

        burn_in_steps = np.array(
            [
                min(
                    i * self.learning_steps + self.curr_burn_in_steps,
                    self.burn_in_steps,
                )
                for i in range(num_sequences)
            ],
            dtype=np.uint8,
        )
        learning_steps = np.array(
            [
                min(self.learning_steps, self.size - i * self.learning_steps)
                for i in range(num_sequences)
            ],
            dtype=np.uint8,
        )
        forward_steps = np.array(
            [
                min(self.forward_steps, self.size + 1 - np.sum(learning_steps[: i + 1]))
                for i in range(num_sequences)
            ],
            dtype=np.uint8,
        )
        assert forward_steps[-1] == 1 and burn_in_steps[0] == self.curr_burn_in_steps

        max_qval = np.max(qval_buffer[max_forward_steps : self.size + 1], axis=1)
        max_qval = np.pad(max_qval, (0, max_forward_steps - 1), "edge")
        target_qval = qval_buffer[np.arange(self.size), actions]

        td_errors = np.abs(
            n_step_reward + n_step_gamma * max_qval - target_qval, dtype=np.float32
        )
        priorities = np.zeros(
            self.block_length // self.learning_steps, dtype=np.float32
        )
        priorities[:num_sequences] = calculate_mixed_td_errors(
            td_errors, learning_steps
        )

        # save burn in information for next block
        self.obs_buffer = self.obs_buffer[-self.burn_in_steps - 1 :]
        self.last_action_buffer = self.last_action_buffer[-self.burn_in_steps - 1 :]
        self.last_reward_buffer = self.last_reward_buffer[-self.burn_in_steps - 1 :]
        self.hidden_buffer = self.hidden_buffer[-self.burn_in_steps - 1 :]
        self.action_buffer.clear()
        self.reward_buffer.clear()
        self.qval_buffer.clear()
        self.curr_burn_in_steps = len(self.obs_buffer) - 1
        self.size = 0

        block = Block(
            obs,
            last_action,
            last_reward,
            actions,
            n_step_reward,
            n_step_gamma,
            hiddens,
            num_sequences,
            burn_in_steps,
            learning_steps,
            forward_steps,
        )
        return [block, priorities, self.sum_reward if self.done else None]


import sys
import copy
import time
import paddle
import paddle.nn as nn
import paddle.distributed as dist
# from paddle.optimizer import lr
# from paddle.optimizer.lr import LRScheduler

sys.path.append('.')
# sys.path.append("/home/aistudio/PaddleOCR-2.6.0/")
# sys.path.append('/home/aistudio/external-libraries')
# from ppocr.utils.logging import get_logger
# from ppocr.utils.utility import set_seed
# from ppocr.modeling.architectures import apply_to_static
# from ppocr.utils.save_load import load_model, save_model

from test_util import save_model, apply_to_static, load_model, save_model, build_optimizer
from test_program import preprocess
from test_evaluation import validate
from test_data import build_dataloader
from test_model import SPTS

import importlib
def build_model(config):
    config = copy.deepcopy(config)
    if not "name" in config:
        from test_model import BaseModel
        arch = BaseModel(config)
    else:
        name = config.pop("name")
        mod = importlib.import_module(__name__)
        arch = getattr(mod, name)(config)
    return arch

def build_loss(config):
    weight = paddle.ones([config['num_classes']])
    weight[config['eos_index']] = config['eos_loss_coef']
    module_class = nn.CrossEntropyLoss(weight=weight, ignore_index=config['padding_index'])
    return module_class

def main(config, device, logger, vdl_writer):
    # init dist environment
    if config['Global']['distributed']:
        dist.init_parallel_env()
    global_config = config['Global']
    epoch_num = global_config['epoch_num']

    # build dataloder
    train_dataloader = build_dataloader(config, 'Train', device, logger)
    if config['Eval']:
        valid_dataloader = build_dataloader(config, 'Eval', device, logger)
    else:
        valid_dataloader = None

    # build model
    model = build_model(config['Architecture'])
    model = apply_to_static(model, config, logger)

    # build loss
    loss_class = build_loss(config['Loss'])

    # build optim
    # from ppocr.optimizer import build_optimizer
    optimizer, lr_scheduler = build_optimizer(
        config['Optimizer'],
        epochs=config['Global']['epoch_num'],
        step_each_epoch=len(train_dataloader),
        model=model)

    # # build metric
    # eval_class = build_metric(config['Metric'])
    # eval_class = validate

    logger.info('train dataloader has {} iters'.format(len(train_dataloader)))

    # load pretrain model
    pre_best_model_dict = load_model(config, model, optimizer)
    # model_after_load = model.state_dict()

    if config['Global']['distributed']:
        model = paddle.DataParallel(model)
    global_step = 0

    for epoch in range(epoch_num):
        model.train()

        # training log
        train_reader_cost = 0.0
        train_run_cost = 0.0
        total_samples = 0
        reader_start = time.time()
        batch_past = 0
        print_freq = 1
        for batch_idx, data in enumerate(train_dataloader):
            image = data['image']
            data['image'] = paddle.concat([i for i in image], axis=0)
            sequence = data['sequence']
            data['sequence'] = paddle.concat([i for i in sequence], axis=0)
            mask = data['mask']
            data['mask'] = paddle.concat([i for i in mask], axis=0)

            train_reader_cost += time.time() - reader_start
            train_start = time.time()

            lr = optimizer.get_lr()
            sequence = data['sequence']
            out_seq_label = sequence[:, 1, :].astype("int")
            outputs = model(data)
            loss = loss_class(outputs, out_seq_label)
            loss.backward()
            optimizer.clear_grad()
            optimizer.step()
            global_step += 1
            lr_scheduler.step()

            train_run_cost += time.time() - train_start
            # acc = utils.accuracy(output, target).item()
            total_samples += data['image'].shape[0]
            batch_past += 1
            if global_step >= 0 and global_step % print_freq == 0:
                msg = "[Epoch {}/{}, iter: {}] lr: {}, loss: {}, avg_reader_cost: {:.5f} sec, avg_batch_cost: {:.5f} sec, avg_samples: {}, avg_ips: {:.5f} images/sec.".format(
                    epoch+1, epoch_num, batch_idx+1,
                    optimizer.get_lr(),
                    loss.item(), train_reader_cost / batch_past,
                    (train_reader_cost + train_run_cost) / batch_past,
                    total_samples / batch_past,
                    total_samples / (train_reader_cost + train_run_cost))
                # just log on 1st device
                if paddle.distributed.get_rank() <= 0:
                    logger.info(msg)
                    # print(msg)
                # sys.stdout.flush()
                train_reader_cost = 0.0
                train_run_cost = 0.0
                total_samples = 0
                batch_past = 0
        lr_scheduler.step()

        if epoch % 2 == 0:
            save_model_dir = global_config['save_model_dir']
            save_model(
                    model,
                    optimizer,
                    save_model_dir,
                    logger,
                    config,
                    is_best=True,
                    prefix='epoch_{}'.format(epoch),
                    # best_model_dict=best_model_dict,
                    epoch=epoch,
                    global_step=global_step)
            save_model(
                    model,
                    optimizer,
                    save_model_dir,
                    logger,
                    config,
                    is_best=True,
                    prefix='latest',
                    # best_model_dict=best_model_dict,
                    epoch=epoch,
                    global_step=global_step)

        # validate(model, valid_dataloader, epoch, config, logger)

if __name__ == '__main__':
    config, device, logger, vdl_writer = preprocess(is_train=True)
    # seed = config['Global']['seed'] if 'seed' in config['Global'] else 1024
    # set_seed(seed)
    main(config, device, logger, vdl_writer)