from bs4 import BeautifulSoup
import requests
import pandas as pd
from random import randint
import re
import spotipy
from spotipy.oauth2 import SpotifyClientCredentials
from random import randint
from time import sleep
from sklearn import datasets
# K-means
import numpy as np
import matplotlib.pyplot as plt
from sklearn import cluster, datasets
from sklearn.preprocessing import StandardScaler
from matplotlib.lines import Line2D
from sklearn.cluster import KMeans

def recommending_artist_improved(artist):
    songs_to_recommend_top_100 = []
    position_top_100 = []
    songs_to_recommend_second_top = []
    position_second_top = []
    for i in range(len(top_100['artist'])):
        text1 = top_100['artist'][i]
        text2 = second_top['artist'][i]
        pattern = artist
        number_1 = re.findall(pattern, text1)
        number_2 = re.findall(pattern, text2)
        if len(number_1)> 0:
            songs_to_recommend_top_100.append(top_100['song'][i])
            position_top_100.append(i)
            number_1.pop(0)
        if len(number_2)> 0:
            songs_to_recommend_second_top.append(second_top['song'][i])
            position_second_top.append(i)
            number_2.pop(0)
    if len(songs_to_recommend_top_100)> 0:
        print('Good choice! This artist has ' + str(len(songs_to_recommend_top_100)) + ' songs:')
        for i in range(len(songs_to_recommend_top_100)):
            print('-' + songs_to_recommend_top_100[i] + ' in position ' + str(position_top_100[i] + 1))
        return
    elif len(songs_to_recommend_second_top)> 0:
        print('Good choice! This artist has ' + str(len(songs_to_recommend_second_top)) + ' songs:')
        for i in range(len(songs_to_recommend_second_top)):
            print('-' + songs_to_recommend_second_top[i] + ' in position ' + str(position_second_top[i] + 1))
        return
    else:
        artist_songs = comparing_table[comparing_table['artist'] == artist].head()
        if len(artist_songs['artist']) > 0:
            random_one = artist_songs['song'].values[0]
            random_two = artist_songs['song'].values[1]
            print("Sorry, this artist isn't now into the top 100, but in my list i have " + str(len(comparing_table[comparing_table['artist'] == artist])) + " " + artist+ "'s songs, for example:\n")
            print("- " + random_one + "\n")
            print("- " + random_two)
            return 0
        else:
            print("Sorry, this artist isn't into the top 100 and in my list there are no songs from this artist")
            return 0

def recommending_song_improved(song):
    songs_to_recommend_top_100 = []
    position_top_100 = []
    songs_to_recommend_second_top = []
    position_second_top = []
    for i in range(len(top_100['song'])):
        if song == top_100['song'][i]:
            songs_to_recommend_top_100.append(top_100['song'][i])
            position_top_100.append(i)
        elif song == second_top['song'][i]:
            songs_to_recommend_second_top.append(second_top['song'][i])
            position_second_top.append(i)
    if len(songs_to_recommend_top_100)> 0:
        random_number = position_top_100[0]
        while random_number == position_top_100[0]:
            random_number = randint(0,99)
        i_recommend = top_100['song'][random_number]
        return print("Good choice! This song is now in the position number " + str(position_top_100[0]+1) + " of our first top from " + top_100['artist'][position_top_100[0]] + ". Here we have another recommendation of this top 100 songs: " + i_recommend + " from " + top_100['artist'][random_number])
    elif len(songs_to_recommend_second_top)> 0:
        random_number = position_second_top[0]
        while random_number == position_second_top[0]:
            random_number = randint(0,99)
        i_recommend = second_top['song'][random_number]
        return print("Good choice! This song is now in the position number " + str(position_second_top[0]+1) + " of our second top from " + second_top['artist'][position_second_top[0]] + ". Here we have another recommendation of this top 100 songs: " + i_recommend + " from " + second_top['artist'][random_number])
    else:
        results = sp.search(q=song, type = 'track', limit=1)
        uri = results["tracks"]["items"][0]["uri"]
        features = sp.audio_features(uri)
        data_features = pd.DataFrame(features)
        numeric_features = data_features._get_numeric_data()
        cluster_prediction = kmeans.predict(numeric_features)[0]
        similar_songs = comparing_table[comparing_table['cluster_type'] == cluster_prediction].sample(1)
        song_recommended = similar_songs['song'].values[0]
        artist_recommended = similar_songs['artist'].values[0]
        print("Your song is not in the top 100, but we can recommend you this song from our database: "+ song_recommended + " from " + artist_recommended+ "... Hope you like it!!!")
        return 0

def asking_song_improved():
    print("Which song do you want?")
    song_name = input()
    value = recommending_song_improved(song_name)
    while value == 0:
        print("Do you want to try again?(yes/no)")
        yes_no = input()
        if yes_no == 'yes':
            print("Which new song do you want to try?")
            song_name = input()
            value = recommending_song_improved(song_name)
        elif yes_no == 'no':
            value = 1
        else:
            print("That's not a valid value")
            value = 1
    return

def asking_artist_improved():
    print("Which artist do you want?")
    artist_name = input()
    value = recommending_artist_improved(artist_name)
    while value == 0:
        print("Do you want to try again?(yes/no)")
        yes_no = input()
        if yes_no == 'yes':
            print("Which new artist do you want to try?")
            song_name = input()
            value = recommending_artist_improved(artist_name)
        elif yes_no == 'no':
            value = 1
        else:
            print("That's not a valid value")
            value = 1
    return

def asking_improved():
    print("If you want to search by song enter 1, and if you want to search by artist enter 2")
    search = input()
    while (search != '1') & (search != '2'):
        print("This isn't a valid value... Try again")
        print("If you want to search by song enter 1, and if you want to search by artist enter 2")
        search = input()
    if search == '1':
        return asking_song_improved()
    elif search == '2':
        return asking_artist_improved()

final_dataframe = pd.read_excel('./files_for_lab/songs_completed.xlsx')
info_dataframe = pd.read_excel('./files_for_lab/songs_info.xlsx')
top_100 = pd.read_excel('./files_for_lab/top_100.xlsx')
second_top = pd.read_excel('./files_for_lab/second_top.xlsx')
comparing_table = pd.read_excel('./files_for_lab/comparing_table.xlsx')

secrets_file = open("secrets.txt","r")

string = secrets_file.read()

secrets_dict={}
for line in string.split('\n'):
    if len(line) > 0:
        secrets_dict[line.split(':')[0]]=line.split(':')[1]

sp = spotipy.Spotify(auth_manager=SpotifyClientCredentials(client_id=secrets_dict['cid'],
                                                           client_secret=secrets_dict['csecret']))

info_dataframe = info_dataframe.drop(['Unnamed: 0'], axis = 1)
scaled_variables = StandardScaler().fit_transform(info_dataframe)
pd.DataFrame(scaled_variables,columns=info_dataframe.columns).head()
kmeans = KMeans(n_clusters=9, random_state=1234)
kmeans.fit(scaled_variables)

asking_improved()


import open3d as op3d, numpy as np, logging

from scipy.io import loadmat
from scipy.spatial import KDTree
from typing import Union, Tuple, List
from pathlib import Path
from frontalization.mesh import Mesh
from utils_mesh import MeshOperator, create_sphere_mesh

path_bfm_model = Path('../BFM/01_MorphableModel.mat')
path_additional_bfm = Path('../BFM/facemodel_info.mat')
path_landmarks_bfm_indeces = Path('../BFM/landmarks68_BFM.anl')
path_bfm_prepared = Path('../BFM/BFM_model_front.mat')

# https://ibug.doc.ic.ac.uk/resources/300-W/
class LandmarksAnnotator():

    def get_left_eye_landmarks_indeces(self):
        return [37, 38, 39, 40, 41, 42]

    def get_right_eye_landmarks_indeces(self):
        return [43, 44, 45, 46, 47, 48]

    def get_nose_vertical_landmarks(self):
        return [28, 29, 30, 31]

    def get_nose_gorizontal_landmarks(self):
        return [32, 33, 34, 35, 36]

    def get_upper_lip_landmarks(self):
        return [49, 50, 51, 52, 53, 54, 55, 61, 62, 63, 64, 65]

    def get_bottom_lip_landmarks(self):
        return [56, 57, 58, 59, 60, 66, 67, 68]

    def get_left_border_of_the_face(self):
        return [1, 2, 3, 4, 5, 6, 7, 8, 9]

    def get_right_border_of_the_face(self):
        return [10, 11, 12, 13, 14, 15, 16, 17]

class BaselFaceModel2009():

    _default_pathes = {"path_bfm_model" : Path('../BFM/01_MorphableModel.mat'),
                       "path_additional_bfm": Path('../BFM/facemodel_info.mat'),
                       "path_landmarks_bfm_indeces" : Path('../BFM/landmarks68_BFM.anl'),
                       "path_bfm_prepared" : Path('../BFM/BFM_model_front.mat')}

    def __init__(self, path_bfm_model, path_bfm_triangulation=None, path_bfm_prepared_model=None,
                 path_landmark_indeces=None, scale_basis:bool=True, create_mean_face_mesh:bool=False):
        self._path_bfm_model = self.check_path(path_bfm_model)
        assert self._path_bfm_model.suffix == '.mat', 'Incorrect format BFM 2009!'
        if path_bfm_triangulation:
            self._path_bfm_additional_info = self.check_path(path_bfm_triangulation)
        if path_landmark_indeces:
            self._path_landmark_label = self.check_path(path_landmark_indeces)
        self._num_sh_basis = 199
        self._num_tex_basis = 199

        # load model
        if path_bfm_prepared_model:
            path_bfm_prepared_model = self.check_path(path_bfm_prepared_model)
            self._init_bfm_model(path_bfm_prepared_model, scale_basis=scale_basis)
        else:
            # without this filepath
            self._init_bfm_model()

        if create_mean_face_mesh:
            self._path_mean_face_shape_obj = self._save_mean_face_mesh_as_obj()

    def check_path(self, path, type='file'):
        if isinstance(path, str):
            path = Path(path)
        if type=='file':
            assert path.is_file(), f'The file {path} not found!'
        else:
            assert path.is_dir(), f'The directory {path} not found!'
        return path

    @property
    def shape_basis(self):
        return self._shape_basis

    def _init_bfm_model(self, path_bfm_prepare_model=None, scale_basis:bool=True):
        self._bfm_model = loadmat(str(self._path_bfm_model))
        self._shape_basis = self._bfm_model['shapePC'].astype(np.float64).reshape((-1, 3, self._num_sh_basis))/1e5  # shape basis
        self._shape_sigmas = self._bfm_model['shapeEV'].astype(np.float64).reshape((-1, self._num_sh_basis))  # corresponding eigen values
        self._texture_basis = self._bfm_model['texPC'].astype(np.float64).reshape((-1, 3, self._num_sh_basis))  # texture basis
        self._texture_sigmas = self._bfm_model['texEV'].astype(np.float64).reshape((-1, self._num_sh_basis))  # eigen value
        if scale_basis:
            self._shape_basis = self._shape_sigmas * self._shape_basis.reshape(-1, self._num_sh_basis)
            self._shape_basis = self.shape_basis.reshape((-1, 3, self._num_sh_basis))
            self._texture_basis = self._texture_sigmas * self._texture_basis.reshape(-1, self._num_sh_basis)
            self._texture_basis = self._texture_basis.reshape((-1, 3, self._num_sh_basis))

        self._tex_mean = self._bfm_model['texMU'].astype(np.float64).reshape((-1, 3))  # mean texture
        self._shape_mean = self._bfm_model['shapeMU'].astype(np.float64).reshape((-1, 3))/1e5 # mean face
        if path_bfm_prepare_model:
            if isinstance(path_bfm_prepare_model, str):
                path_bfm_prepare_model = Path(path_bfm_prepare_model)
            assert path_bfm_prepare_model.is_file(), f'The file {path_bfm_prepare_model} not found!'
            bfm_prepared = loadmat(str(path_bfm_prepare_model)) # load triangulation
            self._triangles = bfm_prepared['tri'].astype(np.int64)-1
            logging.info(f'Extract triangles from prepared BFM model: {path_bfm_prepare_model}!')
        else:
            #self._triangles = bfm_additional_info['tri'].astype(np.int64).reshape((-1, 3))-1
            self._triangles = self._bfm_model['tl'].astype(np.int64) - 1
            self._triangles = self._triangles[:, [1, 0, 2]]

        if hasattr(self, '_path_landmark_label'):
            keypoints_indeces = []
            with self._path_landmark_label.open() as reader:
                for idx_line, line in enumerate(reader):
                    if idx_line == 0:
                        continue
                    index_keypts = int(line.rstrip(' \n'))
                    keypoints_indeces.append(index_keypts)
            self._keypts_indeces = keypoints_indeces
        logging.info(f'The basel face model is load!')

    def get_face_colors(self, tex_coefs:np.ndarray, normalize=True):
        num_basis = tex_coefs.shape[0]
        basis = self._texture_basis.reshape(-1, self._texture_basis.shape[-1])
        basis = basis[:, :num_basis]
        face_colors = basis @ tex_coefs[:, None] + self.mean_texture.reshape(-1, 1)
        if normalize:
            face_colors = face_colors / 255.
        return face_colors.reshape(-1, 3)

    def get_face_shape(self, id_coeff:np.ndarray):
        num_basis = id_coeff.shape[0]
        basis = self._shape_basis.reshape(-1, self._shape_basis.shape[-1])
        basis = basis[:, :num_basis]
        face_shape = basis @ id_coeff[:, None] + self.mean_shape.reshape(-1, 1)
        return face_shape.reshape(-1, 3)

    def get_right_glob_idx_by_left(self, glob_l_idx:int):
        if not hasattr(self, 'left_idx_to_right_idx_sym') or not hasattr(self, 'right_idx_to_left_idx_sym'):
            assert False
        if glob_l_idx not in self.left_idx_to_right_idx_sym:
            return None
        return self.left_idx_to_right_idx_sym[glob_l_idx]

    def get_left_glob_idx_by_right(self, glob_r_idx: int):
        if not hasattr(self, 'left_idx_to_right_idx_sym') or not hasattr(self, 'right_idx_to_left_idx_sym'):
            assert False
        if glob_r_idx not in self.right_idx_to_left_idx_sym:
            return None
        return self.right_idx_to_left_idx_sym[glob_r_idx]

    def get_left_symmetry_index_pt_test(self, glob_left_idx:int):
        if not hasattr(self, 'left_idx_to_right_idx_sym') or not hasattr(self, 'right_idx_to_left_idx_sym'):
            assert False
        if glob_left_idx not in self.left_idx_to_right_idx_sym:
            if glob_left_idx in self.right_idx_to_left_idx_sym:
                print(f'Glob index {glob_left_idx} correspond to right_idx_to_left_idx_sym.')
                return self.right_idx_to_left_idx_sym[glob_left_idx], 'right'
            else:
                print(f'Glob index {glob_left_idx} does not match any point.')
                return None
        return self.left_idx_to_right_idx_sym[glob_left_idx], 'left'

    def get_right_symmetry_index_pt_test(self, glob_right_idx: int):
        if not hasattr(self, 'right_idx_to_left_idx_sym') or not hasattr(self, 'left_idx_to_right_idx_sym'):
            assert False
        if glob_right_idx not in self.right_idx_to_left_idx_sym:
            if glob_right_idx in self.left_idx_to_right_idx_sym:
                print(f'Glob index {glob_right_idx} correspond to left_idx_to_right_idx_sym.')
                return self.left_idx_to_right_idx_sym[glob_right_idx], 'left'
            else:
                print(f'Glob index {glob_right_idx} does not match any point.')
                return None

        return self.right_idx_to_left_idx_sym[glob_right_idx], 'right'

    def get_symmetry_points(self, left_idx:int):
        result = self.get_left_symmetry_index_pt_test(left_idx)
        if result is None:
            print(f'Incorrect global left index: {left_idx}.')
            return None
        right_idx, flag = result
        # need swap
        if flag == 'right':
            print('Need swap left <-> right.')
            temp = left_idx
            left_idx = right_idx
            right_idx = temp

        left_pt, right_pt = self._shape_mean[left_idx], self._shape_mean[right_idx]
        symm_left_right_pts = np.vstack([left_pt[None, :], right_pt[None, :]])
        return symm_left_right_pts

    #TODO: IN COMPLEX CASE NEED CORRECT
    def build_vert_symmetry_pt_correspondence(self, complex:bool=False) -> dict:
        mean_shape = self._shape_mean
        left_part_of_face_mask = (mean_shape[:, 0] < 0).astype(np.bool)
        right_part_of_face_mask = (mean_shape[:, 0] >= 0).astype(np.bool)
        left_part_of_face = mean_shape[left_part_of_face_mask]
        right_part_of_face = mean_shape[right_part_of_face_mask]
        kdtree_r_part_yz_proj = KDTree(data=right_part_of_face[:, 1:], copy_data=True)
        left_idx_to_glob_idx = np.arange(self._shape_mean.shape[0])[left_part_of_face_mask]
        right_idx_to_glob_idx = np.arange(self._shape_mean.shape[0])[right_part_of_face_mask]
        left_idx_to_right_idx_sym = {}
        right_idx_to_left_idx_sym = {}
        for idx_pt_left, pt_left in enumerate(left_part_of_face):
            pt_left_yz = pt_left[1:]
            dist, indeces = kdtree_r_part_yz_proj.query(pt_left_yz, k = np.arange(1, 6))
            # TODO: need check also np.abs(x)
            if complex:
                pts_right_candidates = right_part_of_face[indeces]
                diffs_dists_x0_plane = np.abs(np.abs(pts_right_candidates[:, 0]) - np.abs(pt_left[0]))
                min_right_idx = np.argmin(diffs_dists_x0_plane)
                right_closest_idx = indeces[min_right_idx]
            else:
                right_closest_idx = indeces[0]
            glob_left_idx = left_idx_to_glob_idx[idx_pt_left]
            glob_right_idx = right_idx_to_glob_idx[right_closest_idx]

            left_idx_to_right_idx_sym[glob_left_idx] = glob_right_idx
            right_idx_to_left_idx_sym[glob_right_idx] = glob_left_idx

        self.left_idx_to_right_idx_sym = left_idx_to_right_idx_sym
        self.right_idx_to_left_idx_sym = right_idx_to_left_idx_sym

        return {'left_to_right': left_idx_to_right_idx_sym, 'right_to_left': right_idx_to_left_idx_sym}

    def draw_face_mesh_with_pts(self, added_pts:np.ndarray):
        face_mean_mesh = MeshOperator.get_mesh_by(vertices=self._shape_mean, triangles=self._triangles)
        face_mean_mesh.compute_vertex_normals()
        list_circle_meshs = []
        for pt in added_pts:
            mesh_circle = create_sphere_mesh(pt, radius=0.01, color=[0, 1., 0.])
            list_circle_meshs.append(mesh_circle)
        list_circle_meshs += [face_mean_mesh]
        op3d.visualization.draw_geometries(list_circle_meshs, width=800, height=600)

    @property
    def keypoints_indeces(self):
        return self._keypts_indeces

    @property
    def get_mean_face_mesh_custom(self):
        return Mesh(vertices=self._shape_mean, triangles=self._triangles, textures=self._tex_mean)

    @property
    def mean_shape(self):
        return self._shape_mean

    @property
    def mean_texture(self):
        return self._tex_mean

    @property
    def triangles(self):
        return self._triangles

    @property
    def mean_face_pcd(self):
        return MeshOperator.get_point_cloud_by(self._shape_mean, self._tex_mean)

    @property
    def mean_face_mesh_open3d(self):
        return MeshOperator.get_mesh_by(self._shape_mean, self._triangles, self._tex_mean)

    def get_mean_face_mesh_open3d_with_kpts(self, color_kpts=(0., 1., 0.),
                                                  radius:float=0.005) -> List[op3d.geometry.TriangleMesh]:
        if not hasattr(self, "_keypts_indeces"):
            assert False, "Keypoints indeces not defined!"
        keypts_vertices = self._shape_mean[self._keypts_indeces]
        meshs_keypts = []
        for idx, key_pts in enumerate(keypts_vertices):
            sphere_pt = self._create_sphere_mesh(center=key_pts, radius=radius, color=color_kpts)
            meshs_keypts += [sphere_pt]

        mean_face_mesh = self.mean_face_mesh_open3d
        return [mean_face_mesh] + meshs_keypts

    def _create_sphere_mesh(self, center: np.ndarray, radius: float, color:Union[Tuple, List]):
        mesh_sphere = op3d.geometry.TriangleMesh.create_sphere(radius=radius)
        mesh_sphere.paint_uniform_color(color)
        return mesh_sphere.translate(center)

    def _save_mean_face_mesh_as_obj(self, filepath='/tmp/mean_face_mesh_basel.obj', save_texture_file:bool=True):
        mean_face_mesh = self.mean_face_mesh_open3d
        if isinstance(filepath, str):
            filepath = Path(filepath)
        output = op3d.io.write_triangle_mesh(str(filepath), mesh=mean_face_mesh, write_triangle_uvs=save_texture_file,
                                    print_progress=True)
        logging.info(f'Save mean face shape as {filepath}, Success: {output}.')
        return filepath

if __name__ == '__main__':
    bfm = BaselFaceModel2009(path_bfm_model)
    bfm.build_vert_symmetry_pt_correspondence(complex=True)


import re
from collections import OrderedDict

from rest_framework import serializers
from rest_framework.relations import Hyperlink, PKOnlyObject

from drf_img_crop.serializers import ImageCreateSerializer

from core.models import Profile
from core.api.serializers import ProfileDetailSerializer, ProfileListSerializer

from jinja2 import Environment, BaseLoader

#from drf_hateoas.fields import HateoasUrlField
from drf_hateoas.serializers import HateoasSerializer
from drf_hateoas.fields import HateoasUrlField

from ..models import Blog, Category, Page, Section, HtmlContent, ContentImage
from .validators import TitleValidator

CONTENT_FIELDS = [
        'pk',
        'title',
        'slug',
        'is_visible',
        'creation_date',
        'changed_date',
    ]

HTML_CONTENT_FIELDS = CONTENT_FIELDS + [
        'body',
        'images',
    ]

BASE_KWARGS = {
    'pk': {'read_only': True},
    'slug': {'read_only': True},
    'creation_date': {'read_only': True},
    'changed_date': {'read_only': True},
}

READ_ONLY_KWARGS = {**BASE_KWARGS, **{'title': {'read_only': True}}}

CONTENT_FIELDS_KWARGS = {**BASE_KWARGS, **{'title': {'required': True}}}

BASE_PATH = '/api/content/'

URLS = {
    'page': BASE_PATH + 'page',
    'section': BASE_PATH + 'section',
    'category': BASE_PATH + 'category',
    'blog': BASE_PATH + 'blog',
    'image': BASE_PATH + 'image',
}

def GET_CONTENT_FIELDS_KWARGS(queryset):
    return {**BASE_KWARGS, **{'title': {'required': True, 'validators':[TitleValidator(queryset=queryset)]}}}

#SLUG_REGEX = re.compile('^(\w+)$')

class BaseListSerializer(HateoasSerializer):
    class Meta:
        fields = CONTENT_FIELDS
        extra_kwargs = READ_ONLY_KWARGS

class SectionListSerializer(HateoasSerializer):
    class Meta:
        model = Section
        fields = CONTENT_FIELDS
        extra_kwargs = GET_CONTENT_FIELDS_KWARGS(queryset=model.objects.all())

class SectionSerializer(HateoasSerializer):
    pages = HateoasUrlField(read_only=True, name='Page', many=True,
                view_name='content-api:page-detail')
    class Meta:
        model = Section
        fields = CONTENT_FIELDS + ['pages']
        extra_kwargs = GET_CONTENT_FIELDS_KWARGS(queryset=model.objects.all())

class PageListSerializer(BaseListSerializer):
    class Meta(BaseListSerializer.Meta):
        model = Page

class PageSerializer(HateoasSerializer):
    #images = ContentImageSerializer(many=True, read_only=True)
    #section = SectionSerializer(read_only=True)
    #section = serializers.SerializerMethodField()
    section = HateoasUrlField(read_only=False, name='Section',
        view_name='content-api:section-detail',
        queryset=Section.objects.all(),
        required=False)

    def create(self, validated_data):
        if self.context['section'] is not None:
            validated_data['section'] = self.context['section']
        return super().create(validated_data)

    class Meta:
        model = Page
        fields = CONTENT_FIELDS + [
            'body',
            'section',
        ]
        extra_kwargs = GET_CONTENT_FIELDS_KWARGS(queryset=model.objects.all())

class CategoryListSerializer(BaseListSerializer):
    class Meta(BaseListSerializer.Meta):
        model = Category

class CategorySerializer(serializers.ModelSerializer):
    #blogs = BlogListSerializer(many=True, read_only=True)
    class Meta:
        model = Category
        fields = CONTENT_FIELDS# + ['blogs']
        extra_kwargs = GET_CONTENT_FIELDS_KWARGS(queryset=model.objects.all())

class BlogListSerializer(BaseListSerializer):
    class Meta(BaseListSerializer.Meta):
        model = Blog

class BlogSerializer(HateoasSerializer):
    author = serializers.PrimaryKeyRelatedField(queryset=Profile.objects.all(), required=False)
    #images = ContentImageSerializer(many=True)
    #categories = CategoryListSerializer(read_only=False, many=True)
    class Meta:
        model = Blog
        fields = CONTENT_FIELDS + [
            'body',
            'author',
            #'section',
        ]
        extra_kwargs = GET_CONTENT_FIELDS_KWARGS(queryset=model.objects.all())

class ContentImageListSerializer(BaseListSerializer):
    class Meta(BaseListSerializer.Meta):
        model = ContentImage

class ContentImageSerializer(serializers.ModelSerializer):
    #pages = PageListSerializer(read_only=True, many=True)
    class Meta:
        model = ContentImage
        fields = CONTENT_FIELDS + [
            'image',
            #'pages',
        ]
        extra_kwargs = {
            'image': {'read_only': True},
            'pk': {'read_only': True},
        }

class ContentImageCreateSerializer(ImageCreateSerializer):

    class Meta:
        model = ContentImage
        fields = CONTENT_FIELDS + [
            'image',
            'crop',
        ]
        extra_kwargs = {**{
            'image': {'required': True},
        }, **GET_CONTENT_FIELDS_KWARGS(queryset=model.objects.all())}

#Read only public serializers

class NavbarPageSerializer(serializers.ModelSerializer):
    class Meta:
        model = Page
        fields = [
            'title',
            'slug',
        ]

class NavbarSerializer(serializers.ModelSerializer):
    pages = serializers.SerializerMethodField('get_nav_pages')
    #pages = NavbarPageSerializer(many=True, read_only=True, required=False)

    title = serializers.SerializerMethodField('get_nav_title')
    slug = serializers.SerializerMethodField('get_nav_slug')
    class Meta:
        model = Section
        fields = [
            'title',
            'slug',
            'pages',
        ]
    def to_representation(self, instance):
        result = super(NavbarSerializer, self).to_representation(instance)
        return OrderedDict([(key, result[key]) for key in result if result[key] is not None])

    def get_nav_title(self, obj):
        return obj.nav_title()

    def get_nav_slug(self, obj):
        return obj.nav_slug()

    def get_nav_pages(self, obj):
        if obj.not_dropdown():
            return None
        else:
            return NavbarPageSerializer(obj.nav_pages(), many=True).data

class HtmlPageSerializer(serializers.ModelSerializer):
    body = serializers.SerializerMethodField()
    class Meta:
        model = Page
        fields = [
            'body',
        ]

    def get_body(self, obj):
        templateVars = {}
        for img in obj.images.all():
            templateVars['img__' + str(img.slug)] = '<img src="' + img.image.url +'" />'
        rtemplate = Environment(loader=BaseLoader()).from_string(obj.body)
        return rtemplate.render(templateVars)

class HtmlBlogSerializer(serializers.ModelSerializer):
    class Meta:
        model = Page
        fields = [
            'body',
            'categories',
        ]


import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '1'

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from keras.models import Sequential
from keras.layers import Dense, LSTM
from keras.optimizers import SGD
from keras.utils import to_categorical
from keras.callbacks import EarlyStopping
from keras.losses import BinaryCrossentropy
from sklearn.dummy import DummyClassifier
from sklearn.metrics import confusion_matrix, accuracy_score, precision_score, recall_score

from preprocessing import full_pipeline, stack_windows

data = pd.read_csv('data/cache.csv')
data.fillna(method='ffill', inplace=True)

scaled_data = full_pipeline.fit_transform(data)

# Split into training and testing sets
total_intervals, n_features = scaled_data.shape
train_len = int(total_intervals * 0.8)
train, test = scaled_data[:train_len], scaled_data[train_len:]

window_len = 15
units_ahead = 0 # 0 is the minimum valid value, representing the next unit in the future

X_train, y_train_pct_change = stack_windows(train, window_len, units_ahead, True)
X_test, y_test_pct_change = stack_windows(test, window_len, units_ahead, True)

y_train = to_categorical(y_train_pct_change > 0)
y_test = to_categorical(y_test_pct_change > 0)

# Create and fit the LSTM model
model = Sequential()
model.add(LSTM(128, activation='relu', input_shape=(window_len, n_features)))
model.add(Dense(2, activation = 'softmax'))

opt = SGD(learning_rate=0.01, momentum=0.0)

model.compile(loss='binary_crossentropy', optimizer=opt)

es = EarlyStopping(monitor='val_loss', mode='min', patience=10, verbose=1)
history = model.fit(X_train, y_train, epochs=100, batch_size=100, verbose=1,
                    validation_data=(X_test, y_test), callbacks=[es])

y_pred = model.predict(X_test)

# loss, accuracy = model.evaluate(X_test, y_test)
# print(loss, accuracy)

# Plot loss over time
fig1, ax1 = plt.subplots()

baseline = BinaryCrossentropy()(y_test, np.ones_like(y_test)*0.5)
ax1.plot(history.history['loss'], label='Train Loss')
ax1.plot(history.history['val_loss'], label='Validation Loss')
ax1.axhline(baseline, color='k', ls='--', label='Naive Loss')
ax1.set_xlabel('epochs')
ax1.set_ylabel('loss')
ax1.legend()

# Compare with DummyClassifier
dummy_classifiers = ['stratified', 'most_frequent', 'prior', 'uniform']
for i in dummy_classifiers:
        dummy_classifier = DummyClassifier(strategy = i).fit(X_train, y_train)
        s = dummy_classifier.score(X_test, y_test)
        print('Score for', i ,'classifier is:', s)

y_pred = np.argmax(y_pred, axis=1)
y_test = np.argmax(y_test, axis=1)

print(y_pred)
print(y_test)

# Calculate accuracy, precision and recall
print(confusion_matrix(y_test, y_pred))
print('Accuracy:', accuracy_score(y_test, y_pred))
print('Precision:', precision_score(y_test, y_pred))
print('Recall:', recall_score(y_test, y_pred))

price_scaler = full_pipeline.named_transformers_['price_pipeline'].named_steps['scaler']
y_test_pct_change = price_scaler.inverse_transform(y_test_pct_change[:, np.newaxis])

test_days = len(y_train) + np.arange(len(y_test))

trading = np.cumprod(y_test_pct_change.flatten() * y_pred + 1)
hold = np.cumprod(y_test_pct_change.flatten() + 1)

# Plot comparison of buy & hold vs trading
fig4, ax4 = plt.subplots()
ax4.plot(test_days, hold, 'k-', label='Hold')
ax4.plot(test_days, trading, 'r-', label='Trading')
ax4.legend()

plt.show()


import re

import nltk
from nltk import word_tokenize, TreebankWordDetokenizer
from nltk.corpus import stopwords

#nltk.download()

def extract_line_contains(data: str) -> str:
    return ""

def remove_unwanted_characters(text: str) -> str:
    cleanTxt = text.replace("_", "")
    cleanTxt = cleanTxt.replace("-", "")
    cleanTxt = cleanTxt.replace(".", "")
    cleanTxt = cleanTxt.replace(";", "")

    return cleanTxt

def process_text(text: str) -> str:
    cleaned_text = remove_unwanted_characters(text)
    filtered_text = remove_stopwords(cleaned_text, 'french')
    #correct_spell(cleaned_text)

    return filtered_text

def remove_stopwords(text: str, language: str):
    stop_words = set(stopwords.words(language))

    #regex will search for all words including \n which is needed to preserve the new lines in the detokenizer in line 36
    word_tokens = re.findall(r'\S+|\n', text)
    filtered_words =[]
    for word in word_tokens:
        if word not in stop_words or word.__eq__('\n'):
            filtered_words.append(word)

    filtered_text = TreebankWordDetokenizer().detokenize(filtered_words)
    return filtered_text

# def check_spell(text: str):
#     words = word_tokenize(text)
#     return misspelled

def correct_spell(text: str):
    return text


import os
import utils
import re

scans = utils.read_file(os.getcwd() + "\\input.txt")

def parse_input(scans):
    map = {}
    for y, line in enumerate(scans[:-2]):
        for x, char in enumerate(line):
            if char != " ":
                map[x, y] = char

    instructions = []
    current_num = ""
    for char in scans[-1]:
        if char.isdigit():
            current_num += char
        else:
            instructions.append((int(current_num), char))
            current_num = ""
    instructions.append((int(current_num), "N"))

    return map, instructions

def get_opposite_square(pos, map):
    if pos[2] == 0:
        for x in range(len(scans[pos[1]])):
            if (x, pos[1]) in map.keys():
                return x, pos[1]

    elif pos[2] == 1:
        for y in range(len(scans) - 2):
            if (pos[0], y) in map.keys():
                return pos[0], y

    elif pos[2] == 2:
        for x in range(len(scans[pos[1]]) - 1, -1, -1):
            if (x, pos[1]) in map.keys():
                return x, pos[1]

    elif pos[2] == 3:
        for y in range(len(scans) - 3, -1, -1):
            if (pos[0], y) in map.keys():
                return pos[0], y

    print("error")

def get_forward_square(pos, map):
    forward_pos = (0, 0)
    if pos[2] == 0:
        forward_pos = (pos[0] + 1, pos[1])
    elif pos[2] == 1:
        forward_pos = (pos[0], pos[1] + 1)
    elif pos[2] == 2:
        forward_pos = (pos[0] - 1, pos[1])
    elif pos[2] == 3:
        forward_pos = (pos[0], pos[1] - 1)

    if forward_pos in map.keys():
        return forward_pos
    return get_opposite_square(pos, map)

def navigate(map, instructions):
    pos = (([x for x in range(len(scans[0])) if (x, 0) in map.keys()][0]), 0, 0)

    for curr_instruction in instructions:
        for i in range(curr_instruction[0]):
            coord_next_square = get_forward_square(pos, map)
            if map[coord_next_square] == '#':
                break
            pos = [coord_next_square[0], coord_next_square[1], pos[2]]

        if curr_instruction[1] != 'N':
            new_rotation = (pos[2] + (1 if curr_instruction[1] == "R" else -1)) % 4
            pos = (pos[0], pos[1], new_rotation)

    return pos

if __name__ == "__main__":
    map, instructions = parse_input(scans)
    pos = navigate(map, instructions)
    print(str(1000 * (pos[1] + 1) + 4 * (pos[0] + 1) + pos[2]))


from santaspkg.cost_function import reference_cost_function
from santaspkg.dataset import sample_submission

from pathlib import Path

submission_template = sample_submission.copy(deep=True)

submissions_dir = Path(__file__).parent.parent / 'submissions'
submissions_dir.mkdir(exist_ok=True)

def mk_submit(assignment):
    score = reference_cost_function(assignment)
    filename = str(submissions_dir / f'submission_{score}.csv')
    submission_template['assigned_day'] = assignment
    submission_template.to_csv(filename)


# This Python 3 environment comes with many helpful analytics libraries installed

# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python

# For example, here's several helpful packages to load

import numpy as np # linear algebra

import pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)

# Input data files are available in the read-only "../input/" directory

# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory

import os

for dirname, _, filenames in os.walk('/kaggle/input'):

    for filename in filenames:

        print(os.path.join(dirname, filename))

# You can write up to 5GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using "Save & Run All"

# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session
train_csv = pd.read_csv('/kaggle/input/santander-value-prediction-challenge/train.csv')

test_csv = pd.read_csv('/kaggle/input/santander-value-prediction-challenge/test.csv')
print(train_csv.shape, test_csv.shape)
train_csv.head()
import matplotlib.pyplot as plt

import seaborn as sns
plt.figure(figsize=(12, 8))

plt.scatter(range(train_csv.shape[0]), np.sort(train_csv.target.values))

plt.grid()

plt.xlabel('index', fontsize=12)

plt.ylabel('Target', fontsize=12)

plt.title("Target Distribution", fontsize=14)
plt.figure(figsize=(12,8))

plt.hist(train_csv.target.values, bins=50)

plt.xlabel('Target', fontsize=12)

plt.title("Target Histogram", fontsize=14)
plt.figure(figsize=(12,8))

plt.hist(np.log1p(train_csv.target.values), bins=50)

plt.xlabel('Target', fontsize=12)

plt.title("Target Histogram", fontsize=14)
num_nunique = train_csv.nunique().reset_index()

num_nunique.columns = ['Col_name', 'Value_columns']

num_nunique
only1_nuique = num_nunique[num_nunique.Value_columns == 1]

only1_nuique.shape
from scipy.stats import spearmanr
from tqdm import tqdm, tqdm_notebook

from scipy.stats import spearmanr, pearsonr

import warnings

warnings.filterwarnings('ignore')

col_names = []

cor_value = []

for col in tqdm(train_csv.columns, ncols=100 , leave= True, desc="Spearman r : "):

    if col not in ['ID','target']:

        col_names.append(col)

        cor_value.append(spearmanr(train_csv[col].values, train_csv.target.values)[0])

corrs = pd.DataFrame({'Feature_Name':col_names,'Corr_value':cor_value})

corrs = corrs.sort_values(by = 'Corr_value')
corr_df = corrs[(corrs.Corr_value > 0.1) | (corrs.Corr_value < -0.1)].reset_index()

corr_df.drop('index', axis=1, inplace=True)
corr_df = corr_df.set_index('Feature_Name')

corr_df.plot(kind='barh', figsize = (12,15), title='Correlation of variables')
corr_df = corrs[(corrs.Corr_value > 0.11) | (corrs.Corr_value < -0.11)].reset_index()

corr_df.drop('index', axis = 1, inplace = True)

corr_df = corr_df.set_index('Feature_Name')

corr_df.plot(kind='barh', figsize = (12,15), title='Correlation of variables')
labels = []

values = []

for col in tqdm(train_csv.columns, ncols = 100, desc = 'Pearson r :'):

    if col not in ['ID', 'target']:

        corr_value = pearsonr(train_csv[col].values, train_csv['target'].values)

        values.append(corr_value)

        labels.append(col)

corr_df = pd.DataFrame({'Feature_Name' : labels, 'Feature_Value' : values})

corr_df = corr_df.sort_values(by = 'Feature_Value')
target = train_csv.target

train_csv.drop(list(only1_nuique.Col_name) + ['ID', 'target'], axis = 1, inplace=True)

test_csv.drop(list(only1_nuique.Col_name), axis = 1, inplace=True)
test_id = test_csv.ID

test_csv.drop('ID', axis = 1, inplace=True)
target = np.log1p(target)
from sklearn import ensemble

model = ensemble.ExtraTreesRegressor(n_estimators=200, max_depth=20, max_features=0.5,

                                     n_jobs=-1, random_state=50)

model.fit(train_csv, target)
feature_importances = pd.DataFrame({'Feature_Name' : train_csv.columns, 'feature_importance' : model.feature_importances_})

feature_importances = feature_importances.set_index('Feature_Name')

# sort by importances

feature_importances = feature_importances.sort_values(by = 'feature_importance')

std_feat_importances = np.std([tree.feature_importances_ for tree in model.estimators_],axis=0)

feature_importances.iloc[-50:].plot(kind='barh',figsize = (10,20),color = 'g',

                         xerr=std_feat_importances[-50:], align='center',

                                   title='top 50 Feature Importances')

import lightgbm as lgb

from sklearn import  model_selection, metrics

params = {

    "objective" : "regression",

    "metric" : "rmse",

    "num_leaves" : 30,

    "learning_rate" : 0.01,

    "bagging_seed" : 1884,

    "device" : "gpu",

    "gpu_platform_id" : 0,

    "gpu_device_id" : 0,

    "num_thread" : 8

}

val_data = train_csv[0:500]

val_label = target[0:500]

val_set = lgb.Dataset(val_data, val_label)

train_set = lgb.Dataset(train_csv, label=target)

model = lgb.train(params, train_set ,5000, valid_sets = [val_set],  early_stopping_rounds = 100, verbose_eval = 200)
lgb.plot_importance(model, max_num_features=50, height=0.8, figsize=(12, 20))

plt.title("LightGBM - Feature Importance top 50", fontsize=15)
result = model.predict(test_csv, num_iteration=model.best_iteration)

result = np.expm1(result)
submit = pd.DataFrame({'ID':test_id, 'target':result})

submit.to_csv('submit3_3.csv',index = False)
importance = model.feature_importance(importance_type='split')

feature_name = model.feature_name()
features = pd.DataFrame({'Feature_Name' : feature_name, 'Feature_Important' : importance}).sort_values(by=['Feature_Important'], ascending=[False])[:1000]

train = train_csv[list(features.Feature_Name)]

test = test_csv[list(features.Feature_Name)]
ntrain = train.shape[0]

ntest = test.shape[0]
weight = ((train != 0).sum()/len(train)).values

tmp_train = train[train!=0]

tmp_test = test[test!=0]

train["weight_count"] = (tmp_train*weight).sum(axis=1)

test["weight_count"] = (tmp_test*weight).sum(axis=1)

train["count_not0"] = (train != 0).sum(axis=1)

test["count_not0"] = (test != 0).sum(axis=1)

train["sum"] = train.sum(axis=1)

test["sum"] = test.sum(axis=1)

train["var"] = tmp_train.var(axis=1)

test["var"] = tmp_test.var(axis=1)

train["median"] = tmp_train.median(axis=1)

test["median"] = tmp_test.median(axis=1)

train["mean"] = tmp_train.mean(axis=1)

test["mean"] = tmp_test.mean(axis=1)

train["std"] = tmp_train.std(axis=1)

test["std"] = tmp_test.std(axis=1)

train["max"] = tmp_train.max(axis=1)

test["max"] = tmp_test.max(axis=1)

train["min"] = tmp_train.min(axis=1)

test["min"] = tmp_test.min(axis=1)

train["skew"] = tmp_train.skew(axis=1)

test["skew"] = tmp_test.skew(axis=1)

train["kurtosis"] = tmp_train.kurtosis(axis=1)

test["kurtosis"] = tmp_test.kurtosis(axis=1)
train["weight_count"] = train["weight_count"].fillna(0)

train["count_not0"] = train["count_not0"].fillna(0)

train["sum"] = train["sum"].fillna(0)

train["var"] = train["var"].fillna(0)

train["median"] = train["median"].fillna(0)

train["mean"] = train["mean"].fillna(0)

train["std"] = train["std"].fillna(0)

train["max"] = train["max"].fillna(0)

train["min"] = train["min"].fillna(0)

train["skew"] = train["skew"].fillna(0)

train["kurtosis"] = train["kurtosis"].fillna(0)

test["weight_count"] = test["weight_count"].fillna(0)

test["count_not0"] = test["count_not0"].fillna(0)

test["sum"] = test["sum"].fillna(0)

test["var"] = test["var"].fillna(0)

test["median"] = test["median"].fillna(0)

test["mean"] = test["mean"].fillna(0)

test["std"] = test["std"].fillna(0)

test["max"] = test["max"].fillna(0)

test["min"] = test["min"].fillna(0)

test["skew"] = test["skew"].fillna(0)

test["kurtosis"] = test["kurtosis"].fillna(0)
train.shape
from sklearn import random_projection

NUM_OF_COM = 100

tmp = pd.concat([train,test])

transformer = random_projection.SparseRandomProjection(n_components = NUM_OF_COM)

RP = transformer.fit_transform(tmp)

rp = pd.DataFrame(RP)

columns = ["RandomProjection{}".format(i) for i in range(NUM_OF_COM)]

rp.columns = columns

rp_train = rp[:ntrain]

rp_test = rp[ntrain:]

rp_test.index = test.index

#concat RandomProjection and raw data

train = pd.concat([train,rp_train],axis=1)

test = pd.concat([test,rp_test],axis=1)

train.shape
from sklearn import ensemble

model = ensemble.ExtraTreesRegressor(n_estimators=200, max_depth=20, max_features=0.5,

                                     n_jobs=-1, random_state=50)

model.fit(train, target)
result = model.predict(test)

result = np.expm1(result)

submit = pd.DataFrame({'ID':test_id, 'target':result})

submit.to_csv('submit3_4_extra.csv',index = False)
import xgboost as xgb

import lightgbm as lgb

model_xgb = xgb.XGBRegressor(colsample_bytree=0.055, colsample_bylevel =0.5,

                             gamma=1.5, learning_rate=0.02, max_depth=32,

                             objective='reg:linear',booster='gbtree',

                             min_child_weight=57, n_estimators=1000, reg_alpha=0,

                             reg_lambda = 0,eval_metric = 'rmse', subsample=0.7,

                             silent=1, n_jobs = -1, early_stopping_rounds = 14,

                             random_state =7, nthread = -1)

model_lgb = lgb.LGBMRegressor(objective='regression',num_leaves=144,

                              learning_rate=0.005, n_estimators=720, max_depth=13,

                              metric='rmse',is_training_metric=True,

                              max_bin = 55, bagging_fraction = 0.8,verbose=-1,

                              bagging_freq = 5, feature_fraction = 0.9)

model_xgb.fit(train, target)

model_lgb.fit(train, target)
result = model_xgb.predict(test)

result = np.expm1(result)

submit = pd.DataFrame({'ID':test_id, 'target':result})

submit.to_csv('submit3_4_xgb.csv',index = False)
result = model_lgb.predict(test)

result = np.expm1(result)

submit = pd.DataFrame({'ID':test_id, 'target':result})

submit.to_csv('submit3_4_lgb.csv',index = False)


'''
tree sum

Write a function, tree_sum, that takes in the root of a binary tree that contains number values. The function should return the total sum of all values in the tree.
test_00:

a = Node(3)
b = Node(11)
c = Node(4)
d = Node(4)
e = Node(-2)
f = Node(1)

a.left = b
a.right = c
b.left = d
b.right = e
c.right = f

#       3
#    /    \
#   11     4
#  / \      \
# 4   -2     1

tree_sum(a) # -> 21

test_01:

a = Node(1)
b = Node(6)
c = Node(0)
d = Node(3)
e = Node(-6)
f = Node(2)
g = Node(2)
h = Node(2)

a.left = b
a.right = c
b.left = d
b.right = e
c.right = f
e.left = g
f.right = h

#      1
#    /   \
#   6     0
#  / \     \
# 3   -6    2
#    /       \
#   2         2

tree_sum(a) # -> 10

test_02:

tree_sum(None) # -> 0
'''

# class Node:
#   def __init__(self, val):
#     self.val = val
#     self.left = None
#     self.right = None

# DFS recursive
def tree_sum(root):
  if root is None:
    return 0
  return root.val + tree_sum(root.left) + tree_sum(root.right)

# BFS
from collections import deque

def tree_sum(root):
  if root is None:
    return 0

  queue = deque([root])
  sum = 0
  while queue:
    current = queue.popleft()
    sum += current.val

    if current.left:
      queue.append(current.left)
    if current.right:
      queue.append(current.right)

  return sum


import re
import gast

from tensorflow.python.autograph.pyct import anno
from tensorflow.python.autograph.pyct import origin_info
from tensorflow.python.autograph.pyct import parser
from tensorflow.python.autograph.pyct import transformer
from tensorflow.python.platform import test

class TransformerTest(test.TestCase):

  def _simple_context(self):
    entity_info = transformer.EntityInfo(
        name='Test_fn',
        source_code=None,
        source_file=None,
        future_features=(),
        namespace=None)
    return transformer.Context(entity_info, None, None)

  def assertSameAnno(self, first, second, key):
    self.assertIs(anno.getanno(first, key), anno.getanno(second, key))

  def assertDifferentAnno(self, first, second, key):
    self.assertIsNot(anno.getanno(first, key), anno.getanno(second, key))

  def test_state_tracking(self):

    class LoopState(object):
      pass

    class CondState(object):
      pass

    class TestTransformer(transformer.Base):

      def visit(self, node):
        anno.setanno(node, 'loop_state', self.state[LoopState].value)
        anno.setanno(node, 'cond_state', self.state[CondState].value)
        return super(TestTransformer, self).visit(node)

      def visit_While(self, node):
        self.state[LoopState].enter()
        node = self.generic_visit(node)
        self.state[LoopState].exit()
        return node

      def visit_If(self, node):
        self.state[CondState].enter()
        node = self.generic_visit(node)
        self.state[CondState].exit()
        return node

    tr = TestTransformer(self._simple_context())

    def test_function(a):
      a = 1
      while a:
        _ = 'a'
        if a > 2:
          _ = 'b'
          while True:
            raise '1'
        if a > 3:
          _ = 'c'
          while True:
            raise '1'

    node, _ = parser.parse_entity(test_function, future_features=())
    node = tr.visit(node)

    fn_body = node.body
    outer_while_body = fn_body[1].body
    self.assertSameAnno(fn_body[0], outer_while_body[0], 'cond_state')
    self.assertDifferentAnno(fn_body[0], outer_while_body[0], 'loop_state')

    first_if_body = outer_while_body[1].body
    self.assertDifferentAnno(outer_while_body[0], first_if_body[0],
                             'cond_state')
    self.assertSameAnno(outer_while_body[0], first_if_body[0], 'loop_state')

    first_inner_while_body = first_if_body[1].body
    self.assertSameAnno(first_if_body[0], first_inner_while_body[0],
                        'cond_state')
    self.assertDifferentAnno(first_if_body[0], first_inner_while_body[0],
                             'loop_state')

    second_if_body = outer_while_body[2].body
    self.assertDifferentAnno(first_if_body[0], second_if_body[0], 'cond_state')
    self.assertSameAnno(first_if_body[0], second_if_body[0], 'loop_state')

    second_inner_while_body = second_if_body[1].body
    self.assertDifferentAnno(first_inner_while_body[0],
                             second_inner_while_body[0], 'cond_state')
    self.assertDifferentAnno(first_inner_while_body[0],
                             second_inner_while_body[0], 'loop_state')

  def test_state_tracking_context_manager(self):

    class CondState(object):
      pass

    class TestTransformer(transformer.Base):

      def visit(self, node):
        anno.setanno(node, 'cond_state', self.state[CondState].value)
        return super(TestTransformer, self).visit(node)

      def visit_If(self, node):
        with self.state[CondState]:
          return self.generic_visit(node)

    tr = TestTransformer(self._simple_context())

    def test_function(a):
      a = 1
      if a > 2:
        _ = 'b'
        if a < 5:
          _ = 'c'
        _ = 'd'

    node, _ = parser.parse_entity(test_function, future_features=())
    node = tr.visit(node)

    fn_body = node.body
    outer_if_body = fn_body[1].body
    self.assertDifferentAnno(fn_body[0], outer_if_body[0], 'cond_state')
    self.assertSameAnno(outer_if_body[0], outer_if_body[2], 'cond_state')

    inner_if_body = outer_if_body[1].body
    self.assertDifferentAnno(inner_if_body[0], outer_if_body[0], 'cond_state')

  def test_visit_block_postprocessing(self):

    class TestTransformer(transformer.Base):

      def _process_body_item(self, node):
        if isinstance(node, gast.Assign) and (node.value.id == 'y'):
          if_node = gast.If(
              gast.Name(
                  'x', ctx=gast.Load(), annotation=None, type_comment=None),
              [node], [])
          return if_node, if_node.body
        return node, None

      def visit_FunctionDef(self, node):
        node.body = self.visit_block(
            node.body, after_visit=self._process_body_item)
        return node

    def test_function(x, y):
      z = x
      z = y
      return z

    tr = TestTransformer(self._simple_context())

    node, _ = parser.parse_entity(test_function, future_features=())
    node = tr.visit(node)

    self.assertEqual(len(node.body), 2)
    self.assertIsInstance(node.body[0], gast.Assign)
    self.assertIsInstance(node.body[1], gast.If)
    self.assertIsInstance(node.body[1].body[0], gast.Assign)
    self.assertIsInstance(node.body[1].body[1], gast.Return)

  def test_robust_error_on_list_visit(self):

    class BrokenTransformer(transformer.Base):

      def visit_If(self, node):
        # This is broken because visit expects a single node, not a list, and
        # the body of an if is a list.
        # Importantly, the default error handling in visit also expects a single
        # node.  Therefore, mistakes like this need to trigger a type error
        # before the visit called here installs its error handler.
        # That type error can then be caught by the enclosing call to visit,
        # and correctly blame the If node.
        self.visit(node.body)
        return node

    def test_function(x):
      if x > 0:
        return x

    tr = BrokenTransformer(self._simple_context())

    node, _ = parser.parse_entity(test_function, future_features=())
    with self.assertRaises(ValueError) as cm:
      node = tr.visit(node)
    obtained_message = str(cm.exception)
    expected_message = r'expected "ast.AST", got "\<(type|class) \'list\'\>"'
    self.assertRegex(obtained_message, expected_message)

  def test_robust_error_on_ast_corruption(self):
    # A child class should not be able to be so broken that it causes the error
    # handling in `transformer.Base` to raise an exception.  Why not?  Because
    # then the original error location is dropped, and an error handler higher
    # up in the call stack gives misleading information.

    # Here we test that the error handling in `visit` completes, and blames the
    # correct original exception, even if the AST gets corrupted.

    class NotANode(object):
      pass

    class BrokenTransformer(transformer.Base):

      def visit_If(self, node):
        node.body = NotANode()
        raise ValueError('I blew up')

    def test_function(x):
      if x > 0:
        return x

    tr = BrokenTransformer(self._simple_context())

    node, _ = parser.parse_entity(test_function, future_features=())
    with self.assertRaises(ValueError) as cm:
      node = tr.visit(node)
    obtained_message = str(cm.exception)
    # The message should reference the exception actually raised, not anything
    # from the exception handler.
    expected_substring = 'I blew up'
    self.assertIn(expected_substring, obtained_message)

  def test_origin_info_propagated_to_new_nodes(self):

    class TestTransformer(transformer.Base):

      def visit_If(self, node):
        return gast.Pass()

    tr = TestTransformer(self._simple_context())

    def test_fn():
      x = 1
      if x > 0:
        x = 1
      return x

    node, source = parser.parse_entity(test_fn, future_features=())
    origin_info.resolve(node, source, 'test_file', 100, 0)
    node = tr.visit(node)

    created_pass_node = node.body[1]
    # Takes the line number of the if statement.
    self.assertEqual(
        anno.getanno(created_pass_node, anno.Basic.ORIGIN).loc.lineno, 102)

  def test_origin_info_preserved_in_moved_nodes(self):

    class TestTransformer(transformer.Base):

      def visit_If(self, node):
        return node.body

    tr = TestTransformer(self._simple_context())

    def test_fn():
      x = 1
      if x > 0:
        x = 1
        x += 3
      return x

    node, source = parser.parse_entity(test_fn, future_features=())
    origin_info.resolve(node, source, 'test_file', 100, 0)
    node = tr.visit(node)

    assign_node = node.body[1]
    aug_assign_node = node.body[2]
    # Keep their original line numbers.
    self.assertEqual(
        anno.getanno(assign_node, anno.Basic.ORIGIN).loc.lineno, 103)
    self.assertEqual(
        anno.getanno(aug_assign_node, anno.Basic.ORIGIN).loc.lineno, 104)

class CodeGeneratorTest(test.TestCase):

  def _simple_context(self):
    entity_info = transformer.EntityInfo(
        name='test_fn',
        source_code=None,
        source_file=None,
        future_features=(),
        namespace=None)
    return transformer.Context(entity_info, None, None)

  def test_basic_codegen(self):

    class TestCodegen(transformer.CodeGenerator):

      def visit_Assign(self, node):
        self.emit(parser.unparse(node, include_encoding_marker=False))
        self.emit('\n')

      def visit_Return(self, node):
        self.emit(parser.unparse(node, include_encoding_marker=False))
        self.emit('\n')

      def visit_If(self, node):
        self.emit('if ')
        # This is just for simplifity. A real generator will walk the tree and
        # emit proper code.
        self.emit(parser.unparse(node.test, include_encoding_marker=False))
        self.emit(' {\n')
        self.visit_block(node.body)
        self.emit('} else {\n')
        self.visit_block(node.orelse)
        self.emit('}\n')

    tg = TestCodegen(self._simple_context())

    def test_fn():
      x = 1
      if x > 0:
        x = 2
        if x > 1:
          x = 3
      return x

    node, source = parser.parse_entity(test_fn, future_features=())
    origin_info.resolve(node, source, 'test_file', 100, 0)
    tg.visit(node)

    r = re.compile('.*'.join([
        r'x = 1',
        r'if \(?x > 0\)? {',
        r'x = 2',
        r'if \(?x > 1\)? {',
        r'x = 3',
        r'} else {',
        r'}',
        r'} else {',
        r'}',
        r'return x']), re.DOTALL)

    self.assertRegex(tg.code_buffer, r)
    # TODO(mdan): Test the source map.

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


'''Train CIFAR10 with PyTorch.'''
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import torch.backends.cudnn as cudnn

import torchvision
import torchvision.transforms as transforms

import os
import argparse

from models import *
from utils import progress_bar

parser = argparse.ArgumentParser(description='PyTorch CIFAR10 Training')

# data path
parser.add_argument('--data_path', default='./data', type=str, help='data folder path')

# training parameter
parser.add_argument('--lr', default=0.1, type=float, help='learning rate')
parser.add_argument('--resume', '-r', action='store_false',
                    help='resume from checkpoint')

args = parser.parse_args()

device = 'cuda' if torch.cuda.is_available() else 'cpu'
best_acc = 0  # best test accuracy
start_epoch = 0  # start from epoch 0 or last checkpoint epoch

# Data
print('==> Preparing data..')
transform_train = transforms.Compose([
    transforms.RandomCrop(32, padding=4),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])

transform_test = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])

trainset = torchvision.datasets.CIFAR10(
    root=args.data_path, train=True, download=True, transform=transform_train)
trainloader = torch.utils.data.DataLoader(
    trainset, batch_size=1024, shuffle=True, num_workers=2)

testset = torchvision.datasets.CIFAR10(
    root=args.data_path, train=False, download=True, transform=transform_test)
testloader = torch.utils.data.DataLoader(
    testset, batch_size=1024, shuffle=False, num_workers=2)

classes = ('plane', 'car', 'bird', 'cat', 'deer',
           'dog', 'frog', 'horse', 'ship', 'truck')

# Model
print('==> Building model..')
net = ResNet18()

net = net.to(device)
if device == 'cuda':
    net = torch.nn.DataParallel(net)
    cudnn.benchmark = True

if args.resume:
    # Load checkpoint.
    print('==> Resuming from checkpoint..')
    assert os.path.isdir('checkpoint'), 'Error: no checkpoint directory found!'
    checkpoint = torch.load('./checkpoint/ckpt.pth')
    net.load_state_dict(checkpoint['net'])
    # best_acc = checkpoint['acc']
    best_acc = 0
    start_epoch = checkpoint['epoch']

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=args.lr,
                      momentum=0.9, weight_decay=5e-4)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=200)

upper = torch.tensor([1-0.4914, 1-0.4822, 1-0.4465])
upper /= torch.tensor([0.2023, 0.1994, 0.2010])
lower = torch.tensor([-0.4914, -0.4822, -0.4465])
lower /= torch.tensor([0.2023, 0.1994, 0.2010])

def fgsm_attack(image, epsilon, data_grad):
    # Collect the element-wise sign of the data gradient
    sign_data_grad = data_grad.sign()
    # Create the perturbed image by adjusting each pixel of the input image
    perturbed_image = image + epsilon*sign_data_grad
    # Adding clipping to maintain in the range
    # transpose the batch and channel dimension
    perturbed_image = perturbed_image.transpose(0,1)
    # perturbed_image = perturbed_image.detach().cpu().numpy()
    for channel, _ in enumerate(perturbed_image):
        u_limit, l_limit = upper[channel], lower[channel]
        # perturbed_image[channel] = torch.clamp(perturbed_image[channel].data, l_limit, u_limit)
        perturbed_image[channel] = perturbed_image[channel].data.clamp(l_limit, u_limit)
    perturbed_image = perturbed_image.transpose(0,1)
    # Return the perturbed image
    return perturbed_image

def test_fgsm(model, device, data, target, epsilon=0.05):

    # Send the data and label to the device
    #data, target = data.to(device), target.to(device)
    # Set requires_grad attribute of tensor. Important for Attack
    data.requires_grad = True
    # Forward pass the data through the model
    output = model(data)

    # Calculate the loss
    loss = criterion(output, target)
    # Zero all existing gradients
    model.zero_grad()
    # Calculate gradients of model in backward pass
    loss.backward()

    # Collect datagrad
    data_grad = data.grad.data
    # Call FGSM Attack
    perturbed_data = fgsm_attack(data, epsilon, data_grad)
    #perturbed_data = data
    # Return the accuracy and an adversarial example

    return perturbed_data

# train the model
def train_one_epoch(epoch):
    print('\nEpoch: %d' % epoch)
    net.train()
    train_loss, correct, total, correct_adv, total_adv = 0,0,0,0,0
    for batch_idx, (inputs, targets) in enumerate(trainloader):
        # normal training
        inputs, targets = inputs.to(device), targets.to(device)
        optimizer.zero_grad()
        outputs = net(inputs)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()

        train_loss += loss.item()
        _, predicted = outputs.max(1)
        total += targets.size(0)
        correct += predicted.eq(targets).sum().item()

        # adversarial training
        # adv_input = inputs
        adv_input = test_fgsm(net, device, inputs, targets)#.to(device)
        adv_input = adv_input.clone().detach().to(device)
        optimizer.zero_grad()
        outputs = net(adv_input)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()

        train_loss += loss.item()
        _, predicted = outputs.max(1)
        total_adv += targets.size(0)
        correct_adv += predicted.eq(targets).sum().item()

        progress_bar(batch_idx, len(trainloader), 'Loss: %.3f | Acc: %.3f%% (%d/%d)'
                     % (train_loss/(batch_idx+1), 100.*correct/total, correct, total))
        progress_bar(batch_idx, len(trainloader), 'Loss: %.3f | Acc: %.3f%% (%d/%d)'
                     % (train_loss/(batch_idx+1), 100.*correct_adv/total_adv, correct_adv, total_adv))

# test the accuracy of the model
def test(epoch):
    global best_acc
    #net.eval()
    test_loss = 0
    correct = 0
    total = 0
    total_adv, correct_adv = 0,0
    # with torch.no_grad():
    for batch_idx, (inputs, targets) in enumerate(testloader):
        inputs, targets = inputs.to(device), targets.to(device)
        outputs = net(inputs)
        loss = criterion(outputs, targets)

        test_loss += loss.item()
        _, predicted = outputs.max(1)
        total += targets.size(0)
        correct += predicted.eq(targets).sum().item()

        # adversarial training
        # adv_input = inputs
        adv_input = test_fgsm(net, device, inputs, targets)#.to(device)
        adv_input = adv_input.clone().detach().to(device)
        optimizer.zero_grad()
        outputs = net(adv_input)
        loss = criterion(outputs, targets)

        test_loss += loss.item()
        _, predicted = outputs.max(1)
        total_adv += targets.size(0)
        correct_adv += predicted.eq(targets).sum().item()

        progress_bar(batch_idx, len(testloader), 'Loss: %.3f | Acc: %.3f%% (%d/%d)'
                        % (test_loss/(batch_idx+1), 100.*correct/total, correct, total))
        progress_bar(batch_idx, len(testloader), 'Loss: %.3f | Acc: %.3f%% (%d/%d)'
        % (test_loss/(batch_idx+1), 100.*correct_adv/total_adv, correct_adv, total_adv))

    # Save checkpoint.
    acc = 100.*correct/total
    acc += 100.*correct_adv/total_adv
    acc = acc/2
    if acc > best_acc:
        print('Saving..')
        state = {
            'net': net.state_dict(),
            'acc': acc,
            'epoch': epoch,
        }
        if not os.path.isdir('checkpoint'):
            os.mkdir('checkpoint')
        torch.save(state, './checkpoint/ckpt_adv_scratch.pth')
        best_acc = acc

for epoch in range(start_epoch, start_epoch+200):
    train_one_epoch(epoch)
    test(epoch)
    scheduler.step()


import requests
import json

# id = input("Enter student id: ")
# URL = "http://127.0.0.1:8000/stuinfo/"+id

# r = requests.get(url = URL)

# print(r.json())
URL = "http://127.0.0.1:8000/stucreate/"
data = {
	'name': 'kfjlsadkjf',
	'roll': 1983,
	'city': 'jamalpur'
}
json_data = json.dumps(data)
r = requests.post(url = URL, data = json_data)

data = r.json()
print(data)


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
lines_bars_and_markers/using_span_where.py

Matplotlib > Gallery > Lines, bars and markers > Using span_where
https://matplotlib.org/gallery/lines_bars_and_markers/span_regions.html#using-span-where
"""

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.collections as collections

t  = np.arange(0.0, 2, 0.01)
s1 = np.sin(2*np.pi*t)
# This line exists in the example, but it's not used.
#s2 = 1.2*np.sin(4*np.pi*t)

fig, ax = plt.subplots()
ax.set_title('using span_where')
ax.plot(t, s1, color='black')       # Plot s1
ax.axhline(0, color='black', lw=2)  # Draw a horizontal line for y=0

# Green shady area is applied where s1>0
collection = collections.BrokenBarHCollection.span_where(
    t, ymin=0, ymax=1, where=s1 > 0, facecolor='green', alpha=0.5)
ax.add_collection(collection)

# Green shady area is applied where s1<0
collection = collections.BrokenBarHCollection.span_where(
    t, ymin=-1, ymax=0, where=s1 < 0, facecolor='red', alpha=0.5)
ax.add_collection(collection)

plt.show()


import pytest
from sqlite_utils.utils import column_affinity

EXAMPLES = [
    # Examples from https://www.sqlite.org/datatype3.html#affinity_name_examples
    ("INT", int),
    ("INTEGER", int),
    ("TINYINT", int),
    ("SMALLINT", int),
    ("MEDIUMINT", int),
    ("BIGINT", int),
    ("UNSIGNED BIG INT", int),
    ("INT2", int),
    ("INT8", int),
    ("CHARACTER(20)", str),
    ("VARCHAR(255)", str),
    ("VARYING CHARACTER(255)", str),
    ("NCHAR(55)", str),
    ("NATIVE CHARACTER(70)", str),
    ("NVARCHAR(100)", str),
    ("TEXT", str),
    ("CLOB", str),
    ("BLOB", bytes),
    ("REAL", float),
    ("DOUBLE", float),
    ("DOUBLE PRECISION", float),
    ("FLOAT", float),
    # Numeric, treated as float:
    ("NUMERIC", float),
    ("DECIMAL(10,5)", float),
    ("BOOLEAN", float),
    ("DATE", float),
    ("DATETIME", float),
]

@pytest.mark.parametrize("column_def,expected_type", EXAMPLES)
def test_column_affinity(column_def, expected_type):
    assert expected_type is column_affinity(column_def)

@pytest.mark.parametrize("column_def,expected_type", EXAMPLES)
def test_columns_dict(fresh_db, column_def, expected_type):
    fresh_db.execute("create table foo (col {})".format(column_def))
    assert {"col": expected_type} == fresh_db["foo"].columns_dict


from reportlab.lib.units import inch, cm, pica

class Length:
    def __init__(self, size, unit):
        self.size = size
        if unit in ['inch', 'in']:
            self.unit = 'in'
            self._scale = inch
        elif unit in ['cm', 'mm']:
            self.unit = 'cm'
            self._scale = cm
            if unit == 'mm':
                self.size = self.size / 10
        elif unit in ['pica']:
            self.unit = 'pica'
            self._scale = pica
        else:
            self.unit = 'pt'
            self._scale = 1

    def __call__(self):
        return self.size * self._scale

    def __repr__(self):
        return '{} {}'.format(self.size, self.unit)

class Page:
    def __init__(self, width, height, unit='pt'):
        self.width = Length(width, unit)
        self.height = Length(height, unit)

    @property
    def size(self):
        return (self.width(), self.height())

    def flip(self):
        self.width, self.height = self.height, self.width

def get_paper(size_name, orientation='landscape'):
    if size_name in ['US Letter']:
        paper = Page(11, 8.5, 'in')
    elif size_name in ['US Legal']:
        paper = Page(14, 8.5, 'in')
    elif size_name in ['US Tabloid']:
        paper = Page(17, 11, 'in')
    if orientation == 'portrait':
        paper.flip()
    return paper


import logging
import sys

class StreamToLogger(object):
    """
    Fake file-like stream object that redirects writes to a logger instance.
    """

    def __init__(self, logger, default_output, log_level=logging.INFO, only_on_error=True):
        self.logger = logger
        self.log_level = log_level
        self.default_output = default_output
        self.linebuf = ''
        self.only_on_error = only_on_error

    def write(self, buf):
        if buf != '':
            self.linebuf += buf

    def flush(self):
        if self.only_on_error and 'Traceback (most recent call last)' in self.linebuf:
            for buf in self.linebuf.split('\n'):
                if buf != '':
                    self.logger.log(self.log_level, buf.rstrip())
        else:
            self.default_output.write(self.linebuf)
            self.default_output.flush()
        self.linebuf = ''

def enable(stdout=False, stderror=True, stdout_loglevel=logging.INFO, stderr_loglevel=logging.ERROR):
    if stdout:
        stdout_logger = logging.getLogger('STDOUT')
        sl = StreamToLogger(stdout_logger, sys.stdout, stdout_loglevel)
        sys.stdout = sl

    if stderror:
        stderr_logger = logging.getLogger('STDERR')
        sl = StreamToLogger(stderr_logger, sys.stderr, stderr_loglevel)
        sys.stderr = sl


from torch import nn
import torch

class Encoder(nn.Module):
    def __init__(self, num_channel=3, num_feature=64, latent_dim=100, data_parallel=True):
        super(Encoder, self).__init__()
        features = nn.Sequential(
            # input is (num_channel) x 64 x 64
            nn.Conv2d(num_channel, num_feature, 4, 2, 1, bias=False),
            nn.LeakyReLU(0.2, inplace=True),
            # state size. (num_feature) x 32 x 32
            nn.Conv2d(num_feature, num_feature * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(num_feature * 2),
            nn.LeakyReLU(0.2, inplace=True),
            # state size. (num_feature*2) x 16 x 16
            nn.Conv2d(num_feature * 2, num_feature * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(num_feature * 4),
            nn.LeakyReLU(0.2, inplace=True),
            # state size. (num_feature*4) x 8 x 8
            nn.Conv2d(num_feature * 4, num_feature * 8, 4, 2, 1, bias=False),
            nn.BatchNorm2d(num_feature * 8),
            nn.LeakyReLU(0.2, inplace=True),
        )
        z_mean_map = nn.Sequential(
            # state size. (num_feature*8) x 4 x 4
            nn.Conv2d(num_feature * 8, latent_dim, 4, 1, 0, bias=True),
        )
        z_log_sigma_map = nn.Sequential(
            # state size. (num_feature*8) x 4 x 4
            nn.Conv2d(num_feature * 8, latent_dim, 4, 1, 0, bias=True),
        )
        if data_parallel:
            self.features = nn.DataParallel(features)
            self.z_mean_map = nn.DataParallel(z_mean_map)
            self.z_log_sigma_map = nn.DataParallel(z_log_sigma_map)
        else:
            self.features = features
            self.z_mean_map = z_mean_map
            self.z_log_sigma_map = z_log_sigma_map
        for name, param in self.named_parameters():
            if 'conv' in name and 'weight' in name:
                nn.init.xavier_normal_(param.data)
            elif 'norm' in name and 'weight' in name:
                param.data.fill_(1)
            elif 'norm' in name and 'bias' in name:
                param.data.fill_(0)

    def forward(self, x):
        x = self.features(x)
        mu = self.z_mean_map(x)
        log_sigma = self.z_mean_map(x)
        sigma = torch.exp(log_sigma)
        std_z = torch.randn(mu.size())
        if mu.is_cuda:
            std_z = std_z.cuda()
        z_sample = mu + std_z * sigma
        return mu, log_sigma, sigma, z_sample


import sys
import requests
import json

from django.core.mail import EmailMultiAlternatives
from django.template import RequestContext, TemplateDoesNotExist
from django.template.loader import render_to_string
from django.conf import settings

def send_mail(subject, email_template_name,
              context, to_email, html_email_template_name=None, request=None, from_email=None):
    """
    Sends a django.core.mail.EmailMultiAlternatives to `to_email`.
    """

    print(f'''
    subject: {subject}
    html email template name: {html_email_template_name}
    context: {context}
    to email: {to_email}
    request: {request}
    from email: {from_email}
    ''')
    ctx_dict = {}
    if request is not None:
        ctx_dict = RequestContext(request, ctx_dict)
    # update ctx_dict after RequestContext is created
    # because template context processors
    # can overwrite some of the values like user
    # if django.contrib.auth.context_processors.auth is used
    if context:
        ctx_dict.update(context)

    # Email subject *must not* contain newlines
    from_email = from_email or getattr(settings, 'DEFAULT_FROM_EMAIL')
    if email_template_name:
        message_txt = render_to_string(email_template_name,
                                       ctx_dict)

        email_message = EmailMultiAlternatives(subject, message_txt,
                                               from_email, to_email)
    else:
        try:
            message_html = render_to_string(
                html_email_template_name, ctx_dict)
            email_message = EmailMultiAlternatives(subject, message_html,
                                                   from_email, to_email)
            email_message.content_subtype = 'html'
        except TemplateDoesNotExist:
            pass

    try:
        email_message.send()
    except Exception as e:
        if settings.DEBUG:
            print(f'ERROR: email not sent (utils/utils.py). Reason: {e}')
            print(sys.exc_info())

def check_captcha(request):
    client_key = request.POST['g-recaptcha-response']
    secret_key = settings.RECAPTCHA_PRIVATE_KEY

    captcha_data = {
        'secret':secret_key,
        'response':client_key
    }

    r = requests.post('https://www.google.com/recaptcha/api/siteverify',data=captcha_data)
    response = json.loads(r.text)
    verify = response['success']
    return verify


from typing import Iterable

import pygame
import math

from objects import Ball
from player import Player

SCREEN_X = 1920
SCREEN_Y = 1080
WALL_HEIGHT = 25000

TILE_SIZE = 60
QUALITY1 = TILE_SIZE / 16
QUALITY2 = TILE_SIZE / 256
FOV_MULTIP = 1

REACH = TILE_SIZE / 1

"""
class Player:
    rect: pygame.Rect

    def draw(self, screen: pygame.Surface):
        pygame.draw.rect(screen, (0, 0, 255), self.rect)

    def __init__(self, x, y):
        self.rect = pygame.Rect(0, 0, TILE_SIZE // 2, TILE_SIZE // 2)
        self.rect.center = (x, y)
"""

class GridList:
    tw: int = TILE_SIZE
    columns: int = 32
    rows: int = 0
    offset_x: int = SCREEN_X // 2
    grid: list

    def __init__(self):

        # MAP_1

        self.grid = [
            1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
            1, 0, 0, 0, 0, 0, 0, 0, 3, 0, 0, 0, 0, 3, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 1,
            1, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1,
            1, 0, 0, 0, 3, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1,
            1, 0, 0, 3, 1, 3, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 1,
            1, 0, 0, 0, 3, 0, 0, 0, 3, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 1,
            1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 3, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1,
            1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 1,
            1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 1, 1, 0, 0, 1, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 1,
            1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 1,
            1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1,
            1, 1, 1, 1, 1, 0, 0, 1, 1, 1, 3, 1, 1, 1, 1, 3, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 1,
            1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 1,
            1, 0, 0, 0, 3, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 1,
            1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 1,
            1, 1, 1, 1, 1, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 1,
            1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 1, 1, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 1,
            1, 2, 0, 0, 3, 0, 0, 1, 1, 1, 3, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 1,
            1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 1, 1, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 1,
            1, 1, 1, 1, 1, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 1, 0, 0, 1, 1, 1, 1, 1,
            1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 1, 0, 0, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 1,
            1, 0, 0, 0, 3, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 3, 0, 0, 0, 9,
            1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 1,
            1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
        ]
        self.rows = len(self.grid) // self.columns

    def get_spawn(self) -> tuple:
        index = 0
        for row in range(self.rows):
            for column in range(self.columns):
                if self.grid[index] == 2:
                    return column * self.tw + (self.tw // 2), row * self.tw + (self.tw // 2)
                index += 1

    def get_objects(self) -> Iterable[tuple]:
        index = 0
        for row in range(self.rows):
            for column in range(self.columns):
                if self.grid[index] == 4:
                    yield column * self.tw + (self.tw // 2), row * self.tw + (self.tw // 2)
                index += 1

    def get_tile(self, x, y):
        grid_x = int(math.floor(x) // TILE_SIZE)
        grid_y = int(math.floor(y) // TILE_SIZE)
        index = grid_x
        index += grid_y * SCREEN_X // TILE_SIZE
        return self.grid[index]

    def replace_tile(self, x, y, new_tile):
        grid_x = int(x // TILE_SIZE)
        grid_y = int(y // TILE_SIZE)
        index = grid_x
        index += grid_y * SCREEN_X // TILE_SIZE
        self.grid[index] = new_tile

    def draw(self, screen: pygame.Surface):
        index = 0
        for row in range(self.rows):
            y = row * self.tw
            for column in range(self.columns):
                x = column * self.tw

                if self.grid[index] == 1:
                    pygame.draw.rect(screen, (128, 128, 128), (x, y, self.tw, self.tw))

                elif self.grid[index] == 2:
                    pygame.draw.rect(screen, (0, 0, 0), (x, y, self.tw, self.tw))

                elif self.grid[index] == 3:
                    pygame.draw.rect(screen, (0, 192, 0), (x, y, self.tw, self.tw))

                elif self.grid[index] == 9:
                    pygame.draw.rect(screen, (255, 255, 0), (x, y, self.tw, self.tw))

                index += 1

def main():
    pygame.init()
    pygame.font.init()
    pygame.mixer.init()

    grid = GridList()
    screen = pygame.display.set_mode((SCREEN_X, SCREEN_Y), pygame.FULLSCREEN | pygame.HWACCEL | pygame.DOUBLEBUF)
    raster_screen = screen.copy()
    # raster_screen.set_alpha(50)
    map_screen = pygame.Surface((grid.columns * TILE_SIZE, grid.rows * TILE_SIZE))
    map_zone = pygame.Rect(0, 0, 1200, 1200)
    x_limit, y_limit = map_screen.get_size()
    px, py = grid.get_spawn()
    pdir = 0
    player = Player(px, py)

    pelf = 0

    objects = list()
    for x, y in grid.get_objects():
        obj = Ball(x, y)
        objects.append(obj)

    clock = pygame.time.Clock()
    pressed_keys = set()
    relmouse = 0, 0
    mousewheel = 0
    pygame.event.set_grab(True)
    pygame.mouse.set_visible(False)

    wall_red = 0
    wall_grn = 0
    wall_blu = 0

    running = True
    while running:
        frame_delay = clock.tick(60) / 1000
        running = not bool(pygame.event.get(pygame.QUIT))

        relmouse = 0, 0
        mousewheel = 0

        for event in pygame.event.get():
            if event.type == pygame.KEYDOWN:
                pressed_keys.add(event.key)
            elif event.type == pygame.KEYUP:
                pressed_keys.remove(event.key)
            if event.type == pygame.MOUSEMOTION:
                relmouse = event.rel
            if event.type == pygame.MOUSEWHEEL:
                mousewheel = event.y

        running = not pygame.K_BACKSPACE in pressed_keys and not pygame.K_ESCAPE in pressed_keys

        player.turn(relmouse[0], frame_delay)

        player.turn(
            15 if pygame.K_RIGHT in pressed_keys else
            -15 if pygame.K_LEFT in pressed_keys else
            0,
            frame_delay
        )

        player.change_turn_speed(mousewheel * 0.001)
        player.move(
            forward=pygame.K_w in pressed_keys,
            backward=pygame.K_s in pressed_keys,
            strafe_right=0,
            strafe_left=0
        )

        player.update(grid, frame_delay)

        if pygame.K_LSHIFT not in pressed_keys:
            player.move(
                forward=pygame.K_w in pressed_keys,
                backward=pygame.K_s in pressed_keys,
                strafe_right=0,
                strafe_left=0
            )

        player.update(grid, frame_delay)

        player.change_turn_speed(mousewheel * 0.001)
        player.move(
            forward=0,
            backward=0,
            strafe_right=pygame.K_d in pressed_keys,
            strafe_left=pygame.K_a in pressed_keys
        )

        player.update(grid, frame_delay)

        if pygame.K_LSHIFT not in pressed_keys:
            player.move(
                forward=0,
                backward=0,
                strafe_right=pygame.K_d in pressed_keys,
                strafe_left=pygame.K_a in pressed_keys
            )

            player.update(grid, frame_delay)

        current_x, current_y = player.get_location()
        pdir = player.get_dir()

        if pygame.K_e in pressed_keys:
            pelf = 1
            temp1 = current_x
            temp2 = current_y
            temp1 += math.cos(pdir) * REACH
            temp2 += math.sin(pdir) * REACH

            tile = grid.get_tile(current_x + math.cos(pdir) * REACH, current_y + math.sin(pdir) * REACH)

            if tile == 3:
                grid.replace_tile(temp1, temp2, 0)
            elif tile == 9:
                running = 0
                print("you win i guess")

        raster_screen.fill((0, 0, 0))
        line_x = -8
        angle: float = -0.60
        while angle <= 0.60:
            angle += 0.01
            line_x += 16
            rx = current_x
            ry = current_y
            dist = 0
            rdir: float = (angle * FOV_MULTIP + pdir)
            wall_red = 0
            ray_run = 1

            dx = math.cos(angle) * QUALITY1
            while ray_run == 1:
                rx += math.cos(rdir) * QUALITY1
                ry += math.sin(rdir) * QUALITY1
                dist += dx

                if (grid.get_tile(rx, ry)) == 1:
                    ray_run = 0
                elif (grid.get_tile(rx, ry)) == 3:
                    ray_run = 0
                elif (grid.get_tile(rx, ry)) == 9:
                    ray_run = 0

            ray_run = 1

            dx = math.cos(angle) * QUALITY2
            tile = (grid.get_tile(rx, ry))
            if tile == 1 or tile == 3 or tile == 9:
                while ray_run == 1:
                    rx -= math.cos(rdir) * QUALITY2
                    ry -= math.sin(rdir) * QUALITY2
                    dist -= dx

                    tile = (grid.get_tile(rx, ry))
                    if tile != 1 and tile != 3 and tile != 9:
                        ray_run = 0
                    else:
                        ray_run = 1

            rx += math.cos(rdir) * QUALITY2
            ry += math.sin(rdir) * QUALITY2
            dist += QUALITY2

            if (grid.get_tile(rx, ry)) == 1 or 3 or 9:
                if dist <= 0:
                    dist = 1
                center_y = SCREEN_Y / 2
                wall_height = (1 / dist) * WALL_HEIGHT

                if (grid.get_tile(rx, ry)) == 1:
                    wall_red = 0
                    wall_grn = 0
                    wall_blu = 255
                elif (grid.get_tile(rx, ry)) == 3:
                    wall_red = 0
                    wall_grn = 192
                    wall_blu = 0
                elif (grid.get_tile(rx, ry)) == 9:
                    wall_red = 255
                    wall_grn = 255
                    wall_blu = 0

                pygame.draw.line(
                    raster_screen,
                    (wall_red, wall_grn, wall_blu),
                    (line_x, center_y - wall_height),
                    (line_x, center_y + wall_height),
                    width=16
                )

        for obj in objects:
            obj.draw(raster_screen, current_x, current_y, 0 - pdir)

        map_screen.fill((0, 0, 0))
        screen.fill((0, 0, 0))
        grid.draw(map_screen)
        player.draw(map_screen)

        map_zone.center = player.get_location()
        if map_zone.left < 0:
            map_zone.left = 0
        if map_zone.top < 0:
            map_zone.top = 0
        if map_zone.right > x_limit:
            map_zone.right = x_limit
        if map_zone.bottom > y_limit:
            map_zone.bottom = y_limit

        if pygame.K_m in pressed_keys:
            _map_screen = pygame.transform.scale(map_screen.subsurface(map_zone), (600, 600))
        else:
            _map_screen = pygame.transform.scale(map_screen.subsurface(map_zone), (300, 300))

        screen.blit(raster_screen, (0, 0))
        screen.blit(_map_screen, (0, 0))
        pygame.display.flip()

if __name__ == '__main__':
    main()


import cv2
import omegaconf

import numpy as np
import tensorflow as tf
import tensorflow.keras as tfk

from imgaug import augmenters as iaa
from tensorflow.keras.utils import Sequence
from typing import List

class DataGenerator(Sequence):
    def __init__(self, paths: List,
                labels: List, steps: int):
        self.steps = steps
        self.x_path = paths
        self.batch_size = 32
        self.target_size = (24, 24)
        self.labels = labels

    def __bool__(self):
        return True

    def __len__(self):
        return self.steps

    def __getitem__(self, idx):
        batch_x_path = self.x_path[idx * self.batch_size: (1 + idx) * self.batch_size]
        batch_x = np.asarray([self.load_image(x_path) for x_path in batch_x_path])
        batch_y = self.labels[idx * self.batch_size: (1 + idx) * self.batch_size]
        batch_y = tfk.utils.to_categorical(np.asarray(batch_y), num_classes=4)
        return np.expand_dims(batch_x, axis=-1), batch_y

    def load_image(self, x_path):
        image = self.read_png(x_path)
        if image.shape != self.target_size:
            image_array = cv2.resize(image, self.target_size)
        else:
            image_array = image
        return image_array

    @staticmethod
    def read_png(x_path):
        image = tf.image.decode_png(tf.io.read_file(x_path))
        gray_img = cv2.cvtColor(np.asarray(image), cv2.COLOR_BGR2GRAY)
        return gray_img / 255

    @staticmethod
    def transform_batch_images(batch_x):
        sometimes = lambda aug: iaa.Sometimes(1.0, aug)
        augmenter = iaa.Sequential(
                            [
                              iaa.Fliplr(0.5),
                              iaa.Flipud(0.5),
#                               iaa.geometric.Affine(rotate=(-20, 20), order=1, mode='constant', fit_output=False),
#                               sometimes(iaa.Affine(translate_percent={"x": (-0.2, 0.2), "y": (-0.2, 0.2)}, shear=(-16, 16))),
#                               sometimes(iaa.Crop(px=(0, 25), keep_size=True, sample_independently=False)),
#                               sometimes(iaa.LinearContrast(0.5, 0.75))
                            ],
                            random_order=True,
                        )

        batch_x = augmenter.augment_images(batch_x)
        return batch_x


from flask import current_app as app

class Current_Balance:
    def __init__(self, balance):
        self.balance = balance

class Balance:
    def __init__(self, trans_id, trans_date, trans_ts, user_id, trans, trans_description, balance):
        self.trans_id = trans_id
        self.trans_date = trans_date
        self.trans_ts = trans_ts
        self.user_id = user_id
        self.trans = trans
        self.trans_description = trans_description
        self.balance = balance

    @staticmethod
    def get(user_id):
        rows = app.db.execute('''
SELECT current_balance
FROM Current_Balance
WHERE user_id = :user_id
''',
                              user_id=user_id)

        return Current_Balance(*rows[0])

    @staticmethod
    def get_lastrow():
        rows = app.db.execute('''
SELECT *
FROM Balance
''')
        return Balance(*(rows[-1])) if rows is not None else None

    @staticmethod
    def get_all(user_id):
        rows = app.db.execute('''
SELECT *
FROM Balance
WHERE user_id = :user_id
ORDER BY trans_date DESC
''',
                        user_id=user_id)
        return [Balance(*row) for row in rows]

class Update_balance:
    def __init__(self, trans_id, trans_date, user_id, trans, trans_description, balance):
        self.trans_id = trans_id
        self.trans_date = trans_date
        self.user_id = user_id
        self.trans = trans
        self.trans_description = trans_description
        self.balance = balance

    @staticmethod
    def insert(trans_date, user_id, trans, trans_description, balance):
        rows = app.db.execute('''
INSERT INTO Balance(trans_date, user_id, trans, trans_description, balance)
VALUES (:trans_date, :user_id, :trans, :trans_description, :balance)
''',
                              trans_date = trans_date,
                              user_id = user_id,
                              trans = trans,
                              trans_description=trans_description,
                              balance = balance)

        rows = app.db.execute('''
UPDATE Current_balance
SET current_balance = :balance
WHERE user_id=:user_id
''',
                            user_id=user_id,
                            balance=balance)

        return "Success"


import torch
import torch.nn as nn
from torch.distributions.categorical import Categorical
import numpy as np

class GFN(nn.Module):
    def __init__(self, m, learn_Z=False, architecture=None, h=None):
        """
        :param architecture: nn.Module with input layer of size m and output layer of size m+1
        :param m: input size
        :param learn_Z: if True, then an extra parameter needs to be learned logZ
        :param h: if architecture is not specified then it's the number of hidden units for each of the 2 hidden layers
        """
        super().__init__()
        self.m = m
        if architecture is None:
            assert h is not None
            architecture = nn.Sequential(nn.Linear(m, h), nn.ReLU(), nn.Linear(h, h), nn.ReLU(), nn.Linear(h, m + 1))
        self.logit_maker = architecture
        if learn_Z:
            self.logZ = nn.Parameter(torch.tensor(0.))

    def forward(self, s):
        """
        :param s: (batched) states of the GFN. tensor of size k x m
        :return: P_F(. | s) represented as a tensor of size k x m+1.
        """
        assert s.ndim == 2 and s.shape[1] == self.m
        logits = self.logit_maker(s)
        return logits

class BinaryMaskGFN(GFN):
    """
    The states of the GFN are binary vectors of size m. An action amounts to replacing a zero with a one
    """

    def forward(self, theta):
        """
        :param theta: (batched) states of the GFN, binary tensor of size k x m
        :return: masked out logits: -infinity to invalid actions (those that try to replace a 1 with a 1)
        """
        logits = super().forward(theta)
        logits[torch.cat((theta, torch.zeros(theta.shape[0], 1)), 1) == 1.] = - float("inf")
        return logits

    def on_policy_trajectory_sample(self, temperature=1.):
        """
        :param temperature: softmax temperature
        :return: a complete GFN trajectory (i.e. ending in a terminating state) as a T x m tensor
        The function assumes the source state of the GFN is [0, 0, ..., 0]
        """
        theta = torch.zeros(1, self.m)
        trajectory = [theta.clone().detach()]
        while True:
            logits = self(theta)
            probs = torch.softmax(logits / temperature, 1)
            idx = Categorical(probs).sample()[0]
            if idx == self.m:  # terminating state
                break
            theta[0, idx] = 1.
            trajectory.append(theta.clone().detach())
        return torch.cat(trajectory, 0)

    def uniform_trajectory_sample(self, terminate_prob=0.1):
        """
        :param terminate_prob: probability of ending the trajectory at each step of the trajectory
        :return: a complete GFN trajectory (i.e. ending in a terminating state) as a T x m tensor
        """
        m = self.m
        theta = torch.zeros(1, m)
        trajectory = [theta.clone().detach()]
        while True:
            possible_indices = [i for i, j in enumerate(theta[0].detach().numpy()) if j == 0]
            if len(possible_indices) == 0:
                break
            probs = ((1 - terminate_prob) / len(possible_indices)) * np.ones(len(possible_indices))
            idx = np.random.choice(possible_indices + [m], p=list(probs) + [terminate_prob])
            if idx == m:
                break
            theta[0, idx] = 1.
            trajectory.append(theta.clone().detach())
        return torch.cat(trajectory, 0)

def modified_trajectory_balance_loss(pf, traj, MIN_REW, pb_fn, reward_fn):
    """
    :param pf: nn.Module representing P_F (can be an instantiation of BinaryMaskGFN e.g.)
    :param traj: sampled trajectory
    :param MIN_REW: minimum reward value to replace lower rewards (including -infinity) with
    :param pb_fn: function taking a trajectory (T x m tensor) as input and outputting the backward probs as a (T-1) tensor
    :param reward_fn: function mapping a trajectory to a reward (list of same length)
    :return: loss(traj)
    Works only if all states are connected to the sink state, those that aren't are given an imaginary reward of MIN_REW
    """
    trajectory_logits = pf(traj)
    final_forward_logprob_differences = trajectory_logits[:-1, -1] - trajectory_logits[1:, -1]
    denominators = torch.logsumexp(trajectory_logits, 1)
    final_forward_logprob_differences = final_forward_logprob_differences - (denominators[:-1] - denominators[1:])
    forward_logprobs = trajectory_logits[:-1, :-1][
        torch.arange(traj.shape[0] - 1), (traj[1:, :] - traj[:-1, :]).argmax(1)]  # 1-dim tensor of size T-1
    forward_logprobs = forward_logprobs - denominators[:-1]

    pb = pb_fn(traj)
    log_pb = torch.log(pb)
    pred = final_forward_logprob_differences + log_pb - forward_logprobs

    rewards = reward_fn(traj)
    trajectory_rewards = torch.tensor(rewards).clamp_min(MIN_REW).to(pred.dtype)

    targets = trajectory_rewards[:-1] - trajectory_rewards[1:]
    loss = nn.MSELoss(reduction='sum')(pred, targets)

    return loss, rewards, forward_logprobs.sum()

def trajectory_balance_loss(pf, traj, MIN_REW, pb_fn, reward_fn):
    """
    :param pf: nn.Module representing P_F (can be an instantiation of BinaryMaskGFN e.g.)
    :param traj: sampled trajectory
    :param MIN_REW: minimum reward value to replace lower rewards (including -infinity) with
    :param pb_fn: function taking a trajectory (T x m tensor) as input and outputting the backward probs as a (T-1) tensor
    :param reward_fn: function mapping a trajectory to a reward (list of same length)
    :return: loss(traj)
    needs the GFN to be paramertrized with logZ
    """
    trajectory_logits = pf(traj)
    denominators = torch.logsumexp(trajectory_logits, 1)
    forward_logprobs = trajectory_logits[:-1, :-1][
        torch.arange(traj.shape[0] - 1), (traj[1:, :] - traj[:-1, :]).argmax(1)]  # 1-dim tensor of size length-1
    forward_logprobs = forward_logprobs - denominators[:-1]
    sink_logprob = trajectory_logits[-1, -1] - denominators[-1]
    pb = pb_fn(traj)
    log_pb = torch.log(pb)

    pred = pf.logZ + forward_logprobs.sum() + sink_logprob - log_pb.sum()
    trajectory_rewards = reward_fn(traj[[-1]])  # log-likelihood (maybe + log-prior)
    targets = torch.tensor(trajectory_rewards).clamp_min(MIN_REW).to(pred.dtype)[0]
    loss = nn.MSELoss()(pred, targets)

    trajectory_logprob = forward_logprobs.sum() + sink_logprob

    return loss, trajectory_rewards, trajectory_logprob