import sys
sys.dont_write_bytecode = True

import os
import shutil
import time
import stat
import json
import math
import hashlib
import pdb

import gradio as gr
import modules.extras
import modules.ui
import modules.scripts as scripts
import modules.prompt_parser as prompt_parser
import pprint
import piexif
import piexif.helper

from functools import partial
from PIL import Image
from copy import copy
from modules.shared import opts, cmd_opts
from modules.processing import process_images, Processed
from modules.images import FilenameGenerator
from modules import shared, images, scripts, script_callbacks
from pathlib import Path
from typing import List, Tuple

##########################################################################################

PROMPTBOOK_VERSION = "1.0"

refresh_symbol = '\U0001f504'  # 🔄

def sanitize_prompt(prompt):
    return prompt.strip().strip(',').strip()

class Recipe:
    def __init__(self, parts):
        self.parts = parts

    def serialize(self):
        return {
            "version": PROMPTBOOK_VERSION,
            "parts": [part.serialize() if part else None for part in self.parts]
        }

    def deserialize(data):
        parts = [RecipePart.deserialize(part_json) if part_json else None for part_json in data["parts"]]

        return Recipe(parts)

    def to_json(self):
        return json.dumps(self.serialize())

    def from_json(jsdata):
        data = json.loads(jsdata)
        assert isinstance(data, dict)
        return Recipe.deserialize(data)

    def from_json_array(jsdata):
        data = json.loads(jsdata)
        assert isinstance(data, list)
        return [Recipe.deserialize(entry) for entry in data]

    def evaluate(self, p=None, existing_prompt_action=None):
        prompt_parts = []
        negative_prompt_parts = []

        for part in self.parts:
            if part is None:
                continue

            prompt = part.prompt

            positive = apply_prompt_strength(prompt.positive, part.strength)
            negative = apply_prompt_strength(prompt.negative, part.strength)

            if positive.strip():
                prompt_parts.append(positive)
            if negative.strip():
                negative_prompt_parts.append(negative)

        if p:
            if p.prompt.strip():
                if existing_prompt_action == "Append":
                    prompt_parts.append(p.prompt)
                elif existing_prompt_action == "Prepend":
                    prompt_parts = [p.prompt] + prompt_parts

            if p.negative_prompt.strip():
                if existing_prompt_action == "Append":
                    negative_prompt_parts.append(p.negative_prompt)
                elif existing_prompt_action == "Prepend":
                    negative_prompt_parts = [p.negative_prompt] + negative_prompt_parts

        return join_prompts(prompt_parts), join_prompts(negative_prompt_parts)

class RecipePart:
    def __init__(self, prompt, strength):
        self.prompt: Prompt = prompt
        self.strength = strength

    def serialize(self):
        return {
            "prompt": self.prompt.serialize(slim=True),
            "strength": self.strength
        }

    def deserialize(data):
        prompt = Prompt.deserialize(data["prompt"])
        strength = data["strength"]

        return RecipePart(prompt, strength)

    def to_json(self):
        return json.dumps(self.serialize())

    def from_json(jsdata):
        data = json.loads(jsdata)
        return Recipe.deserialize(data)

class Prompt:
    def __init__(self, name, description, positive, negative, tags=[], sha256=None):
        self.name = name
        self.description = description
        self.positive = sanitize_prompt(positive)
        self.negative = sanitize_prompt(negative)
        self.tags = tags

        if name == "(Custom)":
            self.sha256 = ""
        elif sha256 is None:
            m = hashlib.sha256()
            m.update(self.positive.encode())
            m.update(self.negative.encode())
            self.sha256 = m.hexdigest()
        else:
            self.sha256 = sha256

    def load_from_file(png_filepath):
        image = Image.open(png_filepath)
        return Prompt.load_from_image(image)

    def load_from_image(image):
        _, info = images.read_info_from_image(image)
        return Prompt.from_json(info["promptbook_prompt"])

    def serialize(self, slim=False):
        data = {
            "name": self.name,
            "description": self.description,
            "positive": self.positive,
            "negative": self.negative,
            "sha256": self.sha256,
        }

        if not slim:
            data["version"] = PROMPTBOOK_VERSION
            data["tags"] = self.tags

        return data

    def deserialize(data):
        name = data["name"]
        description = data["description"]
        positive = data["positive"]
        negative = data["negative"]
        tags = data.get("tags", [])
        sha256 = data.get("sha256", None)

        return Prompt(name, description, positive, negative, tags, sha256)

    def to_json(self):
        return json.dumps(self.serialize())

    def from_json(jsdata):
        data = json.loads(jsdata)
        return Prompt.deserialize(data)

####################################################################################################

SPECIAL_PROMPTS = ["(None)", "(Custom)"]

def get_prompt_path(prompt_name):
    return os.path.join(opts.promptbook_prompts_path, f"{prompt_name}.png")

class UiPromptMerge:
    def __init__(self, rows):
        self.rows = rows

    def prompt_choices_outputs(self):
        return [row.dropdown_prompt for row in self.rows]

    def add_prompt_outputs(self):
        return list(sum([(row.dropdown_prompt, row.slider_strength) for row in self.rows], ()))

    def select_button_outputs(self):
        return [row.select_button for row in self.rows]

    def output_prompt_outputs(self):
        return list(sum([(row.dropdown_prompt, row.slider_strength, row.custom_prompt, row.custom_negative_prompt) for row in self.rows], ()))

    @staticmethod
    def load_recipe(recipe_json):
        recipe = Recipe.from_json(recipe_json)
        results = ["(None)", 1.0, "", ""] * int(opts.promptbook_merge_max_rows)

        for index, part in enumerate(recipe.parts):
            if part is None:
                continue
            elif part.prompt.name == "(Custom)":
                results[index * 4] = part.prompt.name
                results[index * 4 + 1] = part.strength
                results[index * 4 + 2] = part.prompt.positive
                results[index * 4 + 3] = part.prompt.negative
            else:
                prompt_path = get_prompt_path(part.prompt.name)
                if os.path.isfile(prompt_path):
                    loaded = Prompt.load_from_file(prompt_path)
                    if loaded.sha256 == part.prompt.sha256:
                        results[index * 4] = part.prompt.name
                        results[index * 4 + 1] = part.strength

        return results

    @staticmethod
    def output_recipe(*rest):
        recipe = UiPromptMerge.make_recipe(rest)
        return recipe.to_json()

    @staticmethod
    def clear_recipe():
        return ["(None)", 1.0, "", ""] * int(opts.promptbook_merge_max_rows)

    @staticmethod
    def evaluate_recipe(*rest):
        recipe = UiPromptMerge.make_recipe(rest)
        prompt, negative_prompt = recipe.evaluate()
        return f"{prompt}\nNegative prompt: {negative_prompt}".strip()

    @staticmethod
    def make_recipe(allvars):
        pairwise = list(zip(allvars[::4], allvars[1::4], allvars[2::4], allvars[3::4]))
        parts = []
        for pair in pairwise:
            prompt_name, strength, custom_prompt, custom_negative_prompt = pair
            if prompt_name != "" and prompt_name != "(None)":
                parts.append(UiPromptMergeRow.make_recipe_part(prompt_name, strength, custom_prompt, custom_negative_prompt))
            else:
                parts.append(None)

        while parts and parts[-1] is None:
            parts.pop()

        return Recipe(parts)

class UiPromptMergeRow:
    def __init__(self, index, row, delete_button, select_button, image_preview, dropdown_prompt, slider_strength, custom_prompt, custom_negative_prompt):
        self.index = index
        self.row = row
        self.delete_button = delete_button
        self.select_button = select_button
        self.image_preview = image_preview
        self.dropdown_prompt = dropdown_prompt
        self.slider_strength = slider_strength
        self.custom_prompt = custom_prompt
        self.custom_negative_prompt = custom_negative_prompt

    @staticmethod
    def make_recipe_part(prompt_name, strength, custom_prompt, custom_negative_prompt):
        if prompt_name == "(Custom)":
            prompt = Prompt("(Custom)", "", custom_prompt, custom_negative_prompt)
        else:
            prompt = Prompt.load_from_file(get_prompt_path(prompt_name))
        return RecipePart(prompt, strength)

def update_row(prompt_name):
    image_file = os.path.join(opts.promptbook_prompts_path, f"{prompt_name}.png")
    if prompt_name in SPECIAL_PROMPTS or not os.path.isfile(image_file):
        image_file = None

    is_custom = prompt_name == "(Custom)"
    return image_file, gr.Row.update(visible=is_custom)

def ui_prompt_merge():
    all_prompt_choices = SPECIAL_PROMPTS
    all_rows = []

    with gr.Row():
        with gr.Column():
            for i in range(0, int(opts.promptbook_merge_max_rows)):
                with gr.Row():
                    with gr.Column():
                        with gr.Row(variant="panel") as row:
                            with gr.Column(scale=1):
                                with gr.Row():
                                    delete_button = gr.Button("❌", elem_id=f"promptbook_prompt_delete_button_{i}")
                                    select_button = gr.Button(f"{i}", elem_id=f"promptbook_prompt_select_button_{i}", variant="primary" if i == 0 else "secondary")
                            with gr.Column(scale=19):
                                with gr.Row():
                                    image_preview = gr.Image(elem_id=f"promptbook_prompt_preview_{i}", type="pil", show_label=False, interactive=False).style(width=64, height=64)
                                    dropdown_prompt = gr.Dropdown(choices=all_prompt_choices, interactive=True, label="Prompt")
                                    slider_strength = gr.Slider(minimum=0.0, maximum=2.0, step=0.05, value=1.0, label="Strength", interactive=True)
                                with gr.Row(visible=False) as custom_row:
                                    custom_prompt = gr.components.Textbox(label="Custom prompt", lines=2)
                                    custom_negative_prompt = gr.components.Textbox(label="Custom negative prompt", lines=2)

                            all_rows.append(UiPromptMergeRow(i, row, delete_button, select_button, image_preview, dropdown_prompt, slider_strength, custom_prompt, custom_negative_prompt))

                dropdown_prompt.change(fn=update_row, inputs=[dropdown_prompt], outputs=[image_preview, custom_row])
                delete_button.click(fn=lambda: ("(None)", 1.0), inputs=None, outputs=[dropdown_prompt, slider_strength])
    with gr.Row():
        create_recipe_button = gr.Button("Generate Recipe", variant="primary")
        load_recipe_button = gr.Button("Load Recipe")
        clear_recipe_button = gr.Button("Clear Recipe")
        evaluate_recipe_button = gr.Button("Evaluate Recipe")
    with gr.Row():
        recipe_json = gr.Textbox(label="Recipe Output")

    prompt_merge_ui = UiPromptMerge(all_rows)
    create_recipe_button.click(fn=UiPromptMerge.output_recipe, inputs=prompt_merge_ui.output_prompt_outputs(), outputs=[recipe_json])
    load_recipe_button.click(fn=UiPromptMerge.load_recipe, inputs=[recipe_json], outputs=prompt_merge_ui.output_prompt_outputs())
    clear_recipe_button.click(fn=UiPromptMerge.clear_recipe, inputs=None, outputs=prompt_merge_ui.output_prompt_outputs())
    evaluate_recipe_button.click(fn=UiPromptMerge.evaluate_recipe, inputs=prompt_merge_ui.output_prompt_outputs(), outputs=[recipe_json])

    return prompt_merge_ui

####################################################################################################

tabs_list = ["promptbook", "generated"]
num_of_prompts_per_page = 0
PROMPTS_PATH = os.path.join(scripts.basedir(), "promptbook/prompts")
GENERATED_PATH = os.path.join(scripts.basedir(), "promptbook/generated")
ICON_PATH = os.path.join(scripts.basedir(), "static/promptbook_icon.png")

def traverse_all_files(curr_path, prompt_list) -> List[Tuple[str, os.stat_result]]:
    f_list = [(os.path.join(curr_path, entry.name), entry.stat()) for entry in os.scandir(curr_path)]
    for f_info in f_list:
        fname, fstat = f_info
        if os.path.splitext(fname)[1] == ".png":
            prompt_list.append(f_info)
        elif stat.S_ISDIR(fstat.st_mode):
            prompt_list = traverse_all_files(fname, prompt_list)
    return prompt_list

def get_all_prompts(sort_by, keyword):
    fileinfos = traverse_all_files(shared.opts.promptbook_prompts_path, [])
    keyword = keyword.strip(" ")
    if len(keyword) != 0:
        fileinfos = [x for x in fileinfos if keyword.lower() in x[0].lower()]
    if sort_by == "date":
        fileinfos = sorted(fileinfos, key=lambda x: -x[1].st_mtime)
    elif sort_by == "path name":
        fileinfos = sorted(fileinfos)

    filenames = [finfo[0] for finfo in fileinfos]
    return filenames

def get_prompt_page(page_index, filenames, keyword, sort_by):
    if page_index == 1 or page_index == 0 or len(filenames) == 0:
        filenames = get_all_prompts(sort_by, keyword)
    page_index = int(page_index)
    length = len(filenames)
    max_page_index = length // num_of_prompts_per_page + 1
    page_index = max_page_index if page_index == -1 else page_index
    page_index = 1 if page_index < 1 else page_index
    page_index = max_page_index if page_index > max_page_index else page_index
    idx_frm = (page_index - 1) * num_of_prompts_per_page
    prompt_list = filenames[idx_frm:idx_frm + num_of_prompts_per_page]

    visible_num = num_of_prompts_per_page if idx_frm + num_of_prompts_per_page < length else length % num_of_prompts_per_page
    visible_num = num_of_prompts_per_page if visible_num == 0 else visible_num

    load_info = "
"
    load_info += f"{length} images in this directory, divided into {int((length + 1) // num_of_prompts_per_page  + 1)} pages"
    load_info += "
"
    return filenames, page_index, prompt_list, visible_num, ""

def show_prompt_info(tabname_box, num, page_index, filenames):
    file = filenames[int(num) + int((page_index - 1) * num_of_prompts_per_page)]
    tm = "
" + time.strftime("%Y-%m-%d %H:%M:%S",time.localtime(os.path.getmtime(file))) + "
"
    return file, tm, num, file, ""

def populate_prompt_info(image):
    prompt = Prompt.load_from_image(image)
    return prompt.name, prompt.description, prompt.positive, prompt.negative, "", "", ','.join(prompt.tags)

def open_promptbook():
    tabname = "promptbook"
    with gr.Tab(tabname):
        with gr.Blocks(analytics_enabled=False):
            with gr.Row():
                first_page = gr.Button('First Page')
                prev_page = gr.Button('Prev Page')
                page_index = gr.Number(value=1, label="Page Index")
                next_page = gr.Button('Next Page')
                end_page = gr.Button('Last Page')
            with gr.Row(elem_id=tabname+'_promptbook'):
                with gr.Row():
                    with gr.Column(scale=3):
                        prompt_gallery = gr.Gallery(show_label=False, elem_id=tabname + "_promptbook_gallery").style(grid=6)
                        with gr.Row():
                            add_prompt_button = gr.Button("Add Prompt", elem_id=tabname + "_promptbook_add_prompt_btn")
                            add_prompt_index = gr.Number(label="Prompt Index", value=0)
                        prompt_merge_ui = ui_prompt_merge()
                    with gr.Column():
                        with gr.Row():
                            sort_by = gr.Radio(value="date", choices=["path name", "date"], label="sort by")
                            keyword = gr.Textbox(value="", label="keyword")
                        with gr.Row():
                            with gr.Column():
                                prompt_name = gr.Textbox(value="", label="Prompt Name", interactive=False)
                                prompt_description = gr.Textbox(value="", label="Prompt Description", interactive=False)
                                prompt_positive = gr.Textbox(label="Positive Prompt", interactive=False, lines=6)
                                prompt_negative = gr.Textbox(label="Negative Prompt", interactive=False, lines=6)
                                prompt_tags = gr.Textbox(value="", label="Prompt Tags", interactive=False)
                                prompt_file_name = gr.Textbox(value="", label="File Name", interactive=False)
                                prompt_file_time = gr.HTML()
                            with gr.Row():
                                try:
                                    send_to_buttons = modules.generation_parameters_copypaste.create_buttons(["txt2img", "img2img", "inpaint", "extras"])
                                except:
                                    pass

                    # hiden items
                    with gr.Row(visible=False):
                        renew_page = gr.Button("Renew Page", elem_id=tabname+"_promptbook_renew_page")
                        visible_img_num = gr.Number()
                        tabname_box = gr.Textbox(tabname)
                        prompt_index = gr.Textbox(value=-1)
                        set_index = gr.Button('set_index', elem_id=tabname+"_promptbook_set_index")
                        filenames = gr.State([])
                        prompt_info_switch = gr.Image(type="pil")
                        info1 = gr.Textbox()
                        info2 = gr.Textbox()
                        load_switch = gr.Textbox(value="load_switch", label="load_switch")
                        turn_page_switch = gr.Number(value=1, label="turn_page_switch")
                        prompt_file_info_raw = gr.Textbox(label="Generate Info", interactive=False, lines=6)

    #turn page
    first_page.click(lambda s:(1, -s) , inputs=[turn_page_switch], outputs=[page_index, turn_page_switch])
    next_page.click(lambda p, s: (p + 1, -s), inputs=[page_index, turn_page_switch], outputs=[page_index, turn_page_switch])
    prev_page.click(lambda p, s: (p - 1, -s), inputs=[page_index, turn_page_switch], outputs=[page_index, turn_page_switch])
    end_page.click(lambda s: (-1, -s), inputs=[turn_page_switch], outputs=[page_index, turn_page_switch])
    load_switch.change(lambda s:(1, -s), inputs=[turn_page_switch], outputs=[page_index, turn_page_switch])
    keyword.submit(lambda s:(1, -s), inputs=[turn_page_switch], outputs=[page_index, turn_page_switch])
    sort_by.change(lambda s:(1, -s), inputs=[turn_page_switch], outputs=[page_index, turn_page_switch])
    page_index.submit(lambda s: -s, inputs=[turn_page_switch], outputs=[turn_page_switch])
    renew_page.click(lambda s: -s, inputs=[turn_page_switch], outputs=[turn_page_switch])

    turn_page_switch.change(
        fn=get_prompt_page,
        inputs=[page_index, filenames, keyword, sort_by],
        outputs=[filenames, page_index, prompt_gallery, visible_img_num, prompt_file_info_raw]
    )
    turn_page_switch.change(fn=None, inputs=[tabname_box], outputs=None, _js="promptbook_turnpage")

    def reload_prompts():
        prompts = get_all_prompts("", "")
        choices = SPECIAL_PROMPTS + [os.path.splitext(os.path.basename(prompt))[0] for prompt in prompts]
        return [gr.update(choices=choices) for i in range(len(prompt_merge_ui.prompt_choices_outputs()))]
    turn_page_switch.change(fn=reload_prompts, inputs=[], outputs=prompt_merge_ui.prompt_choices_outputs())

    set_index.click(show_prompt_info,
                    _js="promptbook_get_current_img",
                    inputs=[
                        tabname_box,
                        prompt_index,
                        page_index,
                        filenames
                    ],
                    outputs=[
                        prompt_file_name,
                        prompt_file_time,
                        prompt_index,
                        prompt_info_switch,
                        prompt_file_info_raw
                    ])
    prompt_info_switch.change(fn=populate_prompt_info, inputs=[prompt_info_switch], outputs=[prompt_name, prompt_description, prompt_positive, prompt_negative, prompt_tags])
    prompt_info_switch.change(fn=modules.extras.run_pnginfo, inputs=[prompt_info_switch], outputs=[info1, prompt_file_info_raw, info2])

    try:
        modules.generation_parameters_copypaste.bind_buttons(send_to_buttons, prompt_file_name, prompt_file_info_raw)
    except:
        pass

    def add_prompt(add_prompt_name, add_index, *rest):
        pairwise = list(zip(rest[::2], rest[1::2]))
        result = []

        for index, pair in enumerate(pairwise):
            prompt_name, strength = pair
            if index == add_index:
                print("Add " + add_prompt_name)
                prompt_name = add_prompt_name
                strength = 1.0
            result += [prompt_name, strength]

        add_index += 1
        if add_index >= len(pairwise):
            add_index = 0

        return [add_index] + result

    add_prompt_vars = prompt_merge_ui.add_prompt_outputs()
    add_prompt_button.click(fn=add_prompt, inputs=[prompt_name, add_prompt_index] + add_prompt_vars, outputs=[add_prompt_index] + add_prompt_vars)

    def update_select_buttons(selected_index):
        return [gr.Button.update(variant="primary" if selected_index == i else "secondary") for i in range(0, int(opts.promptbook_merge_max_rows))]
    add_prompt_index.change(fn=update_select_buttons, inputs=[add_prompt_index], outputs=prompt_merge_ui.select_button_outputs())

    for row in prompt_merge_ui.rows:
        row.select_button.click(fn=lambda i=row.index: i, inputs=None, outputs=[add_prompt_index])
        row.delete_button.click(fn=lambda i=row.index: i, inputs=None, outputs=[add_prompt_index])

def get_all_generated(prompt_name, sort_by, keyword):
    path = os.path.join(opts.promptbook_generated_path, prompt_name + ".json")
    if not os.path.isfile(path):
        return []

    with open(path, "r") as file:
        jsdata = json.load(file)

    # TODO use correct sha256sum
    results = []
    for files in jsdata["files"].values():
        for entry in files:
            if os.path.isfile(entry["filename"]):
                results.append(entry["filename"])

    return results

def get_generated_page(prompt_name, page_index, filenames, keyword, sort_by):
    if page_index == 1 or page_index == 0 or len(filenames) == 0:
        filenames = get_all_generated(prompt_name, sort_by, keyword)
    page_index = int(page_index)
    length = len(filenames)
    max_page_index = length // num_of_prompts_per_page + 1
    page_index = max_page_index if page_index == -1 else page_index
    page_index = 1 if page_index < 1 else page_index
    page_index = max_page_index if page_index > max_page_index else page_index
    idx_frm = (page_index - 1) * num_of_prompts_per_page
    prompt_list = filenames[idx_frm:idx_frm + num_of_prompts_per_page]

    visible_num = num_of_prompts_per_page if idx_frm + num_of_prompts_per_page < length else length % num_of_prompts_per_page
    visible_num = num_of_prompts_per_page if visible_num == 0 else visible_num

    load_info = "
"
    load_info += f"{length} images in this directory, divided into {int((length + 1) // num_of_prompts_per_page  + 1)} pages"
    load_info += "
"
    return filenames, page_index, prompt_list, visible_num, ""

def show_generated_info(tabname_box, num, page_index, filenames):
    file = filenames[int(num) + int((page_index - 1) * num_of_prompts_per_page)]
    tm = "
" + time.strftime("%Y-%m-%d %H:%M:%S",time.localtime(os.path.getmtime(file))) + "
"
    return file, tm, num, file, ""

all_prompts = []

def reload_prompts():
    global all_prompts
    filenames = get_all_prompts("", "")
    all_prompts = [os.path.splitext(os.path.basename(filename))[0] for filename in filenames]

def open_generated():
    tabname = "generated"
    with gr.Tab(tabname):
        with gr.Blocks(analytics_enabled=False):
            with gr.Row():
                first_page = gr.Button('First Page')
                prev_page = gr.Button('Prev Page')
                page_index = gr.Number(value=1, label="Page Index")
                next_page = gr.Button('Next Page')
                end_page = gr.Button('Last Page')
            with gr.Row(elem_id=tabname+'_promptbook'):
                with gr.Row():
                    with gr.Column(scale=3):
                        with gr.Row():
                            prompt_name = gr.Dropdown(label='Prompt', elem_id="promptbook_prompt_name", choices=[])
                            modules.ui.create_refresh_button(prompt_name, reload_prompts, lambda: {"choices": sorted([x for x in all_prompts])}, "refresh_prompts")
                        seed_gallery = gr.Gallery(show_label=False, elem_id=tabname + "_promptbook_gallery").style(grid=6)

                    with gr.Column():
                        with gr.Row():
                            sort_by = gr.Radio(value="date", choices=["path name", "date"], label="sort by")
                            keyword = gr.Textbox(value="", label="keyword")
                        with gr.Row():
                            with gr.Column():
                                generated_file_info = gr.Textbox(label="Generate Info", interactive=False, lines=6)
                                generated_file_name = gr.Textbox(value="", label="File Name", interactive=False)
                                generated_file_time = gr.HTML()
                        with gr.Row(elem_id=tabname + "_promptbook_button_panel"):
                            try:
                                send_to_buttons = modules.generation_parameters_copypaste.create_buttons(["txt2img", "img2img", "inpaint", "extras"])
                            except:
                                pass

                    # hiden items
                    with gr.Row(visible=False):
                        renew_page = gr.Button("Renew Page", elem_id=tabname+"_promptbook_renew_page")
                        visible_img_num = gr.Number()
                        tabname_box = gr.Textbox(tabname)
                        generated_index = gr.Textbox(value=-1)
                        set_index = gr.Button('set_index', elem_id=tabname+"_promptbook_set_index")
                        filenames = gr.State([])
                        generated_info_switch = gr.Image(type="pil")
                        info1 = gr.Textbox()
                        info2 = gr.Textbox()
                        load_switch = gr.Textbox(value="load_switch", label="load_switch")
                        turn_page_switch = gr.Number(value=1, label="turn_page_switch")

    #turn page
    first_page.click(lambda s:(1, -s) , inputs=[turn_page_switch], outputs=[page_index, turn_page_switch])
    next_page.click(lambda p, s: (p + 1, -s), inputs=[page_index, turn_page_switch], outputs=[page_index, turn_page_switch])
    prev_page.click(lambda p, s: (p - 1, -s), inputs=[page_index, turn_page_switch], outputs=[page_index, turn_page_switch])
    end_page.click(lambda s: (-1, -s), inputs=[turn_page_switch], outputs=[page_index, turn_page_switch])
    load_switch.change(lambda s:(1, -s), inputs=[turn_page_switch], outputs=[page_index, turn_page_switch])
    keyword.submit(lambda s:(1, -s), inputs=[turn_page_switch], outputs=[page_index, turn_page_switch])
    sort_by.change(lambda s:(1, -s), inputs=[turn_page_switch], outputs=[page_index, turn_page_switch])
    page_index.submit(lambda s: -s, inputs=[turn_page_switch], outputs=[turn_page_switch])
    renew_page.click(lambda s: -s, inputs=[turn_page_switch], outputs=[turn_page_switch])
    prompt_name.change(lambda s: -s, inputs=[turn_page_switch], outputs=[turn_page_switch])

    turn_page_switch.change(
        fn=get_generated_page,
        inputs=[prompt_name, page_index, filenames, keyword, sort_by],
        outputs=[filenames, page_index, seed_gallery, visible_img_num, generated_file_info]
    )
    turn_page_switch.change(fn=None, inputs=[tabname_box], outputs=None, _js="promptbook_turnpage")

    set_index.click(show_generated_info,
                    _js="promptbook_get_current_img",
                    inputs=[
                        tabname_box,
                        generated_index,
                        page_index,
                        filenames
                    ],
                    outputs=[
                        generated_file_name,
                        generated_file_time,
                        generated_index,
                        generated_info_switch,
                        generated_file_info
                    ])
    generated_info_switch.change(fn=modules.extras.run_pnginfo, inputs=[generated_info_switch], outputs=[info1, generated_file_info, info2])

    try:
        modules.generation_parameters_copypaste.bind_buttons(send_to_buttons, generated_file_name, generated_file_info)
    except:
        pass

TAB_INDEX_GENERATE = 0
TAB_INDEX_SAVE_PROMPT = 1

def merge_processed(processed_list, include_lone_images=True, rows=None, make_grid=True) -> Processed:
    # Temporary list of all the images that are generated to be populated into the grid.
    # Will be filled with empty images for any individual step that fails to process properly
    image_cache = []

    processed_result = None
    for processed in processed_list:
        try:
            for processed_image in processed.images:
                if processed_result is None:
                    # Use our first valid processed result as a template container to hold our full results
                    processed_result = copy(processed)
                    cell_mode = processed_image.mode
                    cell_size = processed_image.size
                    processed_result.images = []
                    if make_grid:
                        processed_result.images.append(Image.new(cell_mode, cell_size))

                image_cache.append(processed_image)
                if include_lone_images:
                    processed_result.images.append(processed_image)
                    processed_result.all_prompts.append(processed.prompt)
                    processed_result.all_seeds.append(processed.seed)
                    processed_result.infotexts.append(processed.infotexts[0])
        except Exception:
            image_cache.append(Image.new(cell_mode, cell_size))

    if not processed_result:
        raise RuntimeError("Unexpected error: draw_xy_grid failed to return even a single processed image")

    if make_grid:
        if rows is None:
            rows = math.ceil(math.sqrt(len(image_cache)))

        # TODO: layout vertically depending on image size
        grid = images.image_grid(image_cache, rows=rows)

        processed_result.images[0] = grid

    return processed_result

####################################################################################################

def join_prompts(p):
    return sanitize_prompt(", ".join([sanitize_prompt(prompt) for prompt in p]))

def apply_prompt_strength(prompt, strength):
    if strength == 1.0 or not prompt.strip():
        return prompt

    parsed = prompt_parser.parse_prompt_attention(prompt)
    for pair in parsed:
        pair[1] *= strength

    return "".join(["(" + pair[0] + ":" + ("%0.3f" % pair[1]) + ")" for pair in parsed])

def apply_recipe(p, recipe, existing_prompt_action):
    prompt, negative_prompt = recipe.evaluate(p, existing_prompt_action)

    p.prompt = prompt
    p.negative_prompt = negative_prompt

    return process_images(p)

# HACK: Better way of getting result filename from a Processed (before script returns)?
def get_resulting_sample_filename(p, processed, seed, prompt, image, image_count, index):
    # Copied from modules/images.py
    namegen = FilenameGenerator(p, seed, prompt, image)
    basename = ""
    extension = opts.samples_format
    path = p.outpath_samples

    save_to_dirs = opts.save_to_dirs

    if save_to_dirs:
        dirname = namegen.apply(opts.directories_filename_pattern or "[prompt_words]").lstrip(' ').rstrip('\\ /')
        path = os.path.join(path, dirname)

    os.makedirs(path, exist_ok=True)

    if seed is None:
        file_decoration = ""
    elif opts.save_to_dirs:
        file_decoration = opts.samples_filename_pattern or "[seed]"
    else:
        file_decoration = opts.samples_filename_pattern or "[seed]-[prompt_spaces]"

    add_number = opts.save_images_add_number or file_decoration == ''

    if file_decoration != "" and add_number:
        file_decoration = "-" + file_decoration

    file_decoration = namegen.apply(file_decoration)

    if add_number:
        basecount = images.get_next_sequence_number(path, basename) - image_count
        fn = f"{basecount + index:05}"
        fullfn = os.path.join(path, f"{fn}{file_decoration}.{extension}")
    else:
        fullfn = os.path.join(path, f"{file_decoration}.{extension}")

    return fullfn

def append_generated(prompt, strength, output_img_filename):
    if prompt.name == "(Custom)" or prompt.sha256 == "":
        return

    path = os.path.join(opts.promptbook_generated_path, prompt.name + ".json")
    if os.path.isfile(path):
        with open(path, "r") as file:
            jsdata = json.load(file)
    else:
        jsdata = {
            "version": PROMPTBOOK_VERSION,
            "files": {}
        }

    if prompt.sha256 not in jsdata["files"]:
        jsdata["files"][prompt.sha256] = []

    jsdata["files"][prompt.sha256].append({"filename": output_img_filename, "strength": strength})

    with open(path, "w+") as file:
        json.dump(jsdata, file)

import base64
import json
import zlib
from PIL import Image, PngImagePlugin, ImageDraw, ImageFont
from fonts.ttf import Roboto

def caption_image_overlay(srcimage, title, textfont=None):
    from math import sin

    title = (title[:30] + '...') if len(title) > 30 else title

    image = srcimage.copy()
    fontsize = 32
    if textfont is None:
        try:
            textfont = ImageFont.truetype(None or Roboto, fontsize)
            textfont = None or Roboto
        except Exception:
            textfont = Roboto

    factor = 4.0
    gradient = Image.new('RGBA', (1, image.size[1]), color=(0, 0, 0, 0))
    for y in range(math.ceil(image.size[1]/3)):
        mag = sin(y/image.size[1]*factor)
        gradient.putpixel((0, y+math.floor(image.size[1]/3*2)), (0, 0, 0, int(mag*255)))
    image = Image.alpha_composite(image.convert('RGBA'), gradient.resize(image.size))

    draw = ImageDraw.Draw(image)

    font = ImageFont.truetype(textfont, fontsize)
    padding = 10

    _, _, w, h = draw.textbbox((0, 0), title, font=font)
    fontsize = min(int(fontsize * (((image.size[0]*0.95)-(padding*4))/w)), 72)
    font = ImageFont.truetype(textfont, fontsize)
    _, _, w, h = draw.textbbox((0, 0), title, font=font)
    draw.text((padding, image.size[1]-(padding*1.5)), title, anchor='ls', font=font, fill=(255, 255, 255, 230))

    icon_size = int(min(image.size[0], image.size[1]) * 0.2)
    icon = Image.open(ICON_PATH).resize((icon_size, icon_size))
    image.paste(icon, (padding, padding), icon)

    return image

class Script(scripts.Script):
    def title(self):
        return "Promptbook Generation"

    def show(self, is_img2img):
        return True

    def ui(self, is_img2img):
        with gr.Tabs():
            with gr.TabItem('Generate') as tab_generate:
                recipe_json = gr.Textbox(label="Recipe JSON")
                dropdown_txt2img_prompt = gr.Dropdown(choices=["Ignore", "Prepend", "Append"], value="Ignore", label="txt2img prompt")
                checkbox_save_grid = gr.Checkbox(label="Save grid", value=True)

            with gr.TabItem('Save Prompt') as tab_save:
                prompt_name = gr.Textbox(label="Prompt name (Leave blank to use positive prompt as name)", value="")
                prompt_description = gr.Textbox(label="Prompt description", value="")
                prompt_tags = gr.Textbox(label="Prompt tags", value="")
                save_prompt_card = gr.Checkbox(label="Save prompt card", value=True)
                with gr.Row():
                    overwrite_existing = gr.Checkbox(label="Overwrite existing", value=False)
                    generate_comparison = gr.Checkbox(label="Generate comparison", value=False)
                example_prompt = gr.Textbox(label="Example prompt", value="", lines=2)
                example_negative_prompt = gr.Textbox(label="Example negative prompt", value="", lines=2)
                append_example_prompts = gr.Checkbox(label="Append example prompts", value=False)
                batch_filename = gr.Textbox(label="Batch file (one positive prompt per line)", value="")
                batch_reset_seed = gr.Checkbox(label="Randomize seed each batch iteration", value=False)

        with gr.Row(visible=False):
            tab_index = gr.Number(value=0)

        tab_generate.select(lambda: (TAB_INDEX_GENERATE), inputs=None, outputs=[tab_index])
        tab_save.select(
            lambda p, pn: (
                TAB_INDEX_SAVE_PROMPT,
                shared.opts.promptbook_default_prompt if not p.strip() else p,
                shared.opts.promptbook_default_negative_prompt if not pn.strip() else pn),
            inputs=[example_prompt, example_negative_prompt],
            outputs=[tab_index, example_prompt, example_negative_prompt])

        return [
            tab_index,
            recipe_json,
            dropdown_txt2img_prompt,
            checkbox_save_grid,
            prompt_name,
            prompt_description,
            prompt_tags,
            overwrite_existing,
            generate_comparison,
            save_prompt_card,
            example_prompt,
            example_negative_prompt,
            append_example_prompts,
            batch_filename,
            batch_reset_seed
        ]

    def run(self,
            p,
            tab_index,
            gen_recipe_json,
            gen_txt2img_prompt,
            gen_save_grid,
            save_prompt_name,
            save_prompt_description,
            save_prompt_tags,
            save_overwrite_existing,
            save_generate_comparison,
            save_prompt_card,
            save_example_prompt,
            save_example_negative_prompt,
            save_append_example_prompts,
            save_batch_filename,
            save_batch_reset_seed):
        modules.processing.fix_seed(p)

        os.makedirs(opts.promptbook_generated_path, exist_ok=True)

        if tab_index == TAB_INDEX_GENERATE:
            return self.generate(p, gen_recipe_json, gen_txt2img_prompt, gen_save_grid)
        elif tab_index == TAB_INDEX_SAVE_PROMPT:
            return self.save_prompt(
                p,
                save_prompt_name,
                save_prompt_description,
                save_prompt_tags,
                save_overwrite_existing,
                save_generate_comparison,
                save_prompt_card,
                save_example_prompt,
                save_example_negative_prompt,
                save_append_example_prompts,
                save_batch_filename,
                save_batch_reset_seed
            )

    def generate(self, p, recipe_json, txt2img_prompt, save_grid):
        if recipe_json.strip().startswith("["):
            recipes = Recipe.from_json_array(recipe_json)
        else:
            recipes = [Recipe.from_json(recipe_json)]

        results = []

        for recipe in recipes:
            processed = apply_recipe(p, recipe, txt2img_prompt)
            results.append(processed)

            unwanted_grid_because_of_img_count = len(processed.images) < 2 and opts.grid_only_if_multiple
            has_grid = (opts.return_grid or opts.grid_save) and not p.do_not_save_grid and not unwanted_grid_because_of_img_count

            for i in range(0, len(processed.all_seeds)):
                filename = get_resulting_sample_filename(p,
                                                         processed,
                                                         processed.all_seeds[i],
                                                         processed.all_prompts[i],
                                                         # ignore grid image
                                                         processed.images[i + 1 if has_grid else 0],
                                                         len(processed.all_seeds),
                                                         i)
                for part in recipe.parts:
                    if part is None:
                        continue
                    append_generated(part.prompt, part.strength, filename)

            if save_grid:
                images.save_image(processed.images[0], p.outpath_grids, "promptbook", prompt=p.prompt, seed=processed.seed, grid=True, p=p)

        if not results:
            raise Exception("No recipes succesfully generated images.")

        merged = merge_processed(results, include_lone_images=True, make_grid=False)

        return merged

    def generate_cover_with_comparison(self, p, i, filename, prompt_description, pos, neg, example_prompt, example_negative_prompt, append_example_prompts):
        p1 = copy(p)
        p1.prompt = example_prompt
        p1.negative_prompt = example_negative_prompt
        before: Processed = process_images(p1)

        p2 = copy(p)
        if append_example_prompts:
            p2.prompt = join_prompts([pos, example_prompt])
            p2.negative_prompt = join_prompts([neg, example_negative_prompt])
        else:
            p2.prompt = join_prompts([example_prompt, pos])
            p2.negative_prompt = join_prompts([example_negative_prompt, neg])
        after: Processed = process_images(p2)

        original_prompt = after.infotexts[0]

        return merge_processed([before, after], include_lone_images=True, rows=1), original_prompt, p2

    def generate_cover(self, p, i, filename, prompt_description, pos, neg, example_prompt, example_negative_prompt, append_example_prompts):
        p1 = copy(p)
        if append_example_prompts:
            p1.prompt = join_prompts([pos, example_prompt])
            p1.negative_prompt = join_prompts([neg, example_negative_prompt])
        else:
            p1.prompt = join_prompts([example_prompt, pos])
            p1.negative_prompt = join_prompts([example_negative_prompt, neg])
        after: Processed = process_images(p1)

        return after, after.infotexts[0], p1

    def save_prompt(self,
                    p,
                    prompt_name,
                    prompt_description,
                    prompt_tags,
                    overwrite_existing,
                    generate_comparison,
                    save_prompt_card,
                    example_prompt,
                    example_negative_prompt,
                    append_example_prompts,
                    batch_filename,
                    batch_reset_seed):
        if prompt_name == "":
            prompt_name = p.prompt
        sanitized_name = images.sanitize_filename_part(prompt_name, replace_spaces=False)
        if batch_filename == "" and not sanitized_name.strip():
            raise RuntimeError("Prompt name was blank.")

        p.batch_count = 1
        # p.batch_size = 1

        if batch_filename != "":
            with open(batch_filename, "r", encoding="utf-8") as file:
                ext = os.path.splitext(batch_filename)
                if ext == ".json":
                    prompts_json = json.load(file)
                    prompts = [(p.get("prompt", ""), p.get("negative_prompt", ""), p.get("tags", [])) for p in prompts_json]
                else:
                    prompts = [(line.strip(), "", []) for line in file]
        else:
            tags = [tag.strip() for tag in prompt_tags.split(',')]
            prompts = [(p.prompt, p.negative_prompt, tags)]

        steps_per_run = 2 if generate_comparison else 1

        shared.total_tqdm.updateTotal(p.steps * p.n_iter * steps_per_run * len(prompts))
        shared.state.job_count = p.n_iter * steps_per_run * len(prompts)

        results = []

        for i, prompt in enumerate(prompts):
            if batch_reset_seed:
                p.seed = -1
                p.subseed = -1
                modules.processing.fix_seed(p)

            pos, neg, tags = prompt

            filename = sanitized_name
            if batch_filename != "":
                filename = images.sanitize_filename_part(pos, replace_spaces=False)

            outpath = os.path.join(opts.promptbook_prompts_path, f"{filename}.png")
            if os.path.exists(outpath) and not overwrite_existing:
                print("File already exists, skipping: " + filename)
                continue

            print(f"Generating prompt card '{filename}'")
            print(f"Positive:{pos}")
            if neg != "":
                print(f"Negative:{neg}")
            if tags:
                print(f"Tags: {', '.join(tags)}")

            if generate_comparison:
                processed, original_prompt, p2 = self.generate_cover_with_comparison(p, i, filename, prompt_description, pos, neg, example_prompt, example_negative_prompt, append_example_prompts)
            else:
                processed, original_prompt, p2 = self.generate_cover(p, i, filename, prompt_description, pos, neg, example_prompt, example_negative_prompt, append_example_prompts)

            if save_prompt_card:
                prompt = Prompt(filename, prompt_description, pos, neg, tags)
                cover = caption_image_overlay(processed.images[0], filename)
                grid_outpath = images.save_image(
                    cover,
                    opts.promptbook_prompts_path,
                    basename="",
                    prompt=p2.prompt,
                    seed=processed.seed,
                    grid=True,
                    p=p2,
                    forced_filename=filename,
                    # don't save the info .txt alongside like if info=<...> were passed
                    # these images are mostly for comparison instead of generating something pretty
                    existing_info={"promptbook_prompt": prompt.to_json(), "parameters": original_prompt},
                    pnginfo_section_name=""
                )

            results.append(processed)

        if not results:
            raise Exception("No promptcards generated! Do all the files in the batch exist already?")

        return merge_processed(results, include_lone_images=True, make_grid=False)

def create_promptbook(pb):
    with gr.Tabs(elem_id="promptbook_tab"):
        open_promptbook()
        open_generated()

def on_ui_tabs():
    global num_of_prompts_per_page
    num_of_prompts_per_page = int(opts.promptbook_page_columns * opts.promptbook_page_rows)
    with gr.Blocks(analytics_enabled=False) as pb:
        create_promptbook(pb)
        gr.Checkbox(opts.promptbook_preload, elem_id="promptbook_preload", visible=False)
        gr.Textbox(",".join(tabs_list), elem_id="promptbook_tabnames_list", visible=False)
    return (pb, "Promptbook", "promptbook"),

DEFAULT_PROMPT = """masterpiece, high quality, highres, 1girl"""
DEFAULT_NEGATIVE_PROMPT = ""
DEFAULT_MAX_ROWS = 15

def on_ui_settings():
    section = ('promptbook', "Promptbook")
    shared.opts.add_option("promptbook_preload", shared.OptionInfo(False, "Preload images at startup", section=section))
    shared.opts.add_option("promptbook_page_columns", shared.OptionInfo(6, "Number of columns on the page", section=section))
    shared.opts.add_option("promptbook_page_rows", shared.OptionInfo(6, "Number of rows on the page", section=section))
    shared.opts.add_option("promptbook_default_prompt", shared.OptionInfo(DEFAULT_PROMPT, "Default prompt to use when generating prompt covers", section=section))
    shared.opts.add_option("promptbook_default_negative_prompt", shared.OptionInfo(DEFAULT_NEGATIVE_PROMPT, "Default negative prompt to use when generating prompt covers", section=section))
    shared.opts.add_option("promptbook_prompts_path", shared.OptionInfo(PROMPTS_PATH, "Path containing prompt .png files", section=section))
    shared.opts.add_option("promptbook_generated_path", shared.OptionInfo(GENERATED_PATH, "Path containing files for generation stats", section=section))
    shared.opts.add_option("promptbook_merge_max_rows", shared.OptionInfo(DEFAULT_MAX_ROWS, "Maximum prompt count for the prompt merger", section=section))

script_callbacks.on_ui_settings(on_ui_settings)
script_callbacks.on_ui_tabs(on_ui_tabs)


"""

Given an integer array nums sorted in non-decreasing order, return an array of the squares of each number sorted in non-decreasing order.

Example 1:

Input: nums = [-4,-1,0,3,10]
Output: [0,1,9,16,100]
Explanation: After squaring, the array becomes [16,1,0,9,100].
After sorting, it becomes [0,1,9,16,100].
Example 2:

Input: nums = [-7,-3,2,3,11]
Output: [4,9,9,49,121]

Constraints:

1 <= nums.length <= 104
-104 <= nums[i] <= 104
nums is sorted in non-decreasing order.

"""

# TESTCASES

"""

[-4,-1,0,3,10]
[-7,-3,2,3,11]

"""

# Solution 1

class Solution:
    def sortedSquares(self, nums: List[int]) -> List[int]:
        return sorted([i**2 for i in nums]) # trivial approach

# Solution 2
""" Trivial and less efficient: Sorting at each line """

class Solution:
    def sortedSquares(self, nums: List[int]) -> List[int]:
        f = []
        for item in nums:
            f.append(item**2)
            f.sort()
        return f

"""

Follow up: Squaring each element and sorting the new array is very trivial, could you find an O(n) solution using a different approach?

"""

# Solution 2
"""
Using 2 pointers: one from left, one from right
Compare each pointer and inserting largest at 0th position
"""

class Solution:
    def sortedSquares(self, nums: List[int]) -> List[int]:
        l = 0
        r = len(nums) - 1
        f = []
        while l <= r:
            ln = nums[l]**2
            rn = nums[r]**2
            if ln >= rn:
                f.insert(0, (ln))
                l += 1
            else:
                f.insert(0, rn)
                r -= 1
        return f

# Solution 2
"""

Using bisect_left to find the correct insert position for each element

"""

class Solution:
    def sortedSquares(self, nums: List[int]) -> List[int]:
        f = []
        for item in nums:
            item = item**2
            f.insert(bisect_left(f, item), item)
        return f


__author__ = 'Danyang'

class Solution:
    def maxProfit(self, prices):

        if len(prices) <= 1:
            return 0

        forward = [0 for _ in xrange(len(prices))]
        lowest_buy_price = prices[0]
        for i in xrange(1, len(prices)):

            forward[i] = max(forward[i-1], prices[i]-lowest_buy_price)

            lowest_buy_price = min(prices[i], lowest_buy_price)

        backward = [0 for _ in xrange(len(prices))]
        highest_sell_price = prices[-1]
        for i in xrange(len(prices)-2, -1, -1):

            backward[i] = max(backward[i+1], highest_sell_price-prices[i])

            highest_sell_price = max(prices[i], highest_sell_price)

        max_profit = 0
        for i in xrange(len(prices)):
            max_profit = max(max_profit, forward[i]+backward[i])
        return max_profit

    def maxProfit_error(self, prices):

        if len(prices) <= 1:
            return 0

        delta_prices = []
        for i in xrange(1, len(prices)):
            delta_prices.append(prices[i]-prices[i-1])

        max_profits = [0, 0]

        max_sub_array = 0
        current_sub_array = 0
        for j in xrange(len(delta_prices)):
            if current_sub_array+delta_prices[j] >= 0:
                current_sub_array += delta_prices[j]
                max_sub_array = max(max_sub_array, current_sub_array)
            else:

                if max_sub_array > max_profits[0]:
                    max_profits[1] = max_profits[0]
                    max_profits[0] = max_sub_array
                elif max_sub_array > max_profits[1]:
                    max_profits[1] = max_sub_array
                max_sub_array = 0
                current_sub_array = 0

        return sum(max_profits)

data_files/123 Best Time to Buy and Sell Stock III.py


import os
from dotenv import load_dotenv
import requests
import io
from PIL import Image
import streamlit as st
from streamlit_drawable_canvas import st_canvas
from utils import footer
import json
import base64

# Load environment variables
load_dotenv()

# Set up page config
st.set_page_config(
    page_title="Image Captioning",
    page_icon="ðŸŽ¨",
    layout="wide",
    initial_sidebar_state="expanded",
)

# Set up the title and a brief description
st.title("Image Captioning")
st.markdown("Capture the essence of your image with a caption!")

st.sidebar.markdown("""
## ðŸŽ‰ Welcome to Image Captioning! ðŸŽ‰

You know how they say a picture is worth a thousand words? Well, we're about to prove them wrong.

Here, we believe that a picture is worth exactly one caption. No more, no less.

So go ahead, upload an image or draw one if you're feeling artistic. Our state-of-the-art AI will then generate a caption that captures the essence of your image in a single sentence.

It's like magic, but without the rabbits and hats. Enjoy!
""")

# Get API token and URL from environment variables
API_TOKEN = os.getenv('API_TOKEN')
API_URL = os.getenv('IMAGE_CAPTIONING_API_URL')
headers = {"Authorization": f"Bearer {API_TOKEN}"}

# Function to send a POST request to the inference server
def query(payload):
    response = requests.post(API_URL, headers=headers, json=payload)
    return response.content

# Function to get caption from inference server
def get_caption_from_inference_server(image):
    image = image.convert("RGB")  # Convert image to RGB
    buffered = io.BytesIO()
    image.save(buffered, format="JPEG")
    img_str = base64.b64encode(buffered.getvalue())
    response = query({"inputs": img_str.decode('utf-8'), })
    response_json = json.loads(response.decode('utf-8'))  # Decode byte string and parse JSON
    caption = response_json[0]['generated_text']  # Extract generated text
    return caption

# Select image source
image_source = st.radio("Select image source", ["Upload image", "Draw image"])

if image_source == "Upload image":
    # Upload image
    uploaded_file = st.file_uploader("Upload an image", type=["png", "jpg", "jpeg"])
    if uploaded_file is not None:
        image = Image.open(uploaded_file)
        st.image(image, caption="Uploaded Image", width=400)
else:
    # Draw image
    st.markdown("Draw an image:")
    stroke_width = st.slider("Stroke width:", 1, 25, 10)
    stroke_color = st.color_picker("Stroke color:", '#e00')
    background_color = st.color_picker("Background color:", '#fff')

    canvas_result = st_canvas(
        fill_color="rgba(255, 165, 0, 0.3)",  # Fixed fill color with some opacity
        stroke_width=stroke_width,
        stroke_color=stroke_color,
        background_color=background_color,
        height=500,  # Adjust the height of the canvas
        width=1000,  # Adjust the width of the canvas
        drawing_mode="freedraw",
        key="canvas",
    )

    # Convert the drawn image to PIL Image
    if canvas_result.image_data is not None:
        drawn_image = Image.fromarray(canvas_result.image_data.astype('uint8'), 'RGBA')
        st.image(drawn_image, caption="Drawn Image", use_column_width=False)

# Button to generate caption
if st.button('Generate Caption'):
    with st.spinner('Generating your caption...'):
        if image_source == "Upload image" and uploaded_file is not None:
            caption = get_caption_from_inference_server(image)
        elif image_source == "Draw image" and canvas_result.image_data is not None:
            caption = get_caption_from_inference_server(drawn_image)
        else:
            st.warning("Please upload or draw an image first.")
            st.stop()
    st.success('Caption generated successfully!')
    st.markdown(f"
{caption}
", unsafe_allow_html=True)

# Display footer
st.markdown(footer,unsafe_allow_html=True)


from setuptools import setup
import os
from glob import glob

package_name = 'cr3_scout_with_srvs'

setup(
    name=package_name,
    version='0.0.0',
    packages=[package_name],
    data_files=[
        #(os.path.join('share',package_name,'launch'), glob('launch/*_launch.py')),
        (os.path.join('share',package_name,'launch'), glob('launch/*_launch.py')),
        (os.path.join('share',package_name,'config'), glob('config/*.yaml')),
        ('share/ament_index/resource_index/packages',
            ['resource/' + package_name]),
        ('share/' + package_name, ['package.xml']),
    ],
    install_requires=['setuptools'],
    zip_safe=True,
    maintainer='qiuyuezuo',
    maintainer_email='qiuyuezuo@todo.todo',
    description='TODO: Package description',
    license='TODO: License declaration',
    tests_require=['pytest'],
    entry_points={
        'console_scripts': [
            'scout = cr3_scout_with_srvs.scout:main',
            'client = cr3_scout_with_srvs.cr3_service_clients:main',
            'scout_with_zoom_camera = cr3_scout_with_srvs.scout_with_zoom_camera:main',
            'camera_intrin_info = cr3_scout_with_srvs.camera_intrin_info:main',
            'scout_with_zoom_camera_horiz = cr3_scout_with_srvs.scout_with_zoom_camera_horiz:main',
            #'cam_info_fake_pub = cr3_scout_with_srvs.cam_info_fake_pub:main'
        ],
    },
)


# https://atcoder.jp/contests/joi2015yo/tasks/joi2015yo_d

if __name__ == "__main__":
    N, M = list(map(int, input().split()))
    D = []
    for _ in range(N):
        D.append(int(input()))
    C = []
    for _ in range(M):
        C.append(int(input()))

    dp = [[0] + [float('inf') for _ in range(N)] for _ in range(M+1)]

    # 各日
    for j in range(M):
        # 各都市
        for i in range(N):
            if dp[j][i] == float('inf'):
                continue
            dp[j+1][i+1] = min(dp[j][i+1], dp[j][i] + D[i] * C[j])

    print(dp[-1][-1])


import os
import unittest
from Logger import Logger

class TestLogger(unittest.TestCase):
    def test_file_handling(self):
        testLog = Logger("testLog")
        ## Check if program can create and open file
        self.assertTrue(testLog.opened)
        returns = testLog.close()
        ## Check if logger correctly signs bool OPENED and returns
        ##     0 as succes.
        self.assertFalse(testLog.opened)
        self.assertEqual(returns,0)
        returns = testLog.close()
        ## Check if logger returns 1 when trying to close already
        ##     closed file
        self.assertEqual(returns,1)
        ## Do cleanup:
        os.remove(testLog.name)

    def test_logging(self):
        testLog = Logger("testLog")
        testPhrase = "TestLine\r\n"
        testLog.save_line(testPhrase)
        testLog.close()
        logfile = open(testLog.name)
        content = logfile.read()
        logfile.close()
        saved = content.split(" : ")
        ## Check if saved data corresponds
        self.assertEqual(saved[1],testPhrase)
        ## cleanup
        os.remove(testLog.name)

from gpsNavigation import gpsModule,gpsPoint
class TestGPSNavigation(unittest.TestCase):
    def test_gps_angles(self):
        gpsMod = gpsModule()

        A = gpsPoint(10,10)
        B = gpsPoint(10.1,10.1)

        distance, azimut = gpsMod.GPSData.getDirAndDist(A,B)
        self.assertEqual(distance,15623.0)
        self.assertEqual(azimut,45.0)

        B = gpsPoint(10.0,10.1)
        distance, azimut = gpsMod.GPSData.getDirAndDist(A,B)
        self.assertEqual(distance,10963.0)
        self.assertEqual(azimut,90.0)

        B = gpsPoint(9.9,10.1)
        distance, azimut = gpsMod.GPSData.getDirAndDist(A,B)
        self.assertEqual(distance,15625.0)
        self.assertEqual(azimut,135.0)

        B = gpsPoint(9.9,10.0)
        distance, azimut = gpsMod.GPSData.getDirAndDist(A,B)
        self.assertEqual(distance,11132.0)
        self.assertEqual(azimut,180.0)

        B = gpsPoint(9.9,9.9)
        distance, azimut = gpsMod.GPSData.getDirAndDist(A,B)
        self.assertEqual(distance,15625.0)
        self.assertEqual(azimut,225.0)

        B = gpsPoint(10.0,9.9)
        distance, azimut = gpsMod.GPSData.getDirAndDist(A,B)
        self.assertEqual(distance,10963.0)
        self.assertEqual(azimut,270.0)

        B = gpsPoint(10.1,9.9)
        distance, azimut = gpsMod.GPSData.getDirAndDist(A,B)
        self.assertEqual(distance,15623.0)
        self.assertEqual(azimut,315.0)

        B = gpsPoint(10.1,10.0)
        distance, azimut = gpsMod.GPSData.getDirAndDist(A,B)
        self.assertEqual(distance,11132.0)
        self.assertEqual(azimut,0)

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


import os
import sys
import numpy as np
import matplotlib.pyplot as plt
from osgeo import gdal

## Mark Koszykowski
## ECE472 - Remote Sensing
## Homework 1 - Problem 1 Task 2

# Defining relative PATHs
GeoTIF_dir = "s2_santafe_spatially_aligned"
composites_dir = "composites"

# Function that loads Spatially Aligned data into a list of numpy arrays for manipulation
def loadImages(path):
    images = []
    labels = []
    geoTransform = None
    projection = None
    for GeoTIF in os.listdir(path):
        labels.append(path + '\\' + GeoTIF)
        dataset = gdal.Open(path + '\\' + GeoTIF)
        geoTransform = dataset.GetGeoTransform()
        projection = dataset.GetProjection()
        img = dataset.ReadAsArray()
        images.append(img)

    return images, labels, geoTransform, projection

# Function to plot histograms of data over time
def createHistogram(data):
    print("Creating Histograms...")
    bands, height, width = data[0].shape
    bands -= 1
    bandNames = ["Red", "Green", "Blue", "NIR", "SWIR1", "SWIR2"]
    bandColors = ["red", "green", "blue", "magenta", "aqua", "darkviolet"]
    fig, axs = plt.subplots(2, 3)
    axs = axs.ravel()
    for c in range(bands):
        values = []
        for image in data:
            for j in range(height):
                for k in range(width):
                    if image[6, j, k] > 0:
                        values.append(image[c, j, k])
        axs[c].hist(x=values, bins=10000, color=bandColors[c])
        axs[c].set_title(bandNames[c] + " Band")
        axs[c].set_xlabel("Intensity")
        axs[c].set_ylabel("Count")
    plt.show()

# Function to create simple cloud mask
#   Similar scheme to alpha band
#   0 -> Cloud Pixel
#   65535 -> Valid Pixel
#  [red, green, blue, nir, swir1, swir2, alpha, cloud mask]
def createCloudMask(data):
    print("Adding Cloud Mask...")
    output = []
    for image in data:
        bands, height, width = image.shape
        # By default, set all pixels to cloudless
        newImage = np.full((bands + 1, height, width), 65535)
        newImage[:bands, :, :] = image
        for j in range(height):
            for k in range(width):
                # Only check for clouds within alpha-valid pixels (necessary for calculating cloudiest image)
                if newImage[6, j, k] > 0:
                    NDSI = (newImage[1, j, k] - newImage[4, j, k]) / (newImage[1, j, k] + newImage[4, j, k])
                    ratio1 = (newImage[2, j, k] - newImage[4, j, k]) / (newImage[2, j, k] + newImage[4, j, k])
                    ratio2 = (newImage[2, j, k] - newImage[5, j, k]) / (newImage[2, j, k] + newImage[5, j, k])
                    # Check thresholds
                    if -.13 <= ratio1 and -.13 <= ratio2 and NDSI <= .4:
                        newImage[7, j, k] = 0
        output.append(newImage)
    return output

# Function that uses standard NDVI to determine greenest scene (does not include alpha/cloud masked pixels)
# NDVI = (NIR - Red) / (NIR + Red)
def findGreenest(data, labels):
    print("Finding Greenest Scene...")
    _, height, width = data[0].shape
    NDVI = []

    for image in data:
        # Create a mask that invalidates cloudy/zero-alpha pixels
        fullMask = np.logical_and((image[6, :, :] > 0), (image[7, :, :] > 0))
        NDVIraster = np.zeros((height, width))
        for j in range(height):
            for k in range(width):
                if fullMask[j, k]:
                    NDVIraster[j, k] = (image[3, j, k] - image[0, j, k]) / (image[3, j, k] + image[0, j, k])
        NDVIavg = np.sum(NDVIraster) / np.count_nonzero(fullMask == True)

        NDVI.append(NDVIavg)

    print(F"Greenest Scene: {labels[np.argmax(NDVI)]}")
    print(F"     Average NDVI: {max(NDVI)}")

# Function that uses standard NDSI to determine snowiest scene (does not include alpha/cloud masked pixels)
# NDSI = (Green - SWIR1) / (Green + SWIR1)
def findSnowiest(data, labels):
    print("Finding Snowiest Scene...")
    _, height, width = data[0].shape
    NDSI = []

    for image in data:
        # Create a mask that invalidates cloudy/zero-alpha pixels
        fullMask = np.logical_and((image[6, :, :] > 0), (image[7, :, :] > 0))
        NDSIraster = np.zeros((height, width))
        for j in range(height):
            for k in range(width):
                if fullMask[j, k]:
                    NDSIraster[j, k] = (image[1, j, k] - image[4, j, k]) / (image[1, j, k] + image[4, j, k])
        NDSIavg = np.sum(NDSIraster) / np.count_nonzero(fullMask == True)

        NDSI.append(NDSIavg)

    print(F"Snowiest Scene: {labels[np.argmax(NDSI)]}")
    print(F"     Average NDSI: {max(NDSI)}")

# Function that counts cloud masked pixels to determine cloudiest scene
# "NDCI" = Number of pixels with clouds / Number of valid pixels
def findCloudiest(data, labels):
    print("Finding Cloudiest Scene...")
    NDCI = []

    for image in data:
        NDCI.append(np.count_nonzero(image[7, :, :] == 0) / np.count_nonzero(image[6, :, :] == 65535))

    print(F"Cloudiest Scene: {labels[np.argmax(NDCI)]}")
    print(F"     Ratio of Cloud Masked Pixels: {max(NDCI)}")

# Function that uses relative luminance to determine brightest scene (does not include alpha/cloud masked pixels)
# "NDBI" (Relative Luminance) = .2126 * Red + .7152 * Green + .0722 * Blue
# Source for calculation: https://stackoverflow.com/questions/596216/formula-to-determine-brightness-of-rgb-color
def findBrightest(data, labels):
    print("Finding Brightest Scene...")
    _, height, width = data[0].shape
    NDBI = []

    for image in data:
        # Create a mask that invalidates cloudy/zero-alpha pixels
        fullMask = np.logical_and((image[6, :, :] > 0), (image[7, :, :] > 0))
        NDBIraster = np.zeros((height, width))
        for j in range(height):
            for k in range(width):
                if fullMask[j, k]:
                    NDBIraster[j, k] = .2126 * image[0, j, k] + .7152 * image[1, j, k] + .0722 * image[2, j, k]
        NDBIavg = np.sum(NDBIraster) / np.count_nonzero(fullMask == True)

        NDBI.append(NDBIavg)

    print(F"Brightest Scene: {labels[np.argmax(NDBI)]}")
    print(F"     Average Relative Luminance: {max(NDBI)}")

# Function to create Mean composite
def makeMean(data, path, geoTransform, projection):
    print("Creating Mean Composite...")
    bands, height, width = data[0].shape
    mean = np.zeros((bands, height, width))
    tot = np.zeros((bands, height, width))

    # Simply adds values of valid pixels to a numpy array while keeping track of how many valid values for each pixel
    for image in data:
        for j in range(height):
            for k in range(width):
                if image[6, j, k] > 0 and image[7, j, k] > 0:
                    mean[:, j, k] += image[:, j, k]
                    tot[:, j, k] += 1

    mean = mean / tot
    createTif(path + '\\' + "mean.tif", mean,
              width, height, bands, geoTransform, projection)

# Function to create Min composite
def makeMin(data, path, geoTransform, projection):
    print("Creating Min Composite...")
    bands, height, width = data[0].shape
    minIm = np.full((bands, height, width), 16384)

    # Iterates through each band of each pixel of each image to find the minimum value for each band in a pixel
    for image in data:
        for j in range(height):
            for k in range(width):
                if image[6, j, k] > 0 and image[7, j, k] > 0:
                    for c in range(bands):
                        minIm[c, j, k] = min(minIm[c, j, k], image[c, j, k])

    createTif(path + '\\' + "min.tif", minIm,
              width, height, bands, geoTransform, projection)

# Function to create Max composite
def makeMax(data, path, geoTransform, projection):
    print("Creating Max Composite...")
    bands, height, width = data[0].shape
    maxIm = np.zeros((bands, height, width))

    # Iterates through each band of each pixel of each image to find the maximum value for each band in a pixel
    for image in data:
        for j in range(height):
            for k in range(width):
                if image[6, j, k] > 0 and image[7, j, k] > 0:
                    for c in range(bands):
                        maxIm[c, j, k] = max(maxIm[c, j, k], image[c, j, k])

    createTif(path + '\\' + "max.tif", maxIm,
              width, height, bands, geoTransform, projection)

# Function to create Median composite
def makeMedian(data, path, geoTransform, projection):
    print("Creating Median Composite...")
    bands, height, width = data[0].shape
    median = np.zeros((bands, height, width))

    # Appends valid values of each band of each pixel to a list of lists, and sets the result equal to the median
    for j in range(height):
        for k in range(width):
            LoL = [[] for c in range(bands)]
            for image in data:
                if image[6, j, k] > 0 and image[7, j, k] > 0:
                    for c in range(bands):
                        LoL[c].append(image[c, j, k])
            for c in range(bands):
                median[c, j, k] = np.median(LoL[c])

    createTif(path + '\\' + "median.tif", median,
              width, height, bands, geoTransform, projection)

# Function to create Greenest composite
def makeGreenest(data, path, geoTransform, projection):
    print("Creating Greenest Composite...")
    bands, height, width = data[0].shape
    greenest = np.zeros((bands, height, width))

    # Appends NDVI values of each pixel to an array, finds the index of the max, and sets result equal to that image's
    # pixel value "packet" (package of 7 bands)
    for j in range(height):
        for k in range(width):
            NDVI = np.zeros((len(data),))
            i = 0
            for image in data:
                if image[6, j, k] > 0 and image[7, j, k] > 0:
                    NDVI[i] = (image[3, j, k] - image[0, j, k]) / (image[3, j, k] + image[0, j, k])
                # If pixel is invalid/cloudy, set NDVI below minimum value
                else:
                    NDVI[i] = -2
                i += 1
            greenest[:, j, k] = data[np.argmax(NDVI)][:, j, k]

    bands, height, width = greenest.shape
    createTif(path + '\\' + "greenest.tif", greenest,
              width, height, bands, geoTransform, projection)

# Function to create 85% Greenest composite
def make85Greenest(data, path, geoTransform, projection):
    print("Creating 85% Greenest Composite...")
    bands, height, width = data[0].shape
    greenest85 = np.zeros((bands, height, width))

    # Appends NDVI values of each pixel to an array, removes invalid values in the array, sorts the array, gets the
    # index of the 85th percentile, and sets result equal to that image's pixel value "packet" (package of 7 bands)
    for j in range(height):
        for k in range(width):
            NDVI = np.zeros((len(data),))
            i = 0
            for image in data:
                if image[6, j, k] > 0 and image[7, j, k] > 0:
                    NDVI[i] = (image[3, j, k] - image[0, j, k]) / (image[3, j, k] + image[0, j, k])
                # If pixel is invalid/cloudy, set NDVI below minimum value
                else:
                    NDVI[i] = -2
                i += 1
            sortedNDVI = np.sort(np.delete(NDVI, np.where(NDVI == -2)))

            greenest85[:, j, k] = data[np.where(NDVI == sortedNDVI[round(.85 * len(sortedNDVI)) - 1])[0][0]][:, j, k]

    bands, height, width = greenest85.shape
    createTif(path + '\\' + "greenest85.tif", greenest85,
              width, height, bands, geoTransform, projection)

# Abstract function to create a GeoTIFF composite from numpy array
def createTif(name, data, width, height, bands, geoTransform, projection):
    bands -= 1  # Exclude cloud mask
    driver = gdal.GetDriverByName('GTiff')
    tif = driver.Create(name, width, height, bands, gdal.GDT_UInt16)
    tif.SetGeoTransform(geoTransform)
    tif.SetProjection(projection)
    for band in range(bands):
        tif.GetRasterBand(band + 1).WriteArray(data[band, :, :])
    tif.FlushCache()
    tif = None
    print("Done, Output: " + name)

# Main execution to complete task 2
if __name__ == "__main__":
    images, labels, geoTransform, projection = loadImages(GeoTIF_dir)
    createHistogram(images)
    images = createCloudMask(images)

    findGreenest(images, labels)
    findSnowiest(images, labels)
    findCloudiest(images, labels)
    findBrightest(images, labels)

    # Create necessary directory
    try:
        os.mkdir(composites_dir)
    except:
        sys.exit("Task 2 Composites already exist")

    makeMean(images, composites_dir, geoTransform, projection)
    makeMin(images, composites_dir, geoTransform, projection)
    makeMax(images, composites_dir, geoTransform, projection)
    makeMedian(images, composites_dir, geoTransform, projection)
    makeGreenest(images, composites_dir, geoTransform, projection)
    make85Greenest(images, composites_dir, geoTransform, projection)


from dataclasses import asdict, dataclass
from enum import Enum
from typing import Callable, Sequence

import numpy as np
import pandas as pd
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score

from .ground_truths import SplitGroundTruths, TrainValTestGroundTruths
from .models import ModelNames
from .predictions import SplitPredictions, TrainValTestPredictions

@dataclass
class SplitScores:
    bust: float
    hip: float
    waist: float

@dataclass
class TrainValTestScores:
    train: SplitScores
    val: SplitScores
    test: SplitScores

@dataclass
class Metrics:
    mae: TrainValTestScores
    mse: TrainValTestScores
    r2: TrainValTestScores

class MetricNames(Enum):
    mae = "mae"
    mse = "mse"
    r2 = "r2"

class MetricsManager:
    def __init__(self) -> None:
        self.metrics: list[Metrics] = []
        self.model_names: list[ModelNames] = []
        self.metrics_table = pd.DataFrame([])

    def update_metrics(
        self,
        ground_truths: TrainValTestGroundTruths,
        predictions: TrainValTestPredictions,
    ) -> pd.DataFrame:
        mae = compute_train_val_test_metrics(
            mean_absolute_error, ground_truths, predictions
        )
        mse = compute_train_val_test_metrics(
            mean_squared_error, ground_truths, predictions
        )
        r2 = compute_train_val_test_metrics(r2_score, ground_truths, predictions)
        self.metrics.append(Metrics(mae, mse, r2))
        self.model_names.append(predictions.model_names)

        all_metric_dfs: list[pd.DataFrame] = []

        for metric in MetricNames:
            train_test_val_scores: TrainValTestScores = getattr(
                self.metrics[-1], metric.value
            )
            metric_df = self.get_single_metric_df(train_test_val_scores, metric)
            all_metric_dfs.append(metric_df)

        table_update = pd.concat(all_metric_dfs, axis=1)
        self.metrics_table = pd.concat([self.metrics_table, table_update])
        return self.metrics_table

    def get_single_metric_df(
        self, train_test_val_scores: TrainValTestScores, metric_name: MetricNames
    ) -> pd.DataFrame:
        metric_df = pd.DataFrame(asdict(train_test_val_scores))
        metric_df.columns = [col + f"_{metric_name.value}" for col in metric_df.columns]

        model_names = self.model_names[-1]
        if list(metric_df.index) == ["bust", "hip", "waist"]:
            metric_df.index = [model_names.bust, model_names.hip, model_names.waist]

        return metric_df

    def show_validation_metrics(self) -> pd.DataFrame:
        validation_columns = [
            col for col in self.metrics_table.columns if col.startswith("val")
        ]
        return self.metrics_table.loc[:, validation_columns]

def compute_train_val_test_metrics(
    score_func: Callable[[Sequence, np.ndarray], float],
    gound_truths: TrainValTestGroundTruths,
    predictions: TrainValTestPredictions,
) -> TrainValTestScores:
    train = compute_split_metrics(score_func, gound_truths.train, predictions.train)
    val = compute_split_metrics(score_func, gound_truths.val, predictions.val)
    test = compute_split_metrics(score_func, gound_truths.test, predictions.test)
    return TrainValTestScores(train, val, test)

def compute_split_metrics(
    score_func: Callable[[Sequence, np.ndarray], float],
    split_ground_truths: SplitGroundTruths,
    split_predictions: SplitPredictions,
) -> SplitScores:
    bust = score_func(split_ground_truths.bust, split_predictions.bust)
    hip = score_func(split_ground_truths.hip, split_predictions.hip)
    waist = score_func(split_ground_truths.waist, split_predictions.waist)

    return SplitScores(bust, hip, waist)


import os
import io
import json
import tarfile
import glob
import time
import logging

import base64

import boto3

from datetime import datetime

lambda_client = boto3.client('lambda')

def lambda_handler(event, context):

    now = datetime.now()
    dt_string = now.strftime("%Y/%m/%d %H:%M:%S")

    print("Sending stuff to Lambda_2")
    event = json.loads(event['body'])
    inputParams = {
        'device_id': event['device_id'],
        'company_name': event['company_name'],
        # 'time': event['time'],
        'time': dt_string,
        'image': event['base64str']
    }

    response = lambda_client.invoke(
        FunctionName = 'arn:aws:lambda:us-east-2:882207224999:function:pi_camera_ppe_detection_v1_L_2',
        InvocationType = 'RequestResponse',
        Payload = json.dumps(inputParams)
    )

    response = lambda_client.invoke(
        FunctionName = 'arn:aws:lambda:us-east-2:882207224999:function:pi_camera_ppe_detection_v1_L_3',
        InvocationType = 'RequestResponse',
        Payload = json.dumps(inputParams)
    )

    return {
        "statusCode": 200,
        "body": json.dumps({'Status': "Successful"})
    }


import math, numpy
import matplotlib.pyplot as plt

# Free off-diagonal density matrix
def rho_free(x, xp, beta):
    return (math.exp(-(x - xp) ** 2 / (2.0 * beta)) /
            math.sqrt(2.0 * math.pi * beta))

def rho_harmonic_trotter(grid, beta):
    return numpy.array([[rho_free(x, xp, beta) * \
                         numpy.exp(-0.5 * beta * 0.5 * (x ** 2 + xp ** 2)) \
                         for x in grid] for xp in grid])

def V(x,cubic,quartic):
	return 0.5*x*x + cubic*x*x*x+ quartic*x*x*x*x
def rho_a_trotter(grid, beta,c,q):
    return numpy.array([[rho_free(x, xp, beta) * \
                         numpy.exp(-0.5 * beta *(V(x,c,q)+ V(xp,c,q))) \
                         for x in grid] for xp in grid])

def pi_quant(x,beta):
	a = math.sqrt(math.tanh(0.5*beta)/math.pi)
	return a*math.exp(-x*x*math.tanh(0.5*beta))
quartic = 1.
cubic = -quartic
x_max = 5.
nx = 100
dx = 2.0 * x_max / (nx - 1)
x = [i * dx for i in range(-(nx - 1) / 2, nx / 2 + 1)]
beta_tmp = 2.0 ** (-6)                   # initial value of beta (power of 2)
beta     = 2.0 ** 2                      # actual value of beta (power of 2)
rho = rho_a_trotter(x, beta_tmp,cubic,quartic)  # density matrix at initial beta
while beta_tmp < beta:
    rho = numpy.dot(rho, rho)
    rho *= dx
    beta_tmp *= 2.0

Z = sum(rho[j,j] for j in range(nx+1))*dx
pi_of_x = [ rho[j,j]/ Z for j in range(nx+1)]
f = open('data_aharm_matrixsquare_beta'+str(beta)+'.dat','w')
for j in range(nx+1):
	f.write(str(x[j])+' '+str(rho[j,j]/Z)+'\n')
f.close()

plt.scatter(x,pi_of_x,label="mat. sq.")
plt.title("beta="+str(beta)+"   cubic="+str(cubic)+"    quartic="+str(quartic))
plt.legend()
plt.show()


import pandas as pd
import os

def wide2long(inputfile, subID, ses):

    c = pd.read_csv(inputfile)
    c["subID"] = subID
    c["ind"] = range(1,len(c)+1)
    c["ses"] = ses
    ## change wide form to long
    return pd.melt(c,id_vars=["subID", "ind", "ses"], var_name="TCK")


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""Bash models for use in scripts."""
from uuid import uuid4
from typing import Union, Optional, List

from pydantic import BaseModel, Field, validator, UUID4

DOCUMENTATION_INDICATION_STR = '#/'
USAGE_STRING = '# usage'

class BashDocumentationLine(BaseModel):
    """Abstraction for documentation in bash script."""
    module_name: str
    documentation: Union[str, List[str]]
    markers: bool = True

    def __str__(self) -> str:
        """Representation of Bash Documentation."""
        contents = DOCUMENTATION_INDICATION_STR + 'Usage:\n\t./' + \
            self.module_name + ' [OPTIONS]\n'
        if isinstance(self.logic, str):
            contents += self.logic
        elif isinstance(self.logic, list):
            for item in self.logic:
                contents += item + '\n'
        markers = ''
        end_markers = ''
        if markers is True:
            markers += ' {{{1'
            end_markers += '# 1}}}'
        return USAGE_STRING + markers + '\n' + DOCUMENTATION_INDICATION_STR +\
            contents + end_markers + '\n'

# pylint: disable=too-few-public-methods
class BashVariable(BaseModel):
    """Construct a BashVariable."""
    key: str
    value: str = ""

    def __str__(self) -> str:
        """Representation of BashVariable."""
        return f'{self.key}={self.value}'

class BashFunction(BaseModel):
    """Construct a Bash function."""
    name: str
    logic: Union[list, str]
    markers: bool = False

    def __str__(self) -> str:
        """Representation of BashFunction."""
        contents = ""
        if isinstance(self.logic, str):
            contents += self.logic
        elif isinstance(self.logic, list):
            for item in self.logic:
                contents += ('\t' + item + '\n')
        markers = ""
        end_markers = ""
        if self.markers is True:
            markers += ' {{{2'
            end_markers += ' # 2}}}'
        return self.name + '() {' + markers + '\n' + contents + '}' + \
            end_markers + '\n'

class BashDocumentation(BaseModel):
    """Provide a Bash documentation script strint."""

class BashEnvironment(BaseModel):
    """Populate the BashEnvironment class."""
    environment_variables: List[BashVariable]
    bash_functions: List[BashFunction]

    class Config:
        """Options to the BaseModel config class."""
        allow_arbitrary_types = True

if __name__ == "__main__":
    print(BashVariable(key='LOGLEVEL', value='INFO'))
    print(BashFunction(name='test_function',
                       logic='ps aux\n',
                       markers=True))


#-*- coding: utf-8 -*-
#!/usr/bin/python3
import os
import sqlite3

patient = input('Name of the person you want to search: ' + str())

conn = sqlite3.connect('../db/pacientes.db')

try:
    c = conn.cursor()
    c.execute('SELECT id FROM pacientes WHERE nombre = ?', (str(patient),))
    namedb = str(c.fetchone()[0]) + '.db'
    conn.commit()
    conn.close()
    if os.path.isfile('../db/registros/'+namedb):
        conn = sqlite3.connect('../db/registros/'+namedb)

        c = conn.cursor()

        c.execute('SELECT * FROM datos')

        print(c.fetchall())

        conn.commit()
        conn.close()
    else:
        print('No records founded')

except TypeError:
    print('No records founded')


"""Returns bots ping, uptime and current load.
"""
from datetime import datetime
from math import ceil
from subprocess import check_output

def handle_update(bot, update, **kwargs):
    sent_time = update.message.date
    now = datetime.now()
    delta = now - sent_time

    ping = ceil(delta.total_seconds() * 1000)
    uptime = check_output(['uptime', '-p']).decode().split(' ', 1)[1]

    output = check_output(['cat', '/proc/loadavg']).split()[1].decode()
    load_5m = float(output) * 100

    report = ("Ping: {}ms\nBot uptime: {}"
              "Load over last 5 minutes: {}%").format(ping, uptime, load_5m)

    bot.sendMessage(update.message.chat_id, text=report)


#!/usr/bin/env python3
"""
Convert Southern Poverty Law Center Cryptocurrency Report to a TagPack.
"""

import os
import csv
from datetime import datetime, date
from typing import List

import yaml

CURRENCY = {
    'Bitcoin Addresses': 'BTC',
    'Ethereum Addresses': 'ETH',
    'Litecoin Addresses': 'LTC',
    'Monero Address': 'XMR'
}

class RawData:
    """
    Download and read data provided by the source.
    """
    def __init__(self, fn: str, url: str):
        self.fn = fn
        self.url = url

    def download(self):
        from urllib.request import urlretrieve
        urlretrieve(self.url, self.fn)

    def read(self) -> List[dict]:
        with open(self.fn, 'r',encoding='utf-8', newline='') as csvfile:
            reader = csv.DictReader(csvfile, delimiter=',', quotechar='"')
            return [row for row in reader]

class TagPackGenerator:
    """
    Generate a TagPack from SPLC Cryptocurrency Report.
    """

    def __init__(self, rows: List[dict], title: str, creator: str, description: str, lastmod: date, source: str):
        self.rows = rows
        self.data = {
            'title': title,
            'creator': creator,
            'description': description,
            'lastmod': lastmod,
            'abuse': 'extremism',
            'category': 'user',  # like in the OFAC TagPack generator
            'source': source,
            'tags': []
        }
        self.source = source

    def generate(self):
        tags = []
        for row in self.rows:
            for column in ['Bitcoin Addresses', 'Ethereum Addresses', 'Litecoin Addresses', 'Monero Address']:
                for address in row[column].split('\n'):
                    if not address or address == 'uses':
                        continue
                    tag = {
                        'address': address.strip(),
                        'currency': CURRENCY[column],
                        'label': row['Entity']
                    }
                    tags.append(tag)
        self.data['tags'] = tags

    def saveYaml(self, fn: str):
        with open(fn, 'w', encoding='utf-8') as f:
            f.write(yaml.dump(self.data, sort_keys=False))

if __name__ == '__main__':
    with open('config.yaml', 'r') as config_file:
        config = yaml.safe_load(config_file)

    raw_data = RawData(config['RAW_FILE_NAME'], config['URL'])
    if not os.path.exists(config['RAW_FILE_NAME']):
        raw_data.download()

    last_mod = datetime.fromtimestamp(os.path.getmtime(config['RAW_FILE_NAME'])).date()
    generator = TagPackGenerator(raw_data.read(), config['TITLE'], config['CREATOR'], config['DESCRIPTION'],
                                 last_mod, config['SOURCE'])
    generator.generate()
    generator.saveYaml(config['TAGPACK_FILE_NAME'])


import json
from typing import Optional

from sommelier.adapters.sockets import (
    WebSocketReader, WebSocketWriter
)
from sommelier.behave_wrapper import FlowListener, response_json_holder
from sommelier.behave_wrapper.tables import carpenter
from sommelier.utils import require_var

class WSocketManager(FlowListener):
    """
    This is a one way web socket manager
    Manager allows to do read, write or both but via 2 different ws
    Tests written in behave are sequential and WS manager fits into this idea
    """

    def __init__(self, host):
        super().__init__()
        require_var(host, "host")
        self.host = host
        self.reader: Optional[WebSocketReader] = None
        self.writer: Optional[WebSocketWriter] = None

    def after_all(self):
        super().after_all()
        self.disconnect_all()

    def create_socket(self, reader):
        cookies = carpenter.builder().double().dict()
        if reader:
            self.reader = WebSocketReader(self.host, cookies)
            self.reader.schedule_start()
        else:
            self.writer = WebSocketWriter(self.host, cookies)
            self.writer.schedule_start()

    def read_data(self):
        response_json_holder.with_description("WS-Read", self.host)
        text = self.reader.read_one()
        data = None
        if text is not None:
            data = json.loads(text)
        response_json_holder.hold(data)

    def write_data(self):
        data = carpenter.builder().double().dict()
        self.writer.write(data)

    def disconnect_all(self):
        if self.reader is not None:
            self.reader.schedule_shutdown()
        if self.writer is not None:
            self.writer.schedule_shutdown()


# ["1", "-", "3", "+", "5", "-", "8"] 배열의 모든 계산 경우는 [-15, -5, -5, 1, 1]이다.
# 괄호로 인한 결합법칙 경우의 수 중에 최대값 구하기.

def solution(arr):
    minv = 0
    maxv = 0
    sumv = 0
    for idx in range(len(arr) - 1, -1, -1):
        if arr[idx] == '+':
            continue
        elif arr[idx] == '-':
            temp_min = minv
            temp_max = maxv
            minv = min(-(sumv + temp_max), -sumv + temp_min)
            minus_v = int(arr[idx + 1])
            maxv = max(-(sumv + temp_min), -minus_v + (sumv - minus_v) + temp_max)
            sumv = 0
        elif int(arr[idx]) >= 0:
            sumv += int(arr[idx])
    maxv += sumv # 첫 수는 반드시 양수이므로 더해준다
    return maxv

print(solution(["1", "-", "3", "+", "5", "-", "8"])) # 1
print((solution(["5", "-", "3", "+", "1", "+", "2", "-", "4"]))) # 3

'''
뒤에서부터 계산하는 이유는 결합 법칙 특성상 마지막 수는 반드시 앞의 부호를 따르기 때문이다.
반대로 첫 수는 반드시 양수임으로 마지막에 더 해준다.
핵심은 작은 수에서 큰 수를 빼는 게 아니라, 큰 수에서 작은 수를 빼도록 처리하는 것.
'''

프로그래머스/동적계획법/사칙연산.py


import os, sys
try:                                            # if running in CLI
    cur_path = os.path.abspath(__file__)
except NameError:                               # if running in IDE
    cur_path = os.getcwd()

while cur_path.split('/')[-1] != 'portfolio':
    cur_path = os.path.abspath(os.path.join(cur_path, os.pardir))
sys.path.insert(1, os.path.join(cur_path, 'lib', 'python3.7', 'site-packages'))

from _connections import db_connection
from datetime import datetime
from progress_bar import progress
import numpy as np
from rh_psql import pg_query, pg_insert

#for script in create_tables['rsi_day_14']:
#    PSQL.client.execute(script)
#    PSQL.client.execute("commit;")

def det_cur_ma(psql):
#    psql = PSQL
    script = "select max(sto_osc_14_id) from portfolio.sto_osc_14"
    _cur = pg_query(psql.client, script)
    _nxt_ids = _cur.values[0][0] + 1

    script = "select rh_id, max(date) from portfolio.sto_osc_14 group by rh_id"
    _cur_comp = pg_query(PSQL.client, script)
    _comp_cur = {k:v for k,v in _cur_comp.values}
    return(_nxt_ids, _comp_cur)

def det_cur_cci(psql):
#    psql = PSQL
    script = "select max(cci_day_20_id) from portfolio.cci_day_20"
    _cur = pg_query(psql.client, script)
    _nxt_ids = _cur.values[0][0] + 1

    script = "select rh_id, max(date) from portfolio.cci_day_20 group by rh_id"
    _cur_comp = pg_query(PSQL.client, script)
    _comp_cur = {k:v for k,v in _cur_comp.values}
    return(_nxt_ids, _comp_cur)

def det_cur_rsi(psql):
#    psql = PSQL
    script = "select max(rsi_day_14_id) from portfolio.rsi_day_14"
    _cur = pg_query(psql.client, script)
    _nxt_ids = _cur.values[0][0] + 1

    script = "select rh_id, max(date) from portfolio.rsi_day_14 group by rh_id"
    _cur_comp = pg_query(PSQL.client, script)
    _comp_cur = {k:v for k,v in _cur_comp.values}
    return(_nxt_ids, _comp_cur)

def rsi(comp_idx, comp_name_conv):
#    comp_idx, comp_name_conv = COMP_IDX, COMP_NAME_CONV

#    nxt_id = 0
    nxt_id, comp_cur = det_cur_rsi(PSQL)
    total_stocks = len(comp_idx)
    cur_date = pg_query(PSQL.client, "SELECT max(date) FROM portfolio.rsi_day_14;")[0].values[0]
    for stock_num, (rh_id) in enumerate(comp_idx):
        progress(stock_num, total_stocks, status = comp_name_conv[rh_id])

        if rh_id in comp_cur.keys() and cur_date <= np.datetime64(comp_cur[rh_id]):
            continue

        comp_data = pg_query(PSQL.client, "SELECT date, close_price FROM portfolio.day_prices where rh_id = '%s';" % (rh_id))
        comp_data.rename(columns = {0:'date', 1:'close_price'}, inplace = True)
        comp_data.sort_values('date', ascending = True, inplace = True)
        comp_data['diff'] = comp_data['close_price'].diff()
        comp_data['gains'] = comp_data['diff'].apply(lambda x: x if x > 0 else 0)
        comp_data['losses'] = comp_data['diff'].apply(lambda x: x if x < 0 else 0)
        comp_data.dropna(inplace = True)
        all_avg_gain = []
        all_avg_loss = []
        dates = []
        all_avg_gain.append(comp_data.iloc[:14]['gains'].mean())
        all_avg_loss.append(comp_data.iloc[:14]['losses'].mean())
        dates.append(comp_data.iloc[14]['date'])
        comp_data = comp_data.iloc[15:]
        for _date, cur_gain, cur_loss in comp_data[['date', 'gains', 'losses']].values:
            dates.append(_date)
            all_avg_gain.append((all_avg_gain[-1]*13 + cur_gain) / 14)
            all_avg_loss.append((all_avg_loss[-1]*13 + cur_loss) / 14)

        for date, gain, loss in zip(dates, all_avg_gain, all_avg_loss):
            if gain == 0:
                rsi = 0
            elif loss == 0:
                rsi = 100
            else:
                rsi = 100 - (100/1+(gain/loss))

            if rh_id in comp_cur.keys() and comp_cur[rh_id] >= date:
                continue

            script = "INSERT INTO portfolio.rsi_day_14(rsi_day_14_id, rh_id, date, rsi) VALUES (%i, '%s', '%s', %.3f);" % (nxt_id, rh_id, datetime.strptime(str(date).split('.')[0], '%Y-%m-%d %H:%M:%S'), rsi)
            pg_insert(PSQL.client, script)
            nxt_id += 1
            comp_cur[rh_id] = date

def cci(comp_idx, comp_name_conv):
#    comp_idx, comp_name_conv = COMP_IDX, COMP_NAME_CONV
    nxt_id, comp_cur = det_cur_cci(PSQL)
    total_stocks = len(comp_idx)
    cur_date = pg_query(PSQL.client, "SELECT max(date) FROM portfolio.cci_day_20;")[0].values[0]
    for stock_num, (rh_id) in enumerate(comp_idx):
        progress(stock_num, total_stocks, status = comp_name_conv[rh_id])

        if rh_id in comp_cur.keys() and cur_date <= np.datetime64(comp_cur[rh_id]):
            continue

        comp_ma = pg_query(PSQL.client, "SELECT date, avg_price from portfolio.ma_day_20 where rh_id = '%s'" % (rh_id))
        comp_ma.rename(columns = {0:'date', 1:'ma'}, inplace = True)
        comp_ma.sort_values('date', ascending = True, inplace = True)
        comp_ma.set_index('date', inplace = True)
        comp_typ_price = pg_query(PSQL.client, "SELECT date, (high_price + low_price + close_price)/3 from portfolio.day_prices where rh_id = '%s'" % (rh_id))
        comp_typ_price.rename(columns = {0:'date', 1:'typical'}, inplace = True)
        comp_typ_price.sort_values('date', ascending = True, inplace = True)
        comp_typ_price.set_index('date', inplace = True)
        comp_cci = comp_typ_price.join(comp_ma)
        comp_cci['typ_avg'] = comp_cci['typical'].rolling(window = 20).mean()
        comp_cci['abs_dev'] = abs(comp_cci['typical'] - comp_cci['typ_avg'])
        comp_cci['mean_dev'] = comp_cci['abs_dev'].rolling(window = 20).mean()
        comp_cci['cci'] = (comp_cci['typical'] - comp_cci['ma']) / (comp_cci['mean_dev'] * .015)
        comp_cci.dropna(inplace = True)
        comp_cci.reset_index(inplace = True)
        for date, cci in comp_cci[['date', 'cci']].values:
            if rh_id in comp_cur.keys() and comp_cur[rh_id] >= date:
                continue
            if cci != cci or cci == np.inf or cci == -np.inf:
                continue
            script = "INSERT INTO portfolio.cci_day_20(cci_day_20_id, rh_id, date, cci) VALUES (%i, '%s', '%s', %.3f);" % (nxt_id, rh_id, datetime.strptime(str(date).split('.')[0], '%Y-%m-%d %H:%M:%S'), cci)
            pg_insert(PSQL.client, script)
            nxt_id += 1
            comp_cur[rh_id] = date

def stoch_osc_k(comp_idx, comp_name_conv):
#    comp_idx, comp_name_conv = COMP_IDX, COMP_NAME_CONV
#    all_next_id, all_comp_cur = det_cur_ema(PSQL)

    cur_date = pg_query(PSQL.client, "SELECT max(date) FROM portfolio.sto_osc_14;")[0].values[0]

    total_stocks = len(comp_idx)
    nxt_id, comp_cur = det_cur_ma(PSQL)
    for stock_num, (rh_id) in enumerate(comp_idx):
        progress(stock_num, total_stocks, status = comp_name_conv[rh_id])

        if rh_id in comp_cur.keys() and cur_date <= np.datetime64(comp_cur[rh_id]):
            continue

#        if rh_id in comp_cur.keys():
#            missing_days = pg_query(PSQL.client, "SELECT date, close_price FROM portfolio.day_prices where rh_id = '%s' and date > '%s';" % (rh_id, all_comp_cur[period][rh_id]))
#            if len(missing_days) == 0:
#                continue
#        else:
        comp_data = pg_query(PSQL.client, "SELECT date, close_price FROM portfolio.day_prices where rh_id = '%s';" % (rh_id))
        comp_data.rename(columns = {0:'date', 1:'close_price'}, inplace = True)
        comp_data.sort_values('date', ascending = True, inplace = True)

        per_high = comp_data['close_price'].rolling(window=14).max()
        per_low = comp_data['close_price'].rolling(window=14).min()

        for date, current_close, lowest_low, highest_high in zip(comp_data.date.values, comp_data.close_price.values, per_low, per_high.values):
            if rh_id in comp_cur.keys() and comp_cur[rh_id] >= date:
                continue

            k = (current_close - lowest_low)/(highest_high - lowest_low) * 100
            if k != k:
                continue
            script = "INSERT INTO portfolio.sto_osc_14(sto_osc_14_id, rh_id, date, k) VALUES (%i, '%s', '%s', %.2f);" % (nxt_id, rh_id, datetime.strptime(str(date).split('.')[0], '%Y-%m-%dT%H:%M:%S'), k)
            pg_insert(PSQL.client, script)
            nxt_id += 1
            comp_cur[rh_id] = date

if __name__ == '__main__':
    PSQL = db_connection('psql')

    COMP_IDX = [i[0] for i in pg_query(PSQL.client, "SELECT rh_id FROM portfolio.stocks;").values]
    COMP_NAME_CONV = pg_query(PSQL.client, 'select rh_id, rh_sym from portfolio.stocks')
    COMP_NAME_CONV = {k:v for k,v in COMP_NAME_CONV.values}

    cci(COMP_IDX, COMP_NAME_CONV)
    rsi(COMP_IDX, COMP_NAME_CONV)