from fastapi import APIRouter, UploadFile, File
from PIL import Image
from io import BytesIO
from schemas.item import Item
from src.embedding.embedding import get_single_image_embedding
from database import index
from constants import DIMENSIONS

router = APIRouter(prefix="/items", tags=["items"])

# Retrieve all clothing items in database
@router.get("/", response_model=list[Item])
def get_items() -> list[Item]:
    try:
        items = index.query(
            vector=[0] * DIMENSIONS,
            top_k=1000,
            include_values=False,
            include_metadata=True,
        )

        return items.matches
    except Exception as e:
        print(e)

# Add clothing image, description, tags, and embeddings to database
@router.post("/", status_code=201)
def add_item(item: Item):
    # Converts Pydantic Item model into Python dictionary
    item_dict = item.model_dump(exclude={"score"})
    try:
        index.upsert(vectors=[item_dict])
    except Exception as e:
        print(e)

@router.post("/similar", response_model=list[Item])
async def similar_items(file: UploadFile = File(...)) -> list[Item]:
    try:
        contents = await file.read()
        with Image.open(BytesIO(contents)).convert("RGB") as image:
            embeddings = get_single_image_embedding(image)[0].tolist()
        items = index.query(
            vector=embeddings,
            top_k=1000,
            include_values=False,
            include_metadata=True,
        )

        return items.matches
    except Exception as e:
        print(e)


import argparse
import os
import signal

from .run import run

def main():
    run(**parse_args())

def parse_args():
    parser = argparse.ArgumentParser(
        prog="dump_lammps",
        description="Create lammps data/dump from simulation trajectory",
    )
    parser.add_argument(
        "--phase",
        dest="phasename",
        metavar="<phase>",
        type=str,
        help="simulation phase to dump (initialization | relaxation | simulation)",
    )
    parser.add_argument(
        "--steps",
        dest="step_range",
        metavar="[<start>:]<end>",
        type=str,
        help="step range to dump",
    )
    parser.add_argument(
        "--apart",
        dest="apart",
        metavar="<factor>",
        type=float,
        help="scale the center of mass of each chain by this factor",
    )
    parser.add_argument(
        "--smooth",
        dest="smooth_window",
        metavar="<window>",
        type=int,
        default=None,
        help="apply gaussian smoothing of given window size",
    )
    parser.add_argument(
        "trajfile", metavar="<trajectory>", type=str, help="input hdf5 file",
    )
    parser.add_argument(
        "datafile", metavar="<data>", type=str, help="output LAMMPS data file",
    )
    parser.add_argument(
        "dumpfile", metavar="<dump>", type=str, help="output LAMMPS dump file",
    )
    args = parser.parse_args()

    args.step_range = parse_step_range(args.step_range)

    return vars(args)

def parse_step_range(arg):
    if arg is None:
        return None
    fragments = arg.split(":")
    if len(fragments) == 1:
        return 0, int(arg)
    if len(fragments) == 2:
        return int(fragments[0]), int(fragments[1])
    raise Exception("invalid step range")

if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        signal.signal(signal.SIGINT, signal.SIG_DFL)
        os.kill(os.getpid(), signal.SIGINT)
    except BrokenPipeError:
        signal.signal(signal.SIGPIPE, signal.SIG_DFL)
        os.kill(os.getpid(), signal.SIGPIPE)


from discord.ext import commands
import aiosqlite

import config

class DatabaseSetup(commands.Cog):
    def __init__(self, bot) -> None:
        self.bot = bot
        self.category = "Owner"

    @commands.Cog.listener()
    async def on_ready(self):
        db = await aiosqlite.connect(config.database)
        cursor = await db.cursor()

        await cursor.execute('''
        CREATE TABLE IF NOT EXISTS "Module Channels" (
            guild_id INTEGER NOT NULL UNIQUE,
            PRIMARY KEY("guild_id")
        );
        ''')
        await cursor.execute('''
        CREATE TABLE IF NOT EXISTS GlobalVars (
            property_name TEXT NOT NULL UNIQUE,
            property_value TEXT,
            PRIMARY KEY(property_name)
        );''')
        await cursor.execute('''
        CREATE TABLE IF NOT EXISTS GuildConfig (
            "guild_id" INTEGER NOT NULL UNIQUE,
            config TEXT DEFAULT "{}",
            PRIMARY KEY("guild_id")
        );''')
        await db.commit()
        await db.close()

def setup(bot):
    bot.add_cog(DatabaseSetup(bot))


import re
import datetime

from rest_framework.response import Response
"""Expresiones Regulares"""
patron_id_cliente=re.compile(r"C[0-9]{1,}")
patron_dir=re.compile(r"[A-Za-z]+ [0-9]+")
patron_nombre=re.compile(r"[A-Za-z]+")
patron_id_sol=re.compile(r"S[0-9]{1,}")
patron_legajo=re.compile(r"CH[0-9]{1,}")
patron_remito=re.compile(r"R[0-9]{1,}")
"""Fin ER"""

def validate_crear_solicitud(body, id_cliente, opcion):
    dir_body=body['direccion']
    nombre=body['nombre_2']
    exito=False
    if re.fullmatch(patron_dir, dir_body) and re.fullmatch(patron_id_cliente, id_cliente) \
        and opcion in (0,1) and re.fullmatch(patron_nombre, nombre):
        exito=True
    return exito

def validate_agregar_bultos(body, id_solicitud):
    tipos_validos=('CH', 'M', 'G')
    tipo_bulto=body['tipo']
    exito=False
    if re.fullmatch(patron_id_sol, id_solicitud) and tipo_bulto in tipos_validos:
        exito=True
    return exito

def validate_remitos_chofer_estado(body):
    legajo=body['legajo']
    estado=body['estado']
    estados_validos=['asignado', 'en_circulacion', 'en_viaje', 'rendido', 'segunda_entrega',\
        'pendiente', 'pagado']
    exito=False
    if re.fullmatch(patron_legajo, legajo) and estado in estados_validos:
        exito=True
    return exito

def validate_asociar_sol_remito(body):
    medios_validos=['OR', 'DES', 'CC']
    nro_remito=body['remito']
    id_solicitud=body['sol']
    medio_pago=body['m_pago']
    exito=False
    if re.fullmatch(patron_remito, nro_remito) and re.fullmatch(patron_id_sol, id_solicitud)\
        and medio_pago in medios_validos:
        exito=True
    return exito

def validate_viaje_fecha(fecha):
    format="%Y-%m-%d"
    exito=True
    try:
        datetime.datetime.strptime(fecha, format)
    except ValueError:
        exito=False
    return exito

def validate_asoc_remito_viaje(body):
    remitos=body['remitos']
    exito=True
    index=0
    while exito is True and index < len(remitos):
        if re.fullmatch(patron_remito, remitos[index]) is None:
            exito=False
        index+=1

    return exito


# """
# write python program that reads json file "keys.json" and prints all dublicates that it found
# """
#
# import json
#
# with open('keys.json', 'r') as file:
# 	keys_data = json.load(file)
#
# vk_data = keys_data["virtual-key codes"]
#
# print(f'1: {vk_data}')
#
# all_items = []
# for item in vk_data:
# 	all_items.append(item['name'])
# 	all_items.append(item['numeric'])
# 	all_items.append(item['decimal'])
# 	all_items.append(item['description'])
#
# print(f'2: {all_items}')
#
# dublicates = []
# for item in all_items:
# 	if all_items.count(item) > 1:
# 		dublicates.append(item)
#
# print(f'3: {dublicates}')

import json
import time

def find_subset(value, data, type=None):
    """
    :param value: value to find in data
    :param data: list of dicts
    :param type: type of value to find
    :return: item from data if value is found, None if value is not found
    """

    if type:  # find subsets of value in data by set type
        if type == 'dec':
            type = 'decimal'
        elif type == 'num':
            type = 'numeric'
        elif type == 'desc':
            type = 'description'

        for item in data:
            if value == item[type]:
                return item  # return subset of value in data
    else:  # try to find subsets in all types automatically
        for item in data:  # go through all subsets in keys dataset
            for key in item:
                if value == item[key]:
                    return item  # return item if value is found in any type
            if value == item['name'][3:]:
                return item  # return item if value is equal to name without "VK_"
    return False  # subset not found

def _find_subset(value, data):
    """
    Find subset of dict in list of dicts
    :param value: value to find
    :param data: list of dicts
    :return: found dict, else False
    """

    for item in data:
        if value == item['name']:
            return item
        elif str(value) == item['numeric']:
            return item
        elif value == item['decimal']:
            return item
        elif value == item['description']:
            return item

    return False

with open('keys.json', 'r') as file:
    keys_data = json.load(file)
vk_data = keys_data['virtual-key codes']

# utils.find_subset() took 0.4837017059326172 seconds
# _find_subset() took 0.452298641204834 seconds

value = 'VK_B'

print(f'find_subset() returned {find_subset(value, vk_data)}')
print(f'_find_subset() returned {_find_subset(value, vk_data)}')

# make 10000 loop to test performance of utils.find_subset()
now = time.time()
for i in range(10000):
    find_subset(0x41, vk_data)
print(f'utils.find_subset() took {time.time() - now} seconds')

# make 10000 loop to test performance of _find_subset()
now = time.time()
for i in range(10000):
    _find_subset("0x41", vk_data)
print(f'_find_subset() took {time.time() - now} seconds')


from torch.utils.data import Dataset,random_split
from PIL import Image
import glob
import pickle
import torch
import numpy as np
import time

class EndGameDataset(Dataset):
    """
    Dataset class for Dense Depth and Mask prediction
    """

    def __init__(self, dataset_paths, dataset_root):
        """
        Loads datasetpaths

        params
        -----
        dataset_paths : path to serialized dataset filepaths
        dataset_root : location of dataset images
        """
        dataset_paths_pickle = open(dataset_paths,"rb")
        self.dataset_paths = pickle.load(dataset_paths_pickle)
        self.dataset_root = dataset_root

    def __getitem__(self, idx):
        """
        Returns the current dataset item based on the index
        """
        filepaths = self.dataset_paths[idx]
        fg_bg = self.get_input(filepaths["fg_bg"])
        bg = self.get_input(filepaths["bg"])
        fg_bg_mask = self.get_target(filepaths["fg_bg_mask"])
        dense_depth = self.get_target(filepaths["dense_depth"])
        return {"fg_bg":fg_bg, "bg":bg, "fg_bg_mask":fg_bg_mask, "dense_depth":dense_depth}

    def __len__(self):
        """
        Returns the length of the dataset
        """
        return len(self.dataset_paths)

    def get_input(self, filepath):
      """
      Loads and returns the model input image
      """
      return self.load_file(filepath)

    def get_target(self, filepath):
      """
      Loads and returns the model target image
      """
      target = self.load_file(filepath)
      target = target[:,:,np.newaxis]
      return target

    def load_file(self, img_path):
        """
        Loads images using PIL
        """
        img = Image.open(self.dataset_root+img_path)
        img_nd = np.array(img)
        if(np.max(img_nd)<1e-6):
          img_nd = img_nd
        else:
          img_nd = img_nd/np.max(img_nd)
        return img_nd

    def get_folder_file_paths(self, folderPath):
      """
      Returns filepaths of the given folder
      """
      file_paths = []
      for file in glob.glob(folderPath+"/**/*.jpg", recursive=True):
        file_paths.append(file)
      return file_paths

class EndGameSubset(Dataset):
    """
    Dataset class for Dense Depth and Mask prediction
    Applies data augmentation and returns the dataset
    """
    def __init__(self, subset, transform=None, target_transform=None):
        """
        Initializes the dataset
        """
        self.subset = subset
        self.transform = transform
        self.target_transform = target_transform

    def __getitem__(self, index):
        """
        Returns current item of the dataset based on the index
        Applies the data augmentation for the dataset images
        """
        result = self.subset[index]
        if self.transform:
            result['fg_bg'] = self.transform["fg_bg"](result['fg_bg'])
            result['bg'] = self.transform["bg"](result['bg'])
        if self.target_transform:
          result['fg_bg_mask'] = self.target_transform["fg_bg_mask"](result['fg_bg_mask'])
          result['dense_depth'] = self.target_transform["dense_depth"](result['dense_depth'])
        return result

    def __len__(self):
        """
        Returns length of the dataset
        """
        return len(self.subset)

def EndGameTrainTestDataSet(dataset_paths, dataset_root, train_split = 70, train_transforms=None, train_target_transforms=None, test_transforms=None, test_target_transforms=None):
  """
  Splits the dataset into train and test

  params
  ------
  dataset_paths : path to serialized dataset filepaths
  dataset_root : location of dataset images
  train_split : percentage of train dataset
  train_transform: Data augmentaion for train - input dataset
  train_transform: Data augmentaion for train - target dataset
  test_transforms: Data augmentaion for test - input dataset
  test_target_transforms: Data augmentaion for test - target dataset

  returns
  ------
  train_dataset
  test_dataset
  """
  dataset = EndGameDataset(dataset_paths, dataset_root)
  train_len = len(dataset)*train_split//100
  test_len = len(dataset) - train_len
  train_set, val_set = random_split(dataset, [train_len, test_len])
  train_dataset = EndGameSubset(train_set, transform=train_transforms, target_transform=train_target_transforms)
  test_dataset = EndGameSubset(val_set, transform=test_transforms, target_transform=test_target_transforms)

  return train_dataset, test_dataset


"""
You are given a network of n nodes, labeled from 1 to n. You are also given times,
a list of travel times as directed edges times[i] = (ui, vi, wi), where ui is the
source node, vi is the target node, and wi is the time it takes for a signal to
travel from source to target.

We will send a signal from a given node k. Return the time it takes for all the n
nodes to receive the signal. If it is impossible for all the n nodes to receive the
 signal, return -1.
"""

import heapq as hq
import collections

class Solution(object):
    def networkDelayTime(self, times, n, k):
        """
        :type times: List[List[int]]
        :type n: int
        :type k: int
        :rtype: int
        """
        adj = collections.defaultdict(list)
        # reorder into adj list
        for time in times:
            adj[time[0]].append((time[1], time[2]))

        distances = {}
        pq = []

        hq.heappush(pq, (0, k))

        while len(pq) > 0:
            dist, node = hq.heappop(pq)
            if node not in distances:
                # add distances here (minimum since based on priority)
                distances[node] = dist
                for to, w in adj[node]:
                    if to not in distances:
                        # dont add here, cuz it may not be minimum
                        new_dist = w + dist
                        hq.heappush(pq, (new_dist, to))

        if len(distances) == n:
            return max(distances.values())
        return -1


"""
For each file named training.log in recursive subdirectories, find the last line containing 'Better eval results!'.
Print the file name, and the file lines between 3 lines above and 1 line bellow the line containing 'Better eval results!'.
"""

from pathlib import Path
from os import system

for path in Path('.').rglob('training.log'):
    print(path)
    with open(path) as file:
        lines = file.readlines()
    last_match_line = None
    for i, line in enumerate(lines):
        if "Better eval results!" in line:
            last_match_line = i
    for log_line in lines[max(0, last_match_line-3) : min(last_match_line+2, len(lines))]:
        print(' | ', log_line.rstrip('\n'))


from airflow.models import DAG
from airflow.contrib.operators.spark_submit_operator import SparkSubmitOperator
from datetime import datetime

default_args = {
    'owner': 'airflow',
    'start_date': datetime(2019, 1, 1)
}

dag = DAG('stackoverflow_stats', default_args=default_args,
          schedule_interval='@daily')

SparkSubmitOperator(
    task_id='get-stats',
    application="/usr/local/airflow/jobs/stats.py",
    dag=dag,
    run_as_user='airflow',
    application_args=['--date', '{{ ds }}'],
    name='Stats DAG for {{ ds }}',
    num_executors=2,
    executor_memory='2g'
)


from selenium.webdriver.common.by import By
from Selectors.links import links_selectors

# Attributes
new_tab_tracked_attr = 'noopener noreferrer'
new_tab_not_tracked_attr = 'noopener noreferrer nofollow'
data_attr_tracked = 'tracked'
data_attr_not_tracked = 'not-tracked'

def rel_attributes(driver, article_url):
    print('\n5933 - Set links to "follow" so they can be tracked')
    driver.get(article_url)

    internal_new_tab_tracked = driver.find_element(By.CSS_SELECTOR, links_selectors.internal_new_tab_tracked_selector)
    internal_new_tab_not_tracked = driver.find_element(By.CSS_SELECTOR, links_selectors.internal_new_tab_not_tracked_selector)
    external_new_tab_tracked = driver.find_element(By.CSS_SELECTOR, links_selectors.external_new_tab_tracked_selector)
    external_new_tab_not_tracked = driver.find_element(By.CSS_SELECTOR, links_selectors.external_new_tab_not_tracked_selector)
    internal_same_tab = driver.find_element(By.CSS_SELECTOR, links_selectors.internal_same_tab_selector)
    external_same_tab = driver.find_element(By.CSS_SELECTOR, links_selectors.external_same_tab_selector)

    # Internal and external same tab links should have no rel attribute by default
    try:
        internal_same_tab.get_attribute('rel')
        message = f'> Internal link {internal_same_tab.text} should have no rel attribute'
        raise AssertionError(message)
    except:
        print(' - Internal same tab link correctly has no rel attribute')

    try:
        external_same_tab.get_attribute('rel')
        message = f'> External link {internal_same_tab.text} should have no rel attribute'
        raise AssertionError(message)
    except:
        print(' - External same tab link correctly has no rel attribute')

    if internal_new_tab_tracked.get_attribute('rel') == new_tab_tracked_attr:
        print(f' - Internal new tab tracked link has expected rel attribute {new_tab_tracked_attr}')
    else:
        message = f' - Internal new tab tracked link has incorrect rel attribute, should be {new_tab_tracked_attr}'
        raise AssertionError(message)

    if internal_new_tab_not_tracked.get_attribute('rel') == new_tab_not_tracked_attr:
        print(f' - Internal new tab not tracked link has expected rel attribute {new_tab_not_tracked_attr}')
    else:
        message = f' - Internal new tab not tracked link has incorrect rel attribute, should be {new_tab_not_tracked_attr}'
        raise AssertionError(message)

    if external_new_tab_tracked.get_attribute('rel') == new_tab_tracked_attr:
        print(f' - External new tab tracked link has expected rel attribute {new_tab_tracked_attr}')
    else:
        message = f' - External new tab tracked link has incorrect rel attribute, should be {new_tab_tracked_attr}'
        raise AssertionError(message)

    if external_new_tab_not_tracked.get_attribute('rel') == new_tab_not_tracked_attr:
        print(f' - External new tab not tracked link has expected rel attribute {new_tab_not_tracked_attr}')
    else:
        message = f' - External new tab not tracked link has incorrect rel attribute, should be {new_tab_not_tracked_attr}'
        raise AssertionError(message)

def runTest(baseUrl, driver, browser):
    print(' - Testing links on {}'.format(browser))

    if not "www.moneyhelper.org.uk" in baseUrl:
        if "moneypensions-qa" in baseUrl:
            article_url = f"{baseUrl}/en/jp-test/how-to-choose-the-right-bank-account"
            rel_attributes(driver, article_url)
    else:
        print('- Skipping 5933 rel attribute test, only tests on test and staging')


from learning_rate_tools import find_lr, get_loss_over_lr
from core.model_loader import create_gpt2_model
from core.enums import GPT2Models
import torch
from datasets import load_dataset
from torch.utils.data import DataLoader
from transformers import default_data_collator
from core.prompter import generate_prompt
import json

device = 'cuda' if torch.cuda.is_available() else 'cpu'
model_name = GPT2Models.gpt2
batch_size = 1
smallest_lr = 1e-5
largest_lr = 1
tokenizer, model = create_gpt2_model(model_name, use_lora=False)
model = model.to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=smallest_lr)

dataset = load_dataset("json", data_files="cmd_assistant_data.json")

dataset = dataset["train"].train_test_split(test_size=0.1)
dataset["validation"] = dataset["test"]
del dataset["test"]
dataset["train"][0]
# causing issues for some reasons
# tokenizer.padding_side = 'left'

def preprocess_function(examples):
    conversation = examples['conversation']
    prompt = generate_prompt(conversation)
    model_inputs = tokenizer(prompt, padding="max_length",
                             truncation=True, return_tensors="pt")
    model_inputs["labels"] = model_inputs["input_ids"].clone()
    return model_inputs

encoded = preprocess_function(dataset['train'][0])
processed_datasets = dataset.map(
    preprocess_function,
    batched=False,
    num_proc=1,
    remove_columns=dataset["train"].column_names,
    load_from_cache_file=False,
    desc="Running tokenizer on dataset",
)

train_dataset = processed_datasets["train"]
eval_dataset = processed_datasets["validation"]
train_dataloader = DataLoader(
    train_dataset, shuffle=True, collate_fn=default_data_collator, batch_size=batch_size, pin_memory=True
)

history = get_loss_over_lr(model, train_dataloader,
                           optimizer, smallest_lr, largest_lr)
lower_lr, higher_lr = find_lr(history)
average_lr = (lower_lr + higher_lr) / 2

with open(f'{model_name}_lr_data.json', 'w') as f:
    json.dump(history, f)


from .parse_weights import WeightedPrompt
from dataclasses import dataclass
from .misc import normalize_weights
from ..utils.debug_vram import debug_vram
from ..utils.load_tensor import load as load_tensor
from diffusers import StableDiffusionPipeline
from functools import partial
import numpy as np
import torch, json
from torch import nn
from os import path
from glob import glob
from threading import Lock
from ..utils.candy import locked
from ..utils.misc import normalize_resolution
from PIL import Image
from math import log, ceil, sqrt
from ..utils.lcs import LCS
import os
from typing import Callable, Iterable, Dict, Union, List
import safetensors
from .model_mgr import LoRA
from ..utils.misc import DEVICE as DEFAULT_DEVICE
def combine_noise(noise: torch.tensor, use_noise:np.ndarray, use_noise_alpha:float):
    dtype = noise.dtype
    device = noise.device
    use_noise = torch.from_numpy(use_noise).to(device=DEFAULT_DEVICE)
    a = 1-use_noise_alpha
    b = use_noise_alpha
    norm = (a**2+b**2)**0.5
    noise = (noise*a+use_noise*b)/norm
    return noise.to(dtype).to(device)

def preprocess_image(img, resolution=None):
    size = normalize_resolution(*img.size, resolution=resolution)
    return img.resize(size, Image.Resampling.LANCZOS)

def image_as_tensor(image):
    image = np.array(image.convert("RGB")).astype(np.float32) / 255.0
    image = image[None].transpose(0, 3, 1, 2) # channel, y, x
    image = torch.from_numpy(image)
    return 2.0 * image - 1.0 # to -1 ~ 1

def preprocess_mask(img, mask):
    w, h=img.size
    return mask.resize((w//8, h//8), Image.Resampling.LANCZOS)
def mask_as_numpy(mask):
    mask = mask.convert("L")
    mask = np.array(mask).astype(np.float32) / 255.0
    mask = np.tile(mask, (4, 1, 1))
    mask = mask[None].transpose(0, 1, 2, 3)  # what does this step do?
    return mask
def mask_as_tensor(mask):
    mask = torch.from_numpy(mask_as_numpy(mask))
    return mask

def preprocess(image):
    w, h = image.size
    w, h = map(lambda x: x - x % 32, (w, h))  # resize to integer multiple of 32
    image = image.resize((w, h), resample=Image.Resampling.LANCZOS)
    image = np.array(image).astype(np.float32) / 255.0
    image = image[None].transpose(0, 3, 1, 2)
    image = torch.from_numpy(image)
    return 2.0 * image - 1.0
NOISE_PRED_BATCH = 1
class IDBucket:
    def __init__(self, max_length, weight, start_id, end_id):
        self.ids = np.zeros((1, max_length), np.int64) + end_id
        self.ids[0, 0] = start_id
        self.weight = weight
        self.used = 0
        self.max_length = max_length-2
    def reinit(self):
        self.used = 0
        end_id = self.ids[0, -1]
        start_id = self.ids[0, 0]
        self.ids = np.zeros((1, self.max_length+2), np.int64) + end_id
        self.ids[0, 0] = start_id
    def add_ids(self, weight, *ids):
        n = len(ids)
        if(n+self.used > self.max_length):
            return False
        if(weight !=self.weight):
            return False
        for idx, i in enumerate(ids):
            self.ids[0, self.used+1] = i
            self.used+=1
        return True
    def pad_wrap(self, target):
        orig_used = self.used
        if(not orig_used):
            return
        while(self.used < target):
            n = self.used%orig_used
            self.add_ids(self.weight, self.ids[0, n+1])
        return

class Reproducable:
    def __init__(self, func, result: Union[Image.Image, List[Image.Image]], *args, **kwargs):
        self.func=func
        self.result=result
        self.args=args
        self.kwargs=kwargs
    def reproduce(self, **kwargs):
        kwa = {}
        kwa.update(self.kwargs)
        kwa.update(kwargs)
        return self.func(*self.args, **kwa)
    @property
    def noise(self) -> np.ndarray:
        return self.kwargs["use_noise"]
class NameAssign:
    def __init__(self):
        self.used = set()
        self.next = dict()
    def __call__(self, name):
        idx = self.next.get(name, 0)
        def concat():
            nonlocal idx, name
            return "%s_%s"%(name, idx)
        while(concat() in self.used):
            idx += 1
        ret = concat()
        self.next[name] = idx+1
        self.used.add(ret)
        return ret
def combine(weights, inputs):
    ret = 0
    for idx, w in enumerate(weights):
        ret+=w*inputs[idx]
    return ret
def split_batch(inp, batch_size=0, torch_cat=False):
    ret = []
    le = len(inp) if isinstance(inp, list) else inp.shape[0]
    for i in range(0, le, batch_size):
        bs = min(batch_size, le-i)
        batch = inp[i:i+bs]
        if(torch_cat):
            batch = torch.cat(batch)
        ret.append(batch)
    return ret
TI_NEAREST_ORIG_WORD = False

class CustomPipeline:
    def __init__(self, sd_pipeline: StableDiffusionPipeline, ti_autoload_path=None, lora_load_path=None):
        self.sd_pipeline = sd_pipeline

        self.raw_tokenize = partial(sd_pipeline.tokenizer,
            padding=False,
            return_tensors="np"
        )

        self.tokenizer = sd_pipeline.tokenizer
        self.text_encoder = sd_pipeline.text_encoder
        self.unet = sd_pipeline.unet
        self.sched = sd_pipeline.scheduler
        self.device = sd_pipeline.device

        self.hijack_lock = Lock()
        self.ti_autoload_path = ti_autoload_path
        self.lora_load_path = lora_load_path
        self.ti_alias = {}
        self.ti_alias["remote_random"] = self.random_prompt
        self.ti_loaded = {}
        self.loras = {}
        self.lora_info = {}
        self.lora_file = {}
        self.ti_ph_assign = NameAssign()
        self.orig_embd_weights = self.text_encoder.text_model.embeddings.token_embedding.weight.numpy(force=True)
        self.auto_load_ti()

        self.auto_load_lora()
        self.inference_lock = Lock()

    def random_prompt(self):
        st, ed = self.tokenizer("")["input_ids"]
        hi = min(st, ed)
        ids = np.random.randint(0, hi, (1, 30))
        return self.tokenizer.batch_decode(ids)
    @torch.no_grad()
    def auto_load_lora(self):
        if(not self.lora_load_path):
            return
        pth = path.join(self.lora_load_path)
        files = list()
        files.extend(glob(path.join(pth, "*.pt")))
        files.extend(glob(path.join(pth, "*.safetensors")))
        for i in files:
            name = path.splitext(path.basename(i))[0]
            if(name in self.loras):
                continue
            self.lora_file[name] = i
            pt = load_tensor(i)
            self.loras[name] = pt
        files = list()
        files.extend(glob(path.join(pth, "*", "*.pt")))
        files.extend(glob(path.join(pth, "*", "*.safetensors")))
        for i in files:
            dn = path.dirname(i)
            name = path.basename(dn)
            if(name in self.loras):
                continue
            self.lora_file[name] = i
            pt = load_tensor(i)
            self.loras[name] = pt

            meta = path.join(dn, "info.json")
            if(path.exists(meta)):
                with open(meta) as f:
                    j = json.load(f)
                    self.lora_info[name] = j

        for k in list(self.loras):
            if (k in self.lora_file):
                f = self.lora_file[k]
                if(not path.exists(f)):
                    self.loras.pop(k)
                    if(k in self.lora_info):
                        self.lora_info.pop(k)

    def wrap_lora(self, *args):
        self.auto_load_lora()
        state_dicts = []
        scales = []
        for w, name in args:
            if(name in self.loras):
                state_dicts.append(self.loras[name])
                scales.append(w)
        if(state_dicts):
            return LoRA.WrapLoRAUNet(self.unet, state_dicts, scales), LoRA.WrapLoRATextEncoder(self.text_encoder, state_dicts, scales)
        else:
            return LoRA.Empty(), LoRA.Empty()
    @torch.no_grad()
    def auto_load_ti(self):
        if(not self.ti_autoload_path):
            return
        tokenizer = self.tokenizer
        text_embeddings = self.text_encoder.text_model.embeddings
        ann = None
        orig_embd_weights = self.orig_embd_weights
        def nearest_words(emb):
            nonlocal ann, orig_embd_weights
            n, m = emb.shape
            if(ann is None):
                from annoy import AnnoyIndex
                ann = AnnoyIndex(m, "euclidean")
                st, ed = tokenizer("")["input_ids"]

                for i in range(max(st, ed)+1):
                    vec = orig_embd_weights[i]
                    print("vec id", i, end="\r")
                    ann.add_item(i, vec)
                ann.build(5)
            ret = []
            for i in range(n):
                nn = ann.get_nns_by_vector(emb[i], 1)[0]
                ret.append(nn)
            return tokenizer.batch_decode(np.array(ret).reshape((1, -1)))
        to_cat = []
        mes = []
        ls = list(glob(path.join(self.ti_autoload_path, "*.pt")))
        ls.extend(glob(path.join(self.ti_autoload_path, "*.safetensors")))

        for i in ls:
            with locked(self.hijack_lock):
                if(i in self.ti_loaded):
                    continue
                if(i.endswith("pt")):
                    pt = torch.load(i)
                else:

                    pt = {}
                    with safetensors.safe_open(i, framework="pt", device="cpu") as f:
                        for k in f.keys():
                            pt[k] = f.get_tensor(k)
                    if("emb_params" in pt):
                        bn = path.splitext(path.basename(i))[0]
                        pt["string_to_param"] = {"bn": pt["emb_params"]}
                        pt["name"] = bn
                    else:
                        print(list(pt))
                key = next(iter(pt["string_to_param"])) # first element
                embeddings = pt["string_to_param"][key].cpu().numpy() # (size, 768)

                placeholder_length, dim = embeddings.shape

                # orig_embd_weights = text_embeddings.token_embedding.weight.cpu().numpy()
                alias_key = pt["name"]
                alias_value = []
                for j in range(placeholder_length):
                    placeholder = self.ti_ph_assign("token")
                    added = tokenizer.add_tokens(placeholder)
                    while(not added):
                        placeholder = self.ti_ph_assign("token")
                        added = tokenizer.add_tokens(placeholder)
                    assert added == 1, "more than 1 id is added for placeholder %s"%placeholder
                    alias_value.append(placeholder)
                alias_value = " ".join(alias_value)
                added_ids = tokenizer(alias_value)["input_ids"]
                st, ed, added_ids = added_ids[0], added_ids[-1], added_ids[1:-1]
                to_cat.append(embeddings)
                mes.append("%s"%(alias_key,))
                if(TI_NEAREST_ORIG_WORD):
                    print(alias_key, "near", nearest_words(embeddings))
                self.ti_alias[alias_key] = alias_value
                self.ti_loaded[i] = i
        if(mes):
            print(*mes, sep=", ")
        if(to_cat):
            embd_weights = np.concatenate([orig_embd_weights]+to_cat)
            self.orig_embd_weights = embd_weights
            embd_weights = torch.tensor(embd_weights, device=DEFAULT_DEVICE)
            del text_embeddings.token_embedding
            text_embeddings.token_embedding = nn.Embedding.from_pretrained(embd_weights)

    def get_txt2img_multiplier(self, prompt, cfg=7.5, steps=20, width=512, height=768, neg_prompt = "", noise_pred_batch_size=NOISE_PRED_BATCH, use_noise=None, use_noise_alpha=1):
        """
        Calculate estimated runtime multiplier.
        """
        weights, ids = self.get_ids(prompt, neg_prompt = neg_prompt)
        ret = len(weights)
        ret *= width*height*steps
        return ret

    def get_img2img_multiplier(self, prompt, orig_image, cfg=7.5, steps=20, alpha=0.7, beta=0.9, neg_prompt="", noise_pred_batch_size=NOISE_PRED_BATCH, eta=0, use_noise=None, use_noise_alpha=1):
        weights, ids = self.get_ids(prompt, neg_prompt = neg_prompt)
        ret = len(weights)
        img = preprocess_image(orig_image)
        width, height = img.size
        ret *= width*height*steps*alpha
        return ret

    def get_inpaint_multiplier(self, prompt, orig_image, mask_image, cfg=7.5, steps=20, neg_prompt="", noise_pred_batch_size=NOISE_PRED_BATCH, use_noise=None, use_noise_alpha=1, mode=1):
        weights, ids = self.get_ids(prompt, neg_prompt = neg_prompt)
        nweights = len(weights)
        w, h = preprocess_image(orig_image).size
        return w*h*nweights*steps

    @torch.no_grad()
    def inpaint(self, prompt, orig_image, mask_image, cfg=7.5, steps=20, neg_prompt="", noise_pred_batch_size=NOISE_PRED_BATCH, use_noise=None, use_noise_alpha=1, mode=1):
        with torch.autocast("cuda"), locked(self.inference_lock):
            text_encoder = self.text_encoder
            unet = self.unet
            sched = self.sched
            vae = self.sd_pipeline.vae

            orig_image = preprocess_image(orig_image)
            orig_tensor = image_as_tensor(orig_image)
            mask_image = preprocess_mask(orig_image, mask_image)
            mask_numpy = mask_as_numpy(mask_image)
            mask_tensor = mask_as_tensor(mask_image)

            # prepare weights
            weights, ids = self.get_ids(prompt, neg_prompt = neg_prompt)
            sqrt_weightn = len(weights)**0.5
            weights = normalize_weights(weights, std=cfg/sqrt_weightn)
            self.debug_ids(weights, ids)
            conds = []
            for i in ids:
                cond = text_encoder(torch.tensor(i).to(device=DEFAULT_DEVICE))[0]
                conds.append(cond)
            cond_batches = split_batch(conds, noise_pred_batch_size, torch_cat=True)
            nconds = len(conds)

            # prepare latents
            latents_dtype = cond.dtype
            orig_tensor = orig_tensor.to(device=self.device, dtype=latents_dtype)
            init_latent_dist = vae.encode(orig_tensor).latent_dist  # vae distribution, vae mean and stddev...
            init_latents = init_latent_dist.sample(generator=None) # vae sample by stddev
            init_latents = 0.18215 * init_latents

            # prepare timestep
            init_timestep = int(steps)
            sched.set_timesteps(steps)
            init_timestep = min(init_timestep, steps)
            timesteps = sched.timesteps[-init_timestep]
            timesteps = torch.tensor([timesteps], device=self.device)

            # prepare added noise
            noise = torch.randn(init_latents.shape, generator=None, device=self.device, dtype=latents_dtype)

            if(use_noise is not None):
                noise = combine_noise(noise, use_noise, use_noise_alpha)
            use_noise = noise.cpu().numpy()
            orig_latents = init_latents
            init_latents = sched.add_noise(init_latents, noise, timesteps)

            latents = init_latents

            timesteps = sched.timesteps[0:].to(self.device)
            if(True):
                it = list(enumerate(self.sd_pipeline.progress_bar(timesteps)))
            else:
                it = enumerate(timesteps)

            # denoising
            if(mode==0):
                mask_numpy = 1-(1-mask_numpy)**(1/steps)
                mask_tensor = torch.from_numpy(mask_numpy).to(device=DEFAULT_DEVICE)
            for i, t in it:
                noise_pred = self._predict_noise(latents, weights, cond_batches, nconds, noise_pred_batch_size, t)
                if(i==len(it)-1):
                    latents_proper = orig_latents
                else:
                    _, pt = it[i+1]
                    latents_proper = sched.add_noise(orig_latents, noise, torch.tensor([pt]))
                latents = sched.step(noise_pred, t, latents).prev_sample
                if(mode==0):
                    pass
                    # latents = latents*mask_tensor+latents_proper*(1-mask_tensor)
                else:
                    current = 1-i/steps
                    mask_current = (mask_numpy>=current).astype(np.float32) # * (mask_numpy**(1/steps))
                    mask_tensor=torch.from_numpy(mask_current).to(device=DEFAULT_DEVICE)
                latents = latents*mask_tensor+latents_proper*(1-mask_tensor)
            latents = 1 / 0.18215 * latents
            image = vae.decode(latents).sample
            del latents
            image = (image / 2 + 0.5).clamp(0, 1)
            image = image.cpu().permute(0, 2, 3, 1).numpy()
            image = self.sd_pipeline.numpy_to_pil(image)
            torch.cuda.empty_cache()
            return Reproducable(self.inpaint, image[0],
                prompt, orig_image, mask_image,
                steps=steps,
                cfg=cfg,
                neg_prompt=neg_prompt,
                use_noise=use_noise,
                use_noise_alpha=1
            )

    @torch.no_grad()
    def img2img(self, prompt, orig_image, cfg=7.5, steps=20, alpha=0.7, beta=1, neg_prompt="", noise_pred_batch_size=NOISE_PRED_BATCH, return_noise=False, use_noise=None, use_noise_alpha=1, print_progress=True, eta=0):
        """
        alpha: the rate of num_denoising_step/steps
        """
        loras = WeightedPrompt(prompt).loras
        loras = [(w, name) for name, w in loras.items()]
        lora_u, lora_t = self.wrap_lora(*loras)
        with torch.autocast("cuda"), locked(self.inference_lock), lora_u, lora_t:
            text_encoder = self.text_encoder
            unet = self.unet
            sched = self.sched
            vae = self.sd_pipeline.vae
            if alpha < 0 or alpha > 1:
                raise ValueError(f"The value of alpha should in [0.0, 1.0] but is {alpha}")

            # prepare orig image
            orig_image = preprocess_image(orig_image)
            orig_tensor = image_as_tensor(orig_image)
            # orig_tensor = preprocess(orig_image)

            # prepare weights
            weights, ids = self.get_ids(prompt, neg_prompt = neg_prompt)
            sqrt_weightn = len(weights)**0.5
            weights = normalize_weights(weights, std=cfg/sqrt_weightn)
            self.debug_ids(weights, ids)
            cond_num = len(weights)
            conds = []
            for i in ids:
                cond = text_encoder(torch.tensor(i).to(device=DEFAULT_DEVICE))[0]
                conds.append(cond)
            cond_batches = dict()
            for i in range(0, cond_num, noise_pred_batch_size):
                bs = min(noise_pred_batch_size, cond_num-i)
                cond = torch.cat(conds[i:i+bs])
                cond_batches[i] = cond

            # prepare latents
            latents_dtype = cond.dtype
            orig_tensor = orig_tensor.to(device=self.device, dtype=latents_dtype)
            init_latent_dist = vae.encode(orig_tensor).latent_dist  # vae distribution, vae mean and stddev...
            init_latents = init_latent_dist.sample(generator=None) # vae sample by stddev
            init_latents = 0.18215 * init_latents

            # prepare timestep
            sched.set_timesteps(steps)
            offset = sched.config.get("steps_offset", 0)
            init_timestep = int(steps * alpha) + offset
            init_timestep = min(init_timestep, steps)
            timesteps = sched.timesteps[-init_timestep]
            timesteps = torch.tensor([timesteps], device=self.device)
            if (True):
                init_numpy = init_latents.cpu().numpy()
                n = 0
                iter_count = 10
                for i in range(iter_count):
                    _noise = np.random.normal(size=init_numpy.shape)
                    if ((_noise*init_numpy).sum()<0):
                        n -= _noise
                    else:
                        n += _noise
                n /= sqrt(iter_count)
                print("std/mean", np.std(n), np.mean(n))
                noise = torch.from_numpy(n).to(device=self.device, dtype=latents_dtype)
            else:
                noise = torch.randn(init_latents.shape, generator=None, device=self.device, dtype=latents_dtype)

            if(use_noise is not None):
                use_noise = torch.tensor(use_noise).to(device=DEFAULT_DEVICE)
                noise = use_noise*use_noise_alpha+noise*(1-use_noise_alpha)
                sqrsum = use_noise_alpha**2+(1-use_noise_alpha)**2
                norm = sqrsum**0.5
                noise = noise/norm
            init_noise = noise.cpu().numpy()
            orig_latents = init_latents

            init_latents = sched.add_noise(init_latents, noise, timesteps)
            if(eta):
                noise = torch.randn(init_latents.shape, generator=None, device=self.device, dtype=latents_dtype)
                sqrsum = eta**2 + (1-eta)**2
                norm = sqrsum**0.5
                init_latents = init_latents*(1-eta) + noise*eta
                init_latents/=norm
            # denoising
            extra_step_kwargs={}
            latents = init_latents
            t_start = max(steps - init_timestep + offset, 0)
            timesteps = sched.timesteps[t_start:].to(self.device)
            rate_latent = beta
            rate_latent = beta**(1/steps/alpha)
            if(print_progress):
                it = list(enumerate(self.sd_pipeline.progress_bar(timesteps)))
            else:
                it = list(enumerate(timesteps))
            low_freq_fix = True
            hi_freq_fix = False
            for i, t in it:
                noise_pred = self._predict_noise(latents, weights, cond_batches, cond_num, noise_pred_batch_size, t)
                latents: torch.Tensor = sched.step(noise_pred, t, latents).prev_sample
                if(beta!=1):
                    if(i+1<len(it)):
                        i1, t1 = it[i+1]
                        latents_proper = sched.add_noise(orig_latents, noise, torch.tensor([t1]))
                    else:
                        latents_proper = orig_latents
                    if (low_freq_fix):
                        now_mean = latents.mean(axis=-1, keepdim=True).mean(axis=-2, keepdim=True)
                        tar_mean = latents_proper.mean(axis=-1, keepdim=True).mean(axis=-2, keepdim=True)
                        if (True):
                            latents_proper = latents_proper-tar_mean+now_mean
                        else:
                            latents = latents+(tar_mean-now_mean)*(1-rate_latent)
                    if (hi_freq_fix):
                        now_std = latents.std(axis=-1, keepdim=True).std(axis=-2, keepdim=True)
                        tar_std = latents_proper.std(axis=-1, keepdim=True).std(axis=-2, keepdim=True)
                        tar_mean = latents_proper.mean(axis=-1, keepdim=True).mean(axis=-2, keepdim=True)
                        if (True):
                            latents_proper = (latents_proper-tar_mean)/tar_std*now_std + tar_mean

                    latents = (latents*rate_latent)+(latents_proper*(1-rate_latent))
            latents = 1 / 0.18215 * latents
            image = vae.decode(latents).sample

            image = (image / 2 + 0.5).clamp(0, 1)
            image = image.cpu().permute(0, 2, 3, 1).numpy()
            image = self.sd_pipeline.numpy_to_pil(image)
            torch.cuda.empty_cache()
            return Reproducable(self.img2img, image[0], orig_image=orig_image, cfg=cfg, steps=steps, alpha=alpha, beta=beta, neg_prompt=neg_prompt, noise_pred_batch_size=noise_pred_batch_size, use_noise=init_noise, use_noise_alpha=1)

    def _predict_noise(self, latents, weights, cond_batches, cond_num, noise_pred_batch_size, t):
        unet = self.unet
        noises = []
        for idx in range(0, cond_num, noise_pred_batch_size):
            bs = min(noise_pred_batch_size, cond_num-idx)
            lat = torch.cat([latents]*bs)
            cond = cond_batches[idx]
            noise_pred = unet(lat, t, encoder_hidden_states=cond).sample
            for jdx in range(bs):
                noises.append((weights[idx+jdx], noise_pred[jdx]))
        noise_pred = None
        for w, noise in noises:
            if(noise_pred is None):
                noise_pred = noise*w
            else:
                noise_pred = noise_pred+noise*w
        return noise_pred
    def get_txt2img_interpolation_ids(self, prompt0, prompt1):
        SEP = ","
        wp0 = WeightedPrompt(prompt0).as_dict()
        wp1 = WeightedPrompt(prompt1).as_dict()
        keys = set(wp0.keys)+set(wp1.keys)
        for k in keys:
            w0 = k if k in wp0 else min(k, 0)
            w1 = k if k in wp1 else min(k, 0)
            s0 = wp0.get(k, "").split(SEP)
            s1 = wp1.get(k, "").split(SEP)
            lcs = LCS(s0, s1)

    def get_txt2img_interpolation_multiplier(self, prompt0, prompt1, cfg=7.5, steps=15, nframes=10, width=512, height=512, noise_pred_batch_size=NOISE_PRED_BATCH, use_noise=None, use_noise_alpha=1):
        weights0, ids0 = self.get_ids(prompt0)
        weights1, ids1 = self.get_ids(prompt1)
        nweights = len(weights0)+len(weights1)
        nfr = nframes if isinstance(nframes, int) else len(nframes)
        return nweights*width*height*steps*nfr
    @torch.no_grad()
    def txt2img_interpolation(self, prompt0, prompt1, cfg=7.5, steps=15, nframes=10, width=512, height=512, noise_pred_batch_size=NOISE_PRED_BATCH, return_noise = False, use_noise=None, use_noise_alpha=1):
        with torch.autocast("cuda"), locked(self.inference_lock):
            text_encoder = self.text_encoder
            unet = self.unet
            sched = self.sched
            vae = self.sd_pipeline.vae

            weights0, ids0 = self.get_ids(prompt0)
            weights1, ids1 = self.get_ids(prompt1)
            weights0 = normalize_weights(weights0, std=cfg/sqrt(len(weights0)))
            self.debug_ids(weights0, ids0)
            weights1 = normalize_weights(weights1, std=cfg/sqrt(len(weights1)))
            self.debug_ids(weights1, ids1)
            # cat_weights0 = np.concatenate([weights0, np.minimum(weights1, 0)], axis=0)
            cat_weights0 = np.concatenate([weights0, weights1*0], axis=0)
            # cat_weights1 = np.concatenate([np.minimum(weights0, 0), weights1], axis=0)
            cat_weights1 = np.concatenate([weights0*0, weights1], axis=0)
            len_embd = ids0[0].shape[-1]
            cat_ids = np.concatenate([ids0, ids1], axis=0).reshape((-1, len_embd))
            cat_ids = torch.tensor(cat_ids).to(device=DEFAULT_DEVICE)
            print("DEBUG:", cat_ids.shape)

            conds = text_encoder(cat_ids)[0]
            conds_batches = split_batch(conds, noise_pred_batch_size)

            latents_shape = (1, unet.in_channels, height // 8, width // 8)
            latents = torch.randn(latents_shape, generator=None, device=self.device, dtype=conds.dtype)
            if(use_noise is not None):
                latents = combine_noise(latents, use_noise, use_noise_alpha)
            init_noise = latents.cpu().numpy()

            if(True):
                print("DEBUG: using noise", init_noise.std(), init_noise.mean())
            # frame_latents = latents.repeat_interleave(nframes, dim=0)

            frame_weigths_batch = []
            if(isinstance(nframes, int)):
                for i in range(nframes):
                    rate = i/(nframes-1)
                    weight = cat_weights0*(1-rate)+cat_weights1*rate
                    weight = normalize_weights(weight, std=None)
                    frame_weigths = weight
                    frame_weigths_batch.append(split_batch(frame_weigths, noise_pred_batch_size))
            elif(isinstance(nframes, list)):
                for rate in nframes:
                    print("DEBUG: use rate", rate)
                    weight = cat_weights0*(1-rate)+cat_weights1*rate
                    weight = normalize_weights(weight, std=None)
                    print("DEBUG: weights", weight)
                    frame_weigths = weight
                    frame_weigths_batch.append(split_batch(frame_weigths, noise_pred_batch_size))
            frame_latents = [latents for i in frame_weigths_batch]
            print(frame_weigths_batch)
            sched.set_timesteps(steps)
            timesteps_tensor = sched.timesteps.to(self.device)
            for i, t in enumerate(self.sd_pipeline.progress_bar(timesteps_tensor)):
                for iframe, frame_latent in enumerate(frame_latents):
                    noise_preds = []
                    for bdx, cond in enumerate(conds_batches):
                        weights_batch = frame_weigths_batch[iframe][bdx]
                        lat = frame_latent.repeat_interleave(len(weights_batch), dim=0)
                        noise_pred_batch = unet(lat, t, encoder_hidden_states=cond).sample
                        for bdx, w in enumerate(weights_batch):
                            noise_pred = noise_pred_batch[bdx]
                            noise_preds.append((w, noise_pred))
                    ws = [i for i, j in noise_preds]
                    nps = [j for i, j in noise_preds]
                    noise_pred = combine(ws, nps)
                    latents = sched.step(noise_pred, t, frame_latent).prev_sample
                    frame_latents[iframe] = latents
            images = []
            for frame_latent in frame_latents:
                latent = 1/0.18215*frame_latent
                image = vae.decode(latent).sample
                image = (image / 2 + 0.5).clamp(0, 1)
                image = image.cpu().permute(0, 2, 3, 1).numpy()
                image = self.sd_pipeline.numpy_to_pil(image)
                images.append(image[0])

            torch.cuda.empty_cache()
            return Reproducable(self.txt2img_interpolation, images, prompt0, prompt1, cfg=cfg, steps=steps, nframes=nframes, width=width, height=height, noise_pred_batch=NOISE_PRED_BATCH, use_noise=init_noise, use_noise_alpha=1)

    def get_multi_diffusion_multiplier(self, masks, prompts, steps=25, cfg=7.5, width=512, height=768, neg_prompt = ""):
        n_noises = 0
        for idx, p in enumerate(prompts):
            weight, ids = self.get_ids(p, neg_prompt=neg_prompt)
            n_noises+=len(ids)*len(masks)
        return n_noises*width*height*steps

    @torch.no_grad()
    def multi_diffusion(self, masks, prompts, steps=25, cfg=7.5, width=512, height=768, neg_prompt = "", use_noise=None, use_noise_alpha=1, noise_pred_batch_size=NOISE_PRED_BATCH):
        with torch.autocast("cuda"), locked(self.inference_lock):
            tencoder = self.text_encoder
            unet = self.unet
            vae = self.sd_pipeline.vae
            sched = self.sched

            conds = []
            _masks = []
            for idx, mask in enumerate(masks):
                mask = mask_as_numpy(mask.resize((width//8, height//8), Image.LANCZOS))
                prompt = prompts[idx]
                weights, idss = self.get_ids(prompt, neg_prompt=neg_prompt)
                for jdx, w in enumerate(weights):
                    ids = idss[jdx]
                    cond = tencoder(torch.tensor(ids).to(device=DEFAULT_DEVICE))[0]
                    conds.append(cond)
                    _masks.append(mask*w)
            masks_arr = np.stack(_masks, axis=-1)
            center: np.ndarray = masks_arr-masks_arr.mean(axis=-1, keepdims=True)
            # norm = center/(center.std(axis=-1, keepdims=True)+1e-6)*cfg
            norm = center/(center.std()+1e-6)*cfg

            masks_arr = norm+1/norm.shape[-1]
            print(masks_arr.shape)
            # assert np.all(masks_arr.sum(axis=-1,keepdims=True)==1), masks_arr.sum(axis=-1)
            masks = [torch.from_numpy(masks_arr[:, :,:,:,idx]).to(self.device) for idx in range(masks_arr.shape[-1])]
            # for mask in masks:
            #     print("weight", mask[0, 0, 0, 0])

            latents_shape = (1, unet.in_channels, height // 8, width // 8)
            latents = torch.randn(latents_shape, generator=None, device=self.device, dtype=cond.dtype)

            sched.set_timesteps(steps)
            timesteps_tensor = sched.timesteps.to(self.device)

            for i, t in enumerate(self.sd_pipeline.progress_bar(timesteps_tensor)):
                noises = []
                for idx, cond in enumerate(conds):
                    noise_pred = unet(latents, t, encoder_hidden_states=cond).sample
                    noises.append(noise_pred)
                noise = torch.zeros_like(latents)
                for idx, mask in enumerate(masks):
                    # print(mask.shape)
                    noise += noises[idx]*mask

                latents = sched.step(noise, t, latents).prev_sample

            latents = 1 / 0.18215 * latents
            image = vae.decode(latents).sample

            image = (image / 2 + 0.5).clamp(0, 1)
            image = image.cpu().permute(0, 2, 3, 1).numpy()
            image = self.sd_pipeline.numpy_to_pil(image)

            torch.cuda.empty_cache()
            return Reproducable(self.multi_diffusion, image[0], masks, prompts, steps=steps, cfg=cfg, width=width, height=height, neg_prompt=neg_prompt)

    @torch.no_grad()
    def txt2img(self, prompt, cfg=7.5, steps=20, width=512, height=768, neg_prompt = "", noise_pred_batch_size=NOISE_PRED_BATCH, return_noise=False, use_noise=None, use_noise_alpha=1):
        loras = WeightedPrompt(prompt).loras
        loras = [(w, name) for name, w in loras.items()]
        lora_u, lora_t = self.wrap_lora(*loras)
        with torch.autocast("cuda"), locked(self.inference_lock), lora_u, lora_t:

            text_encoder = self.text_encoder
            unet = self.unet
            sched = self.sched
            vae = self.sd_pipeline.vae
            weights, ids = self.get_ids(prompt, neg_prompt = neg_prompt)
            sqrt_weightn = len(weights)**0.5
            weights = normalize_weights(weights, std=cfg/sqrt_weightn)
            self.debug_ids(weights, ids)
            cond_num = len(weights)
            conds = []
            for i in ids:
                print("te device", text_encoder.device)
                cond = text_encoder(torch.tensor(i).to(device=DEFAULT_DEVICE))[0]
                conds.append(cond)
            for idx, w in enumerate(weights):
                cond = conds[idx]
                token_ids = ids[idx]
                # tokens = self.tokenizer.batch_decode(token_ids)

            latents_shape = (1, unet.in_channels, height // 8, width // 8)
            latents = torch.randn(latents_shape, generator=None, device=self.device, dtype=cond.dtype)
            if(use_noise is not None):
                latents = combine_noise(latents, use_noise, use_noise_alpha)
            init_noise = latents.cpu().numpy()
            sched.set_timesteps(steps)
            timesteps_tensor = sched.timesteps.to(self.device)

            # scale the initial noise by the standard deviation required by the scheduler
            latents = latents * sched.init_noise_sigma
            cond_batches = dict()
            for i in range(0, cond_num, noise_pred_batch_size):
                bs = min(noise_pred_batch_size, cond_num-i)
                cond = torch.cat(conds[i:i+bs])
                cond_batches[i] = cond

            # accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
            extra_step_kwargs = {}
            # if accepts_eta:
            # extra_step_kwargs["eta"] = eta

            for i, t in enumerate(self.sd_pipeline.progress_bar(timesteps_tensor)):
                noise_pred = self._predict_noise(latents, weights, cond_batches, cond_num, noise_pred_batch_size, t)

                latents = sched.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample
            latents = 1 / 0.18215 * latents
            image = vae.decode(latents).sample

            image = (image / 2 + 0.5).clamp(0, 1)
            image = image.cpu().permute(0, 2, 3, 1).numpy()
            image = self.sd_pipeline.numpy_to_pil(image)

            torch.cuda.empty_cache()
            debug_vram("after txt2img")
            return Reproducable(self.txt2img, image[0], prompt, cfg=cfg, steps=steps, width=width, height=height, neg_prompt=neg_prompt, noise_pred_batch_size=noise_pred_batch_size, use_noise=init_noise, use_noise_alpha=1)

    def debug_ids(self, weights, ids):
        stid, edid = self.raw_tokenize("").input_ids[0]
        _ids = []
        for _ in ids:
            no_pad = []
            for id in _[0]:
                # print(_, stid, edid)
                if(id!=stid and id != edid):
                    no_pad.append(id)
            _ids.append([no_pad])
        ids = np.array(_ids)
        for idx, w in enumerate(weights):
            prompt = self.tokenizer.batch_decode(ids[idx])
            print("len=", len(ids[idx][0]), "weight =", w, "prompt =", prompt)

    def get_ids_v1(self, prompt, neg_prompt = ""):
        self.auto_load_ti()
        tokenizer = self.tokenizer
        max_length = self.tokenizer.model_max_length

        stid, edid = self.raw_tokenize("").input_ids[0]
        buckets = []
        def add_ids(weight, *ids):
            nonlocal buckets
            for b in buckets:
                if(b.add_ids(weight, *ids)):
                    return b
            bucket = IDBucket(max_length, weight, stid, edid)
            assert (bucket.add_ids(weight, *ids))
            buckets.append(bucket)
            return
        def add_sentence(weight, sentence):
            aliases = list(self.ti_alias.items())
            for k, v in sorted(aliases, key=lambda x:-len(x[0])):
                if(k in sentence):
                    print("%s -> %s"%(k, v))
                    clear_sentences = sentence.split(k)
                    for idx, i in enumerate(clear_sentences):
                        if(idx):
                            add_sentence(weight, v)
                        add_sentence(weight, i)
                    return
            if(", " in sentence):
                for s in sentence.split(", "):
                    add_sentence(weight, s)
                return
            ids = self.raw_tokenize(sentence).input_ids[0][1:-1] # 1:-1 ignore startoftext, endoftext
            if(len(ids)<=max_length-2):
                add_ids(weight, *ids)
            else:
                for id in ids:
                    add_ids(weight, id)
        for w, s in WeightedPrompt(prompt + "/*"+neg_prompt+"*/"):
            add_sentence(w, s)
        for b in buckets:
            b.pad_wrap(min((max_length-2)/3, b.used*2))
        # buckets = [b for b in buckets if b.used]
        # f_used = lambda u:(log(u)-log(max_length))/log(max_length)+2
        f_used = lambda u:u**0.5
        weights = np.array([b.weight*f_used(b.used) for b in buckets])
        # fewer token guidance may lead to worse quality and too strong guidance
        ids = [b.ids for b in buckets]
        return weights, ids
    def get_ids_v2(self, prompt, neg_prompt = ""):
        self.auto_load_ti()
        tokenizer = self.tokenizer
        max_length = self.tokenizer.model_max_length

        stid, edid = self.raw_tokenize("").input_ids[0]
        buckets = []
        self.auto_load_ti()
        tokenizer = self.tokenizer
        max_length = self.tokenizer.model_max_length

        stid, edid = self.raw_tokenize("").input_ids[0]
        buckets = []
        def add_ids(weight, *ids):
            nonlocal buckets
            for b in sorted(buckets, key=lambda x:x.used):

                if(b.add_ids(weight, *ids)):
                    return b
            bucket = IDBucket(max_length, weight, stid, edid)
            assert (bucket.add_ids(weight, *ids))
            buckets.append(bucket)
            return bucket
        def add_sentence(weight, sentence, spl = True):
            aliases = list(self.ti_alias.items())
            for k, v in sorted(aliases, key=lambda x:-len(x[0])):
                if(k in sentence):
                    print("%s -> %s"%(k, v))
                    clear_sentences = sentence.split(k)
                    for idx, i in enumerate(clear_sentences):
                        if(idx):
                            add_sentence(weight, v)
                        add_sentence(weight, i)
                    return

            if(spl and ", " in sentence):
                for idx, s in enumerate(sentence.split(", ")):
                    if(idx):
                        add_sentence(weight, ", "+s, spl=False)
                    else:
                        add_sentence(weight, s, spl=False)
                return

            ids = self.raw_tokenize(sentence).input_ids[0][1:-1] # 1:-1 ignore startoftext, endoftext
            if(len(ids)<=max_length-2):
                add_ids(weight, *ids)
            else:
                for id in ids:
                    add_ids(weight, id)
        if(isinstance(prompt, str)):
            wp = WeightedPrompt(prompt + "/*"+neg_prompt+"*/")
        else:
            wp = prompt
        for w, s in wp:
            add_sentence(w, s)

        nids = {}
        for b in buckets:
            nids[b.weight] = nids.get(b.weight, 0)+b.used
        buckets = []
        for w, nid in nids.items():
            n = max(ceil(nid/max_length*1.08), 1)
            for i in range(n):
                buckets.append(IDBucket(max_length, w, stid, edid))

        for w, s in wp:
            add_sentence(w, s)
        f_used = lambda u:u**0.1

        weights = np.array([b.weight*f_used(b.used) for b in buckets])
        # fewer token guidance may lead to worse quality and too strong guidance
        ids = [b.ids for b in buckets]
        return weights, ids
    def get_ids(self, prompt, neg_prompt = ""):
        self.auto_load_ti()
        tokenizer = self.tokenizer
        max_length = self.tokenizer.model_max_length

        stid, edid = self.raw_tokenize("").input_ids[0]
        buckets = []
        self.auto_load_ti()
        tokenizer = self.tokenizer
        max_length = self.tokenizer.model_max_length

        stid, edid = self.raw_tokenize("").input_ids[0]
        buckets = []
        def add_ids(weight, *ids):
            nonlocal buckets
            for b in sorted(buckets, key=lambda x:x.used):

                if(b.add_ids(weight, *ids)):
                    return b
            bucket = IDBucket(max_length, weight, stid, edid)
            assert (bucket.add_ids(weight, *ids))
            buckets.append(bucket)
            return bucket
        def add_sentence(weight, sentence, spl = True):
            aliases = list(self.ti_alias.items())
            for k, v in sorted(aliases, key=lambda x:-len(x[0])):
                if(callable(v)):
                    v = v()
                if(k in sentence):
                    print("%s -> %s"%(k, v))
                    clear_sentences = sentence.split(k)
                    for idx, i in enumerate(clear_sentences):
                        if(idx):
                            add_sentence(weight, v)
                        add_sentence(weight, i)
                    return
            sep = ","
            if(spl and sep in sentence):
                for idx, s in enumerate(sentence.split(sep)):
                    if(idx):
                        add_sentence(weight, sep+s, spl=False)
                    else:
                        add_sentence(weight, s, spl=False)
                return

            ids = self.raw_tokenize(sentence).input_ids[0][1:-1] # 1:-1 ignore startoftext, endoftext
            if(len(ids)<=max_length-2):
                add_ids(weight, *ids)
            else:
                print("TOO LONG", sentence)
                for id in ids:
                    add_ids(weight, id)
        if(isinstance(prompt, str)):
            wp = WeightedPrompt(prompt + "/*"+neg_prompt+"*/")
        else:
            wp = prompt
        extra_bucket = True
        while(extra_bucket):
            LIMIT = 30
            extra_bucket = False

            for idx, b in enumerate(buckets):
                b.reinit()
            for w, s in wp:
                n = len(buckets)
                add_sentence(w, s)
                m = len(buckets)
                extra_bucket = extra_bucket or (n!=m)
            if(len(buckets)>=LIMIT):
                f_used = lambda u:u**0.1
                weights = np.array([b.weight*f_used(b.used) for b in buckets])
                ids = [b.ids for b in buckets]
                self.debug_ids(weights, ids)
                assert len(buckets) < LIMIT
        f_used = lambda u:u**0.1

        weights = np.array([b.weight*f_used(b.used) for b in buckets])
        # fewer token guidance may lead to worse quality and too strong guidance
        ids = [b.ids for b in buckets]
        return weights, ids

if(__name__=="__main__"):
    from modules.diffusion.model_mgr import pipe as diffusion_pipe, models, get_model
    mask0 = Image.open(r"E:\Pics\gradient.png").convert("L")
    mask1 = Image.fromarray(255-np.array(mask0))
    prompt0 = "2girls, pink hair, cat ears"
    prompt1 = "2girls, blue hair, cat ears"
    nprompt = "worst quality"
    model = get_model()
    im = model.multi_diffusion([mask0, mask1], [prompt0, prompt1], width=512, height=512, neg_prompt=nprompt)
    im.show()


from bs4 import BeautifulSoup
import requests

URL = 'https://web.archive.org/web/20200518073855/https://www.empireonline.com/movies/features/best-movies-2/'

response = requests.get(URL)
movies_article = response.text

soup = BeautifulSoup(movies_article, 'html.parser')

title_tags = soup.find_all(name='h3', class_='title')

movies = [tag.getText() for tag in title_tags]
movies = movies[::-1] # Reverse the order of the list

with open('movies.txt', 'w', encoding='utf8') as f:
    f.write('\n'.join(movies))


from django.contrib import admin
from django.urls import path, re_path, include
from django.conf.urls.static import static
from django.conf import settings
from django.views import static

from blog.views import Index, Friends, Detail, Archive, CategoryList, CategoryView, TagList, TagView, About

urlpatterns = [
    path('admin/', admin.site.urls),

    # 首页
    path('', Index.as_view(), name='index'),

    # 友情链接
    path('friends/', Friends.as_view(), name='friends'),

    # 后台 markdown 编辑器配置
    path('mdeditor/', include('mdeditor.urls')),

    # 文章详情
    re_path(r'article/av(?P<pk>\d+)', Detail.as_view(), name='detail'),

    # 文章归档
    path('article/', Archive.as_view(), name='archive'),

    # 分类统计
    path(r'category/', CategoryList.as_view(), name='category'),

    # 文章分类
    re_path(r'category/cg(?P<pk>\d+)', CategoryView.as_view(), name='article_category'),

    # 标签统计
    path(r'tag/', TagList.as_view(), name='tag'),

    # 文章标签
    re_path(r'tag/tg(?P<pk>\d+)', TagView.as_view(), name='article_tag'),

    # 关于本站
    path('about/', About.as_view(),name='about'),

    # 静态文件
    re_path(r'^static/(?P<path>.*)$', static.serve, {'document_root': settings.STATIC_ROOT }, name='static')
]

# 设置后台名称
admin.site.site_header = '国光博客后台'
admin.site.site_title = 'Django Blog 后台'

if settings.DEBUG:
    urlpatterns += static(settings.MEDIA_URL, document_root=settings.MEDIA_ROOT)


import os
import cv2
import matplotlib.pyplot as plt
from pathlib import Path
import numpy as np

def detect(dataPath, clf):
    """
    Please read detectData.txt to understand the format. Load the image and get
    the face images. Transfer the face images to 19 x 19 and grayscale images.
    Use clf.classify() function to detect faces. Show face detection results.
    If the result is True, draw the green box on the image. Otherwise, draw
    the red box on the image.
      Parameters:
        dataPath: the path of detectData.txt
      Returns:
        No returns.
    """
    # Begin your code (Part 4)
    # raise NotImplementedError("To be implemented")
    with open(dataPath) as dtxt:
        image = np.zeros(shape=[512, 512], dtype=np.uint8)
        imgname = ""
        for line in dtxt.read().splitlines():
            lst = line.split(" ")
            if len(lst) == 2:
                imgname = lst[0]
                for root, dirs, files in os.walk(str(Path(dataPath).parent)):
                  if imgname in files:
                    image = cv2.imread(os.path.join(root, imgname))
                    # print("FOUND! Image Shape: ", image.shape)
            elif len(lst) == 4:
                x, y, w, h = int(lst[0]), int(lst[1]), int(lst[2]), int(lst[3])
                startpoint = (x, y)
                endpoint = (x+w, y+h)
                thickness = 3
                crop = image[y:y+h, x:x+w].copy()
                face = cv2.resize(crop, (19, 19))
                gray = cv2.cvtColor(face, cv2.COLOR_BGR2GRAY)
                if clf.classify(gray):
                   color = (0, 255, 0)
                  #  print("Is Face!")
                else:
                   color = (0, 0, 255)
                  #  print("Not face!")
                image = cv2.rectangle(image, startpoint, endpoint, color, thickness)
                cv2.imshow(imgname, image)
                cv2.waitKey(0)
        cv2.destroyAllWindows()
    dtxt.close()
    # End your code (Part 4)


#*****************************************************************************************
# Author: Luis Ernesto Colchado Soncco
# Email: luis.colchado@ucsp.edu.pe / luisernesto.200892@gmail.com
# Description: Generate Interpolation maps
#******************************************************************************************
import os, sys

from tqdm import tqdm

currentdir = os.path.dirname(os.path.realpath(__file__))
parentdir = os.path.dirname(currentdir)
sys.path.append(parentdir)

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
from datetime import timedelta
import shapefile as shp
from Models.Attention.ANN import ANN
from Models.Attention.attetion_layer import GraphAttentionLayer
from Models.commom import *
from Models.cGANSL.generator import Generator
import Models.IDW as IDW
import warnings
from tqdm import tqdm
warnings.filterwarnings('ignore')

try:
    DATASOURCE = sys.argv[1]
    if not DATASOURCE in ['BE', 'SP']:
        raise 'Invalid parameters: 1: BE or SP'
except Exception as e:
    raise 'Not found parameters: input DATASOURCE: SP or BE'

class setting:
    z_dim = 100
    if DATASOURCE == 'BE':
        dir_results = '../Experimentation/Results/BE'
        stations = ['Aotizhongxin', 'Dingling', 'Changping', 'Dongsi', 'Gucheng', 'Guanyuan',  'Huairou',
                   'Nongzhanguan', 'Shunyi', 'Wanliu', 'Wanshouxigong', 'Tiantan']
        knn = 3
        dir_cgansl = f"../Models/cGANSL/Backup/cgansl_advloss1_spl2_knn{knn}"
        sel_test_station = "Shunyi"
        selected_features = ['temp', 'pres', 'dewp', 'wd', 'ws', 'ndvi', 'ntl', 'dem']
        dir_cgansl_r2 = f"../Experimentation/Results/{DATASOURCE}/cgansl_advloss1_spl2_knn{knn}"
        dir_data = f"../Preprocessing/Results/data_train.csv"
        start_date = pd.to_datetime("2015-02-01")
        end_date = pd.to_datetime("2015-02-01")
        condition_features = ['temp', 'pres', 'dewp', 'wd', 'ws', 'ndvi', 'ntl', 'dem']
        dir_shp = "shp/beijing.shp"
        name_city = 'Beijing'
    else:
        dir_results = '../Experimentation/Results/SP'
        stations = ['Osasco', 'Pico do Jaraguá', 'Marg.Tietê-Pte Remédios',
                  'Cid.Universitária-USP-Ipen', 'Pinheiros', 'Parelheiros', 'Ibirapuera',
             'Congonhas', 'Santana', 'Parque D.Pedro II', 'Itaim Paulista']
        knn = 10
        dir_cgansl = f"../Models/cGANSL/Backup/cgansl_advloss0.2_spl0.8_knn{knn}"
        sel_test_station = "Itaim Paulista"
        selected_features = ['temp', 'press', 'hum', 'wd', 'ndvi', 'ntl', 'dem']
        dir_cgansl_r2 = f"../Experimentation/Results/{DATASOURCE}/cgansl_advloss0.2_spl0.8_knn{knn}"
        dir_data = f"../Preprocessing/Results/data_train_sp_.csv"
        start_date = pd.to_datetime("2018-07-01")
        end_date = pd.to_datetime("2018-07-01")#pd.to_datetime("2018-08-31")
        condition_features = ['temp', 'press', 'hum', 'wd', 'ndvi', 'ntl', 'dem']
        dir_shp = "shp/sao_paulo.shp"
        name_city = 'Sao Paulo'
    z_dim = 100
    dir_cgansl_r2 = f"{dir_cgansl_r2}/station_{sel_test_station}/r2_test.npy"
    dir_attention_r2 = f"../Experimentation/Results/{DATASOURCE}/attention_kernel_affine/station_{sel_test_station}/r2_test.npy"
    dir_attention = f"../Models/Attention/Backup/{DATASOURCE}_affine"
    for_normalization = selected_features+['PM25']
    input_g = len(selected_features)+(knn*2)+z_dim
    hidden_g = 100
    output_g = 1
    f_sig = torch.nn.Sigmoid()
    f_relu = torch.nn.LeakyReLU()
    idw_cols_train = ['lat', 'lon', 'PM25']
    idw_cols_test = ['lat', 'lon']
    p = 1

def load_map_shp(shp_path):
    sf = shp.Reader(shp_path, encoding='latin-1')
    return sf

def load_attention(epoch):
    number_features = len(setting.selected_features)
    dir_attention = f"{setting.dir_attention}/attention_layer_{epoch}.pt"
    attention_layer = GraphAttentionLayer(number_features, number_features, k=10,
                                    no_feature_transformation=True, graph=None,
                                    kernel="affine", nonlinearity_1=None, nonlinearity_2=nn.LeakyReLU())
    attention_layer.load_state_dict(torch.load(dir_attention))
    attention_layer.eval()

    dir_ann = f"{setting.dir_attention}/ann_{epoch}.pt"
    ann = ANN(number_features*2+1, 1)
    ann.load_state_dict(torch.load(dir_ann))
    ann.eval()
    return attention_layer, ann

def load_cgansl_g(epoch):
    dir_g = f"{setting.dir_cgansl}/g_model_{DATASOURCE}_{epoch}_{setting.sel_test_station}.pt"
    g = Generator(setting.input_g, setting.hidden_g, setting.output_g, f=setting.f_sig, f_relu=setting.f_relu)
    g.load_state_dict(torch.load(dir_g))
    g.eval()
    return g

def preparing_data(labeled_data, data_, name_model='attention'):
    if name_model == 'attention':
        data_ = data_[~data_['ntl'].isna()]
        scaled_data = data_.copy()
        scalers = {}
        for column in setting.for_normalization:
            scaled, scaler = normalization_interpolation(labeled_data, data_, column)
            scaled_data[column] = scaled.reshape(len(data_), 1)
            scalers.update({column: scaler})
        return scaled_data, scalers

def interpolation(x_tensor, name_model, model, ann=None):
    if name_model == 'attention':
        _, output_atten, _ = model(x_tensor)
        pm25_inter = ann(output_atten, DATASOURCE=DATASOURCE)
    elif name_model == 'cgansl':
        pm25_inter = model(x_tensor)
    return pm25_inter.detach().numpy()[-1]

if __name__ == '__main__':

    values_knn = []
    distance_features = []
    for k in range(setting.knn):
        values_knn.append(f'pm25_{k}')
        distance_features.append(f'dist_{k}')
    features_training = setting.condition_features+values_knn+distance_features

    data_frame = pd.read_csv(setting.dir_data)
    r2_attention = np.load(setting.dir_attention_r2)
    r2_cgansl = np.load(setting.dir_cgansl_r2)

    sf = load_map_shp(setting.dir_shp)

    max_atten_index = 500 if DATASOURCE == 'BE' else 328 #SP328#np.argmax(r2_attention)
    max_cgan_index = np.argmax(r2_cgansl)

    print(max_cgan_index)
    for j, v in enumerate(r2_cgansl):
        print(j, v)

    atten_layer, ann_model = load_attention(max_atten_index)
    cgansl_model = load_cgansl_g(max_cgan_index)

    labeled_data = data_frame[~data_frame['PM25'].isna()]
    lat_min, lat_max = min(data_frame['lat']), max(data_frame['lat'])
    lon_min, lon_max = min(data_frame['lon']), max(data_frame['lon'])

    scaled_data, scalers = preparing_data(labeled_data, data_frame, 'attention')
    scaled_data_cgansl, scalers_cgansl = preparing_data(data_frame, data_frame, 'attention')

    d_ = setting.start_date
    with tqdm(total=(setting.end_date - setting.start_date).days) as bar:
        while d_ <= setting.end_date:

            curr_data = scaled_data[scaled_data['current_date']==str(d_.date())]
            lab_data = curr_data[~curr_data['PM25'].isna()]
            unlab_data = curr_data[curr_data['PM25'].isna()]
            lab_data = lab_data[~lab_data['ntl'].isna()]
            x_lab = lab_data[setting.for_normalization].values

            curr_data_cgansl = scaled_data_cgansl[scaled_data_cgansl['current_date']==str(d_.date())]
            lab_data_cgansl = curr_data_cgansl[~curr_data_cgansl['PM25'].isna()]
            unlab_data_cgansl = curr_data_cgansl[curr_data_cgansl['PM25'].isna()]

            idw_train = lab_data[setting.idw_cols_train]
            idw_tree = IDW.tree(idw_train.values[:, :2], idw_train.values[:, -1])
            for i in unlab_data.index:
                x_unlab = unlab_data[setting.for_normalization].loc[i].values
                concat_data =  np.concatenate([x_lab, x_unlab.reshape(1, len(x_unlab))])
                x_tensor = to_tensor(concat_data, is_cuda=False)
                pm25 = interpolation(x_tensor=x_tensor, name_model='attention',
                                    model=atten_layer, ann=ann_model)
                pm25 = scalers['PM25'].inverse_transform(pm25.reshape(-1, 1))
                unlab_data.loc[i, "PM25_atten"] = pm25[-1]

                x_unlab_cgansl = unlab_data_cgansl[features_training].loc[i].values
                z = np.random.rand(1, setting.z_dim)
                x_unlab_cgansl = x_unlab_cgansl.reshape(1, len(x_unlab_cgansl))
                concat_cgansl = np.concatenate([z, x_unlab_cgansl], axis=1)
                x_tensor = to_tensor(concat_cgansl, is_cuda=False)
                pm25 = interpolation(x_tensor=x_tensor, name_model='cgansl',
                                    model=cgansl_model)
                pm25 = scalers_cgansl['PM25'].inverse_transform(pm25.reshape(-1, 1))
                unlab_data.loc[i, "PM25_cgansl"] = pm25[-1]

                x_test = unlab_data[setting.idw_cols_test].loc[i].values.reshape(1, 2)
                z = idw_tree(x_test, k=setting.knn, p=setting.p)
                pm25_idw = scalers['PM25'].inverse_transform(z.reshape(-1, 1))
                unlab_data.loc[i, "PM25_IDW"] = pm25_idw[-1]

            fig, axs = plt.subplots(figsize=(25, 15), nrows=2, ncols=3)
            fig.suptitle(f'{setting.name_city} ({str(d_.date())})', fontsize=22)
            x = []
            y = []
            for shape in sf.shapeRecords():
                for i in shape.shape.points[:]:
                    if i[0] <= lon_max and i[0] >= lon_min and i[1] <= lat_max and i[1] >= lat_min:
                        y.append(i[0])
                        x.append(i[1])
            axs[0, 0].plot(x,y, c='black', linewidth=0.02)
            axs[0, 1].plot(x, y, c='black', linewidth=0.02)
            axs[0, 2].plot(x, y, c='black', linewidth=0.02)
            axs[1, 0].plot(x, y, c='black', linewidth=0.02)
            axs[1, 1].plot(x, y, c='black', linewidth=0.02)

            lab_data['PM25_inv'] = scalers['PM25'].inverse_transform(lab_data['PM25'].values.reshape(-1, 1))

            im = axs[0, 0].scatter(x=lab_data['lat'], y=lab_data['lon'], c=lab_data['PM25_inv'], marker='s', s=150, cmap='YlOrBr')
            ax = axs[0, 0]
            ax.set_title('PM2.5 by AQM stations', fontsize=18)
            fig.colorbar(im, ax=ax)

            im = axs[0, 1].scatter(x=unlab_data['lat'], y=unlab_data['lon'], c=unlab_data['PM25_IDW'], marker='s', s=150, cmap='YlOrBr')
            ax = axs[0, 1]
            ax.set_title('PM2.5 by Inverse Distance Weighting', fontsize=18)
            fig.colorbar(im, ax=ax)

            im = axs[0, 2].scatter(x=unlab_data['lat'], y=unlab_data['lon'], c=unlab_data['PM25_atten'], marker='s', s=150, cmap='YlOrBr')
            ax = axs[0, 2]
            ax.set_title('PM2.5 by Neural Network with Attention Layer', fontsize=18)
            fig.colorbar(im, ax=ax)

            im = axs[1, 0].scatter(x=unlab_data['lat'], y=unlab_data['lon'], c=unlab_data['PM25_cgansl'], marker='s', s=150, cmap='YlOrBr')
            ax = axs[1, 0]
            ax.set_title('PM2.5 by Generative Model (cGAN & Spatial Learning)', fontsize=18)
            fig.colorbar(im, ax=ax)

            im = axs[1, 1].scatter(x=curr_data['lat'], y=curr_data['lon'], c=curr_data['ndvi'], marker='s', s=150, cmap='YlGn')
            ax = axs[1, 1]
            ax.set_title('NDVI (Vegetation Index)', fontsize=18)
            fig.colorbar(im, ax=ax)

            im = axs[1, 2].scatter(x=curr_data['lat'], y=curr_data['lon'], c=curr_data['ntl'], marker='s', s=150, cmap='gray')
            ax = axs[1, 2]
            ax.set_title('NTL (~Population Level)', fontsize=18)
            fig.colorbar(im, ax=ax)

            plt.savefig(f"maps/{DATASOURCE}/inter_{DATASOURCE}_{d_}", bbox_inches = 'tight')
            bar.update(1)
            d_ += timedelta(days=1)


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

fig = plt.figure()

ingBuffer = []
mark_list = ["4", "8", "p", "*", "x", "D", "4", "8", "p", "x"]
loop = len(mark_list)
num = 8

for i in range(loop):
    x = np.random.randn(num)
    y = np.random.randn(num)
    img = plt.plot(x, y, linewidth=0, marker=mark_list[i], markersize=15)
    ingBuffer.append(img)
    ts = 300
    ani = animation.ArtistAnimation(fig, ingBuffer, interval=ts, repeat=False)

# FuncAnimation sample

fig = plt.figure(figsize=(6, 2))
x = np.arange(0, 10, 0.01)

def update(k, fig_name, Amp):
    plt.cla()
    y = Amp+np.sin(x-k)
    plt.plot(x, y, 'r')
    plt.xlim(-1.2*Amp, 1.2*Amp)
    plt.title(fig_name + 'Frame k=' + str(k))

ani = animation.FuncAnimation(fig, update, fargs=('Animation', 1.8), interval=100, frames=32, repeat=False)


import dash_blueprint as dbp
import dash
from dash.dependencies import Input, Output
import dash_html_components as html
import dash_core_components as dcc

from typing import List
from pydantic.dataclasses import dataclass
import dataclasses

@dataclass
class User:
    id: int
    name: str = 'John Doe'
    # friends: List[int] = dataclasses.field(default_factory=lambda: [0])

users = [
    User(1, 'John'),
    User(2, 'Jow'),
    User(3, 'James'),
]

new_field_mapping = {
    'string': dbp.InputGroup,
    'integer': dbp.NumericInput,
    'float': dbp.NumericInput,
}

view_field_mapping = {
    'string': dbp.Text,
    'integer': dbp.Text,
    'float': dbp.Text
}

list_field_mapping = {
    'string': html.Td,
    'integer': html.Td,
    'float': html.Td
}

def new_object_form(model):
    """Create fields that allow for the form entry of a new object

    Args:
        model (dataclass): Pydantic model for the object

    Returns:
        List[object]: list of dash html objects
    """
    rows = []
    for name, props in model.__pydantic_model__.schema()['properties'].items():

        rows.append(
            dbp.FormGroup(
                inline=True,
                label=props['title'],
                children=new_field_mapping[props['type']]()
            )
        )
    return dbp.FormGroup(id='user-formgroup', children=rows)

def view_object_form(model, obj):
    """Create a read-only form representation of the object

    Args:
        model (dataclass): Pydantic model for the object
        obj (object): representation of the
    """
    rows = []
    for name, props in model.__pydantic_model__.schema()['properties'].items():
        rows.append(
            dbp.FormGroup(
                inline=True,
                label=props['title'],
                children=view_field_mapping[props['type']](children=getattr(obj, name))
            )
        )
    return dbp.FormGroup(id='user-formgroup-ro', children=rows)

def list_objects_table(model, objs):
    """Create a table of objects with values

    Args:
        model (dataclass): Pydantic model for the object
        objs (List[object]): iterable of objects to display
    """
    rows = []
    schema = model.__pydantic_model__.schema()
    display_props = list(schema['properties'].keys())
    for obj in objs:
        row = []
        for prop in display_props:
            row.append(list_field_mapping[schema['properties'][prop]['type']](children=getattr(obj, prop)))
        rows.append(dbp.Tr(row))
    return dbp.HTMLTable(
            id="table",
            sort_column=0,
            # filter_columns=[False, True, True],
            sort_columns=[True, True, True, False, True],
            # interactive=True,
            # show_more_less=True,
            children=[
                html.Thead(
                    dbp.Tr(children=[
                        html.Th(children=schema['properties'][prop]['title']) for prop in display_props
                    ])
                ),
                dbp.Tbody(children=rows),
            ],
        )

app = dash.Dash(__name__)

app.scripts.config.serve_locally = True
app.css.config.serve_locally = True

app.layout = html.Div(
    [

        dcc.Location(id="url", refresh=False),
        # new_object_form(User),
        # view_object_form(User, users[0]),
        list_objects_table(User, users),
        # dash_blueprint.Toaster(id="first-toaster"),
        # dash_blueprint.Toaster(
        #     id="second-toaster", toasterId="second-toaster", position="bottom-left"
        # ),
        # dash_blueprint.Toaster(
        #     id="third-toaster-same-as-first", position="bottom-right"
        # ),
    ]
)

if __name__ == "__main__":
    app.run_server(debug=True)


import pygame
from settings import *
from tile import Tile
from player import Player
from weapon import Weapon
from debug import debug_text
from support import *
from ui import UI
from particles import AnimationPlayer
from magic import MagicPlayer
from enemy import Enemy

class Level:
    def __init__(self):
        self.display_surface = pygame.display.get_surface()

        self.visible_sprites = YSortCameraGroup()
        self.obstacle_sprites = pygame.sprite.Group()

        self.current_attack = None
        self.attack_sprites = pygame.sprite.Group()
        self.attackable_sprites = pygame.sprite.Group()

        self.create_map()

        # interface
        self.ui = UI()

        # particle
        self.animation_player = AnimationPlayer()
        self.magic_player = MagicPlayer(self.animation_player)

    def create_map(self):
        layouts = {
            "boundary": import_csv_layout("level/level_data/map_boundaries.csv"),
            "desert_elements": import_csv_layout(
                "level/level_data/map_desert_elements.csv"
            ),
            "forest_elements": import_csv_layout(
                "level/level_data/map_forest_elements.csv"
            ),
            "swamp_elements": import_csv_layout(
                "level/level_data/map_swamp_elements.csv"
            ),
            "tundra_elements": import_csv_layout(
                "level/level_data/map_tundra_elements.csv"
            ),
            'entities': import_csv_layout("level/level_data/map_enemiesmodified.csv")
        }
        resource = create_graphics_dict()

        for style, layout in layouts.items():
            for row_index, row in enumerate(layout):
                for col_index, col in enumerate(row):
                    if col != "-1":
                        x = col_index * TILESIZE
                        y = row_index * TILESIZE
                        if style == "boundary":
                            Tile((x, y), [self.obstacle_sprites], "invisible")
                        if style == 'entities':
                            if col == '68':
                                monster_name = 'fire'
                            elif col == '38':
                                monster_name = 'ghost'
                            elif col == '41':
                                monster_name = 'bee'
                            else:
                                monster_name = None
                            if monster_name:
                                Enemy(monster_name,(x,y),[self.visible_sprites, self.attackable_sprites],self.obstacle_sprites, self.damage_player)

                        else:
                            if style in resource:
                                resource_index = int(col)
                                if 0 <= resource_index < len(resource[style]):
                                    surf = resource[style][resource_index]
                                    Tile(
                                        (x, y),
                                        [self.visible_sprites, self.obstacle_sprites],
                                        "object",
                                        surf,
                                    )
                    if col == "p":
                        self.player = Player(
                            (1374, 4187), [self.visible_sprites], self.obstacle_sprites, self.create_attack, self.destroy_attack,self.create_magic
                        )
        self.player = Player(
            (1374, 4187),
            [self.visible_sprites],
            self.obstacle_sprites,
            self.create_attack,
            self.destroy_attack,
            self.create_magic
        )

    def create_attack(self):
        self.current_attack = Weapon(self.player, [self.visible_sprites, self.attack_sprites])

    def player_attack_logic(self):
        if self.attack_sprites:
            for attack_sprite in self.attack_sprites:
                collision_sprites = pygame.sprite.spritecollide(attack_sprite, self.attackable_sprites, False)
                if collision_sprites:
                    for target_sprite in collision_sprites:
                        if target_sprite.sprite_type == 'grass':
                            pos = target_sprite.rect.center
                            self.animation_player.create_grass_particles(pos, [self.visible_sprites])
                            target_sprite.kill()
                        else:
                            target_sprite.get_damage(self.player, attack_sprite.sprite_type)

    def create_magic(self,style,strength,cost):
        if style == 'heal':
            self.magic_player.heal(self.player,strength,cost,[self.visible_sprites])

        if style == 'flame':
            self.magic_player.flame(self.player,cost,[self.visible_sprites])

    def destroy_attack(self):
        if self.current_attack:
            self.current_attack.kill()
        self.current_attack = None

    def player_attack_logic(self):
        if self.attack_sprites:
            for attack_sprite in self.attack_sprites:
                collision_sprites = pygame.sprite.spritecollide(attack_sprite, self.attackable_sprites, False)
                if collision_sprites:
                    for target_sprite in collision_sprites:
                            target_sprite.get_damage(self.player, attack_sprite.sprite_type)

    def damage_player(self, damage, attack_type):
        if self.player.vulnerable:
            if self.player.target_health >= damage:
                self.player.target_health -= damage
            else:
                self.player.target_health = 0
            self.player.vulnerable = False
            self.player.hit_time = pygame.time.get_ticks()

    def run(self):
        self.visible_sprites.enemy_update(self.player)
        self.visible_sprites.custom_draw(self.player)
        self.visible_sprites.update()

        self.player_attack_logic()
        self.ui.display(self.player)

class YSortCameraGroup(pygame.sprite.Group):
    def __init__(self):
        super().__init__()
        self.display_surface = pygame.display.get_surface()
        self.half_width = self.display_surface.get_size()[0] // 2
        self.half_height = self.display_surface.get_size()[1] // 2
        self.offset = pygame.math.Vector2()
        self.zoom_scale = 4

        # podłoże
        self.floor_surf = pygame.image.load("graphics/tilemap/map.png").convert()

        self.floor_rect = self.floor_surf.get_rect(topleft=(0, 0))
        self.scaled_surf = pygame.transform.scale(
            self.floor_surf,
            (
                self.floor_rect.width * self.zoom_scale,
                self.floor_rect.height * self.zoom_scale,
            ),
        )

    def custom_draw(self, player):
        self.offset.x = player.rect.centerx - self.half_width
        self.offset.y = player.rect.centery - self.half_height

        floot_offset_pos = self.floor_rect.topleft - self.offset
        self.display_surface.blit(self.scaled_surf, floot_offset_pos)

        for sprite in sorted(self.sprites(), key=lambda sprite: sprite.rect.centery):
            offset_pos = sprite.rect.topleft - self.offset
            self.display_surface.blit(sprite.image, offset_pos)

            if hasattr(sprite, 'monster_name'):
                if (sprite.health != sprite.max_health):
                    # Pasek zdrowia
                    health_bar_length = 50
                    health_bar_height = 7
                    health_ratio = sprite.health / sprite.max_health
                    filled_health = health_bar_length * health_ratio
                    health_bar_color = (255, 0, 0)  # Zielony kolor paska zdrowia
                    pygame.draw.rect(
                        self.display_surface,
                        health_bar_color,
                        (offset_pos[0], offset_pos[1] - 25, filled_health, health_bar_height),
                    )
                    pygame.draw.rect(
                        self.display_surface,
                        (0, 0, 0),
                        (offset_pos[0] + filled_health, offset_pos[1] - 25, health_bar_length - filled_health,
                        health_bar_height),
                    )

    def enemy_update(self, player):
        enemy_sprites = [sprite for sprite in self.sprites() if hasattr(sprite, 'sprite_type') and sprite.sprite_type == 'enemy']
        for enemy in enemy_sprites:
            enemy.enemy_update(player)