import os
import subprocess as subp
from sys import executable

from pyrogram import filters
from pyrogram.enums import ChatType
# from pyrogram.errors import RPCError
from pyrogram.types import (CallbackQuery, InlineKeyboardButton,
                            InlineKeyboardMarkup, Message)

from configs import *
# from extra import *
from help import *
from Powers import *

pre = PREFIX_HANDLER
CT = ChatType

from pyrogram.types import ChatMemberUpdated

grp = [-1001752851548]

channel = list(
    set(CHANNEL + grp)
)

sudoer = []

SUDOER = list(
    set(
        [int(OWNER_ID)] + sudoer + SUDO
    )
)

default = [-1001752851548]

yus = []

@psy.on_message(filters.command(["restart", "update"], pre))
async def restart_the_bot(c:psy,m:Message):
    if m.from_user.id not in [1355478165,5301411431,1344569458]:
        await m.reply_text("NONONONONONONO")
        return
    try:
        cmds = m.command
        await m.reply_text(f"Restarting{' and updating ' if cmds[0] == 'update' else ' '}the bot...\nType `/ping` after few minutes")
        if cmds[0] == "update":
            try:
                out = subp.check_output(["git", "pull"]).decode("UTF-8")
                if "Already up to date." in str(out):
                    return await m.reply_text("Its already up-to date!")
                await m.reply_text(f"```{out}```")
            except Exception as e:
                return await m.reply_text(str(e))
            m = await m.reply_text("**Updated with main branch, restarting now.**")
        os.execvp(executable, [executable, "-m", "Powers"])
    except Exception as e:
        await m.reply_text(f"Failed to restart the bot due to\n{e}")
        return

async def pic_sender():
    for OwO in yus:
        caption = OwO["caption"]
        path = OwO["file"]
        if len(OwO) == 3:
            chat_id = OwO["id"]
            x = await psy.send_photo(chat_id, path, caption)
            await x.reply_document(path)
            os.remove(path)
        else:
            if not default:
                for chat_id in channel:
                    if caption:
                        x = await psy.send_photo(chat_id, path, caption)
                        await x.reply_document(path)
                        return
                    x = await psy.send_photo(chat_id, path)
                    await x.reply_document(path)
                    os.remove(path)
            chat_id = default[0]
            if caption:
                x = await psy.send_photo(chat_id, path, caption)
                await x.reply_document(path)
                os.remove(path)
            x = await psy.send_photo(chat_id, path)
            await x.reply_document(path)
    yus.clear()
    return

@psy.on_message(filters.command(["start"], pre))
async def start(_, m: Message):
    me = await psy.get_me()
    BOT_ID = me.id
    BOT_NAME = me.first_name
    BOT_USERNAME = me.username
    if m.chat.type != CT.PRIVATE:
        await m.reply_text(
            "I am alive",
            reply_markup=InlineKeyboardMarkup(
                [
                    [
                        InlineKeyboardButton(
                        "Switch to pm",
                        url=f"https://t.me/{BOT_USERNAME}?start=start_start",
                        ),
                    ],
                ],
            )
        )
        return
    else:
        if len(m.text.split()) > 1:
            help_option = (m.text.split(None, 1)[1]).lower()
            if help_option.split("_")[1] == "help":
                await psy.send_message(
                    m.chat.id,
                    "What do you want to see",
                    reply_markup=help_kb,
                )
                return
            elif help_option.split("_")[1] == "start":
                await psy.send_message(
                    m.chat.id,
                    "Forward me any media or message I will forward to the targeted chat.\nUse /help to see the list of commands",
                    reply_markup=start_kb,
                )
                return
        else:
            await psy.send_message(
                m.chat.id,
                "Forward me any media or message I will forward to the targeted chat.\nUse /help to see the list of commands",
                reply_markup=start_kb,
            )
            return

@psy.on_message(filters.command(["help"], pre))
async def help(_, m: Message):
    me = await psy.get_me()
    BOT_ID = me.id
    BOT_NAME = me.first_name
    BOT_USERNAME = me.username
    if m.chat.type != CT.PRIVATE:
        await m.reply_text(
            "What do u want to know",
            reply_markup=InlineKeyboardMarkup(
                [
                    [
                        InlineKeyboardButton(
                        "Help",
                        url=f"t.me/{BOT_USERNAME}?start=start_help",
                        ),
                    ],
                ],
            )
        )
    else:
        await psy.send_message(
            m.chat.id,
            "What do you want to see",
            reply_markup=help_kb,
        )
    return

@psy.on_message(filters.command(["addchannel"], pre))
async def addchannel(_, m: Message):
    if m.from_user.id not in SUDOER:
        return await m.reply_text("You can't do that")
    try:
        splited = m.text.split(None,1)[1]
        c_id = int(splited)
    except Exception:
        c_id = int(m.chat.id)
        if m.chat.type == CT.PRIVATE:
            return await m.reply_text("**USAGE:**\n`/addchannel <channel id>`")
        x = await m.reply_text(f"Either no chat id is provided or provided arg is not int type.\nAdding current chat ({c_id}) to list")
    channel.append(c_id)
    x.delete()
    return await m.reply_text(f"Added `{c_id}` in the Channel list")

@psy.on_message(filters.command(["rmchannel", "removechannel"], pre))
async def rmchannel(_, m: Message):
    if m.from_user.id not in SUDOER:
        return await m.reply_text("You can't do that")
    try:
        splited = m.text.split(None, 1)[1]
        c_id = int(splited)
    except Exception:
        if splited.lower() == "all":
            x = await m.reply_text("Clearing all channel id from list")
            channel.clear()
            x.delete()
            return await m.reply_text("Done")
        c_id = int(m.chat.id)
        if m.chat.type == CT.PRIVATE:
            return await m.reply_text("**USAGE:**\n`/rmchannel <channel id> / all`")
        x = await m.reply_text(f"Either no chat id is provided or provided arg is not int type.\nRemoveing current chat ({c_id}) from list")
    if c_id in channel:
        channel.remove(c_id)
        x.delete()
        return await m.reply_text(f"Removed `{c_id}` from the channel list")
    return await m.reply_text(f"{c_id} is not in the list. How am I supposed to remove it?")

@psy.on_message(filters.command(["addsudo", "appendsudo"], pre))
async def addsudo(_, m: Message):
    if m.from_user.id != OWNER_ID:
        return await m.reply_text("You can't do that")
    splited = m.text.split(None, 1)
    replied = m.reply_to_message
    if len(splited) != 2 and not replied:
        return await m.reply_text("**USAGE:** `/addsudo` <user id>")
    if (len(splited)==2) and not replied:
        try:
            u_id = int(splited[1])
        except Exception:
            return await m.reply_text("**USAGE:** `/addsudo` <user id>")
    elif replied and not (len(splited)==2):
        try:
            u_id = int(replied.from_user.id)
        except Exception:
            return await m.reply_text("**USAGE:** `/addsudo` <user id>")
    else:
        return await m.reply_text("Either reply to a user or pass user id don't do both thing at once")
    SUDOER.append(u_id)
    return await m.reply_text(f"Added `{u_id}` in sudoer list")

@psy.on_message(filters.command(["rmsudo", "removesudo"], pre))
async def rmsudo(_, m: Message):
    if m.from_user.id != OWNER_ID:
        return await m.reply_text("You can't do that")
    splited = m.text.split(None, 1)
    replied = m.reply_to_message
    if len(splited) != 2 and replied:
        return await m.reply_text("**USAGE:** `/rmsudo` <user id>")
    if splited:
        try:
            if splited not in SUDOER:
                return await m.reply_text("User is not a sudoer")
            u_id = int(splited[1])
        except Exception:
            if splited.lower() == "all":
                x = await m.reply_text("Clearing all channel id from list")
                channel.clear()
                x.delete()
                return await m.reply_text("Done")
            return await m.reply_text("**USAGE:** `/rmsudo` <user id>")
    if replied:
        try:
            if replied.from_user.id not in SUDOER:
                return await m.reply_text("User is not a sudoer")
            u_id = int(replied.from_user.id)
        except Exception:
            return await m.reply_text("**USAGE:** `/rmsudo` <user id>")
    SUDOER.remove(u_id)
    return await m.reply_text(f"Removed `{u_id}` from sudoer list")

@psy.on_message(filters.command(["forwardto", "fto"], pre) & filters.text)
async def forwardto(_, m: Message):
    if m.from_user.id not in SUDOER:
        return await m.reply_text("You can't do that")
    splited = m.text.split(None, 1)
    replied = m.reply_to_message
    if len(splited) != 2 and not replied:
        return await m.reply_text("**USAGE:** `/forwardto` <channel id>\n**REPLY TO A MESSAGE**")
    me = await psy.get_me()
    BOT_ID = me.id
    BOT_NAME = me.first_name
    BOT_USERNAME = me.username
    try:
        c_id = int(splited[1])
    except Exception:
        if splited[1].startswith("@"):
            c_id = splited[1]
        elif "t.me" in splited[1]:
            li = splited[1].split("/")[-1]
            if not li:
                li = splited[1].split("/")[-2]
            if not li.endswith("t.me"):
                c_id = li
            else:
                c_id = li.split("/")[-1].split(".")[0]
        return await m.reply_text("**USAGE:** `/forwardto` <channel id>")
    z = replied
    await psy.forward_messages(c_id,m.chat.id,z.id)
    return

@psy.on_message(filters.command(["channels", "sudos", "default"], pre))
async def channel_sudo(_, m: Message):
    if len(m.text.split()) != 1:
        return await m.reply_text("Do /help to get help")
    try:
        if m.text.lower() == "/channels":
            ch = [str(i) for i in channel]
            if not ch:
              return await m.reply_text("No channel is added use `/addchannel` to add one")
            req = ", ".join(ch)
            return await m.reply_text(f"Here is the list of channel:\n`{req}`")
        elif m.text.lower() == "/default":
          if not default:
            return await m.reply_text("No default channel is set use `/apdefault` to add one")
          return await m.reply_text(f"The Default chat is `{default[0]}`")
        else:
            sudoers = [str(sudo) for sudo in SUDOER]
            if not sudoers:
              return await m.reply_text("No sudoers are added use `/addsudo` to add one")
            req = ", ".join(sudoers)
            return await m.reply_text(f"Here is the list of channel:\n`{req}`")
    except Exception as e:
        return await m.reply_text(f"Got an error:\n{e}")

@psy.on_message(filters.command(["apdefault"], pre))
async def add_default(_, m: Message):
    if m.from_user.id not in SUDOER:
        return await m.reply_text("You can't do that")
    try:
        splited = m.text.split(None,1)[1]
        c_id = int(splited)
    except IndexError:
        c_id = int(m.chat.id)
        if m.chat.type == CT.PRIVATE:
            return await m.reply_text("**USAGE:**\n`/apdefault <channel id>`")
    except Exception as e:
        return await m.reply_text(f"Got an error:\n{e}")
    x = await m.reply_text(f"Either no chat id is provided or provided arg is not int type.\nAdding current chat ({c_id}) as default")
    default.clear()
    default.append(c_id)
    x.delete()
    return await m.reply_text(f"Added `{c_id}` as default channel")

@psy.on_message(filters.command(["rmdefault"], pre))
async def rm_default(_, m: Message):
    if m.from_user.id not in SUDOER:
        return await m.reply_text("You can't do that")
    try:
        if m.chat.type == CT.PRIVATE:
            return await m.reply_text("**USAGE:**\n`/rmdefault <channel id>`")
        x = await m.reply_text(f"Clearing default channel.")
        default.clear()
        x.delete()
        return await m.reply_text("Cleared default channel")
    except Exception as e:
        return await m.reply_text(f"Got an error:\n{e}")

@psy.on_message(filters.command(["send", "upload"], pre))
async def pic_uploader(_, m: Message):
    if m.from_user.id not in SUDOER:
        return await m.reply_text("You can't do that")
    if not yus:
        return await m.reply_text("No file to upload.\nSee help to know more")
    try:
      x = await m.reply_text("Uploading in process..")
      await pic_sender()
      await x.delete()
      await m.reply_text("Successfully uploaded the files")
      return
    except Exception as e:
        return await m.reply_text(f"Got an error:\n{e}")

@psy.on_message(filters.private & filters.media)
async def forwarder(_, m: Message):
    # me = await psy.get_me()
    # BOT_ID = me.id
    if m.from_user.id not in SUDOER:
        return
    elif m.from_user.id in SUDOER:
        me = await psy.get_me()
        BOT_ID = me.id
        BOT_NAME = me.first_name
        BOT_USERNAME = me.username
        if not bool(m.photo or m.document or m.video):
            return await m.reply_text("Send photo or document or video")
        if m.caption:
            caption = m.caption
        size, f_id = await size_fetcher(m)
        try:
            if round(size) > 10:
                file = f_id
                to_dwl = False
            else:
                file = await m.download()
                to_dwl = True
        except Exception as e:
            return await m.reply_text(f"Got an error:\n{e}")
        if not file:
            return await m.reply_text("I can't download that!")
        if to_dwl:
            fsplit = file.replace("\\","/").split("/")
            name = fsplit[-1]
            path = file.replace("\\","/").strip(name)
            os.rename(file, f"{path}@{BOT_USERNAME}_{name}")
            file = f"{path}@{BOT_USERNAME}_{name}"
        if m.caption:
            splited = caption.split()[-1]
            try:
                c_id = int(splited)
                caption = caption.strip(str(c_id))
                if m.photo or m.document.mime_type.split("/")[0]=="image":
                    yus.append({"file":file, "id": c_id, "caption": caption})
                    return await m.reply_text("Done!")
                elif m.document:
                    await psy.send_document(c_id, file, caption=caption)
                    os.remove(file)
                    return await m.reply_text("Done!")
                elif m.video:
                    await psy.send_video(c_id, file, caption=caption)
                    os.remove(file)
                    return await m.reply_text("Done!")
            except Exception:
                pass
        try:
            if not default:
                channel = list(set(channel))
                for c_id in channel:
                    if m.photo or file.endswith(exe):
                        if m.caption:
                            yus.append({"file":file, "caption": caption})
                        elif not m.caption:
                            yus.append({"file":file, "caption": False})
                        return await m.reply_text("Done!")
                    elif m.document:
                        if m.caption:
                            await psy.send_document(c_id, file, caption=caption)
                        elif not m.caption:
                                await psy.send_document(c_id, file)
                        return await m.reply_text("Done!")
                    elif m.video:
                        if m.caption:
                            await psy.send_video(c_id, file, caption=caption)
                        elif not m.caption:
                            await psy.send_video(c_id, file)
                        return await m.reply_text("Done!")
                    if to_dwl and not (m.photo or file.endswith(exe)):
                        os.remove(file)
            c_id = default[0]
            if m.photo or file.endswith(exe):
                if m.caption:
                    yus.append({"file":file, "caption": caption})
                elif not m.caption:
                    yus.append({"file":file, "caption": False})
                return await m.reply_text("Done!")
            elif m.document:
                if m.caption:
                    await psy.send_document(c_id, file, caption=caption)
                elif not m.caption:
                    await psy.send_document(c_id, file)
                return await m.reply_text("Done!")
            elif m.video:
                if m.caption:
                    await psy.send_video(c_id, file, caption=caption)
                elif not m.caption:
                    await psy.send_video(c_id, file)
            if to_dwl and not (m.photo or file.endswith(exe)):
                os.remove(file)
                return await m.reply_text("Done!")
        except Exception as e:
            os.remove(file)
            return await m.reply_text(f"Got an error:\n{e}")
    else:

        return

@psy.on_callback_query(filters.regex("^help_lao"),69)
async def help_lelo_guyz(_,q: CallbackQuery):
    data = q.data.split("_")
    typo = data[-1]
    if typo == "back":
        await q.edit_message_text(
            "What do you want to see",
            reply_markup=help_kb
        )
        return
    elif typo == "bc":
        await q.edit_message_text(
            helpmsg2,
            reply_markup=help_kb2
        )
        return
    elif typo == "re":
        if q.from_user.id not in SUDOER:
            await q.answer("Not for you give 500 stars on repo to convert it into free4all cmd",True)
            return
        await q.edit_message_text(
            helpmsg1,
            reply_markup=help_kb2
        )
        return

@psy.on_callback_query(group=0)
async def callbacks(_,q: CallbackQuery):
    try:
        data = q.data
        if data == "close":
            await q.message.delete()
            await q.answer("Closed")
            return
        elif data == "help":
            await q.edit_message_text(
                "What do you want to see",
                reply_markup=help_kb
            )
            await q.answer("Help")
            return
        elif data == "back":
            await q.edit_message_text(
                "Forward me any media or message I will forward to the targeted chat.\nUse /help to see the list of commands",
                reply_markup=start_kb
            )
            return await q.answer("Back")
    except Exception as e:
        c_id = q.message.chat.id
        await q.message.delete()
        return await psy.send_message(
            c_id,
            f"Got an error while handeling the callback quer:\n{e}")

async def size_fetcher(m: Message):
    """Return the file size in mb's and id"""
    if m.video:
        size = m.video.file_size
        id = m.video.file_id
    elif m.document:
        size = m.document.file_size
        id = m.document.file_id
    elif m.photo:
        size = m.photo.file_size
        id = m.photo.file_id
    x = (size/1024)/1024
    return x, id


from django.shortcuts import render
from .models import Image

# Create your views here.
def display_images(request):
    photos = Image.photo_display

    return render(request, 'all-photos/photo.html', { "photos": photos})
def search_results(request):
    '''
    search function to display search search_results
    args:
    order defines category
    '''
    if 'image' in request.GET and request.GET["image"]:
        search_term = request.GET.get("image")
        searched_images = Image.search_by_category(search_term)
        message = f"{search_term}"

        return render(request, 'all-photos/search.html', {"message": message, "images": searched_images})

    else:
        message = "You haven't searched for any term"
        return render(request, 'all-photos/search.html', {"message": message})
def image(request,image_id):
    try:
        image = Image.objects.get(id=image_id)
    except DoesNotExist:
        raise Http404()
    return render(request,"all-photos/item.html",{"image":image})
def filter_by_location(request,location_id):
   """
   Function that filters images by location
   """
   images = Image.filter_by_location(id=location_id )
   return render (request, 'location.html', {"images":images})


from multiprocessing import cpu_count

from PyQt6.QtWidgets import QDialog, QMainWindow
from PyQt6.QtCore import Qt

from ...logika.wielowatkowosc.analizator_zdjec import AnalizatorZdjec
from ..designer.ustawienia_zasobow_designer import Ui_Dialog

class UstawieniaZasobow(QDialog):
    '''Okno ustawień zużycia zasobów'''

    def __init__(self, glowne_okno: QMainWindow, analizator: AnalizatorZdjec) -> None:
        super().__init__()

        self.resize(800, 200)
        self.__glowne_okno = glowne_okno
        self.__analizator = analizator

        self.ui = Ui_Dialog()
        self.ui.setupUi(self)
        self.setWindowTitle("Ustawienia zasobów")

        self.ui.suwak.setMaximum(cpu_count())
        self.ui.suwak.valueChanged.connect(lambda: self.ui.wartosc.setText(str(self.ui.suwak.value())))
        self.ui.suwak.valueChanged.connect(lambda: self.__analizator.zmien_limit_procesow(self.ui.suwak.value()))

    def otworz(self):
        '''Otwiera okno z ustawieniami zużycia zasobów'''

        self.setStyleSheet(self.__glowne_okno.styleSheet())
        self.ui.suwak.setValue(self.__analizator.maks_liczba_procesow)

        self.setWindowModality(Qt.WindowModality.ApplicationModal)
        self.exec()


"""This module contains the classes for Gradescope assignments."""
from __future__ import annotations

import logging
import re
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, TypedDict

import requests
from bs4 import BeautifulSoup
from dataclasses_json import dataclass_json
from requests_toolbelt import MultipartEncoder

from gapper.connect.api.mixins import SessionHolder
from gapper.connect.api.utils import OSChoices

_assignment_logger = logging.getLogger("gapper.connect.api.assignment")

NO_IMAGE_REGEX = re.compile(r"gon\.image *= *null *;?")
IMAGE_REGEX = re.compile(
    r"gon\.image *= *{.*\"name\" *: *\"gradescope/autograders:.*-(?P<docker_id>\d+)\""
)

class DockerStatusJson(TypedDict):
    """The json response of the docker status of upload."""

    id: int
    assignment_id: int
    name: str
    status: str
    stdout: str
    stderr: str
    created_at: str
    updated_at: str

@dataclass_json
@dataclass
class GSAssignmentEssential(SessionHolder):
    """The essential information of a Gradescope assignment."""

    cid: str
    aid: str
    docker_id: str | None

    def __init__(
        self,
        cid: str,
        aid: str,
        docker_id: str | None = None,
        *,
        session: requests.Session | None = None,
    ) -> None:
        """Create a Gradescope assignment essential object.

        :param cid: The course id of the Gradescope assignment.
        :param aid: The assignment id of the Gradescope assignment.
        :param docker_id: The docker id of the assignment.
        :param session: The session to use.
        """
        super().__init__(session)
        self.cid = cid
        self.aid = aid
        self.docker_id = docker_id
        self._logger = _assignment_logger.getChild(f"GSAssignmentEssential_{self.aid}")

    async def upload_autograder(self, path: Path, os_choice: OSChoices) -> None:
        """Upload the autograder to the assignment.

        :param path: The path to the autograder zip file.
        :param os_choice: The choice of the operating system.
        """
        if not path.exists():
            raise FileNotFoundError(f"File {path} does not exist when uploading")
        if not path.is_file() or not path.suffix == ".zip":
            raise ValueError(f"File {path} is not a zip file")

        autograder_config = self._session.get(
            "https://www.gradescope.com/courses/"
            + self.cid
            + "/assignments/"
            + self.aid
            + "/configure_autograder"
        )
        autograder_config_resp = BeautifulSoup(autograder_config.text, "html.parser")
        autograder_form = autograder_config_resp.find(
            "form", attrs={"class": "js-autograderForm"}
        )
        authenticity_token = autograder_form.find(
            "input", attrs={"name": "authenticity_token"}
        ).get("value")

        autograder_dict: Dict[str, str] = {
            "utf8": "â",
            "_method": "patch",
            "authenticity_token": authenticity_token,
            "configuration": "zip",
            "autograder_zip": (path.name, path.read_bytes(), "text/plain"),
            "base_image_id": str(os_choice.value.id),
            "assignment[image_name]": os_choice.value.image,
        }

        multipart = MultipartEncoder(fields=autograder_dict)

        response = self._session.post(
            f"https://www.gradescope.com/courses/{self.cid}/assignments/{self.aid}/",
            data=multipart,
            headers={"Content-Type": multipart.content_type},
        )

        if response.status_code != requests.codes.ok:
            raise ValueError(f"Upload failed with status code {response.status_code}")

    def get_active_docker_id(self) -> str | None:
        """Get the active docker id of the assignment."""
        if self.docker_id is None:
            autograder_config = self._session.get(
                "https://www.gradescope.com/courses/"
                + self.cid
                + "/assignments/"
                + self.aid
                + "/configure_autograder"
            )

            if autograder_config.status_code != requests.codes.ok:
                self._logger.debug(
                    f"Failed to get autograder config. The status code is {autograder_config.status_code}"
                )
                return None

            self._logger.debug("Got autograder config")
            self._logger.debug(autograder_config.text)

            if NO_IMAGE_REGEX.search(autograder_config.text) is not None:
                self._logger.debug("No image found")
                return None

            id_match = IMAGE_REGEX.search(autograder_config.text)
            if id_match is None:
                self._logger.debug("No image found")
                self._logger.error(
                    "Cannot find docker image id. Please notify the developer of this error."
                )
                return None

            self.docker_id = id_match.groupdict()["docker_id"]

        return self.docker_id

    def get_docker_build_status(self) -> DockerStatusJson | None:
        """Get the docker build status of the assignment."""
        docker_id = self.get_active_docker_id()
        if docker_id is None:
            return None

        return self._session.get(
            f"https://www.gradescope.com/courses/{self.cid}/assignments/{self.aid}/docker_images/{docker_id}.json"
        ).json()

    def __eq__(self, other: Any) -> bool:
        """Check if two GSAssignments are equal."""
        if isinstance(other, GSAssignmentEssential):
            return self.cid == other.cid and self.aid == other.aid
        return False

    def __hash__(self) -> int:
        return hash(f"{self.cid}{self.aid}")

@dataclass_json
@dataclass
class GSAssignment(GSAssignmentEssential):
    """The Gradescope assignment with detailed information."""

    name: str
    points: str
    submissions: str
    percent_graded: str
    published: bool
    release_date: str
    due_date: str
    hard_due_date: str | None

    def __init__(
        self,
        cid: str,
        name: str,
        aid: str,
        points: str,
        submissions: str,
        percent_graded: str,
        published: bool,
        release_date: str,
        due_date: str,
        hard_due_date: str | None,
        docker_id: str | None = None,
        *,
        session: requests.Session | None = None,
    ) -> None:
        """Create a Gradescope assignment object.

        :param cid: The course id of the Gradescope assignment.
        :param name: The name of the Gradescope assignment.
        :param aid: The assignment id of the Gradescope assignment.
        :param points: The points of the Gradescope assignment.
        :param submissions: The submissions of the Gradescope assignment.
        :param percent_graded: The percent graded of the Gradescope assignment.
        :param published: Whether the Gradescope assignment is published.
        :param release_date: The release date of the Gradescope assignment.
        :param due_date: The due date of the Gradescope assignment.
        :param hard_due_date: The late due date of the Gradescope assignment.
        :param docker_id: The docker id of the Gradescope assignment.
        :param session: The session to use.
        """
        super().__init__(cid, aid, docker_id, session=session)
        self.name = name
        self.points = points
        self.submissions = submissions
        self.percent_graded = percent_graded
        self.published = published
        self.release_date = release_date
        self.due_date = due_date
        self.hard_due_date = hard_due_date
        self._logger = _assignment_logger.getChild(f"GSAssignment_{self.aid}")


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

def main():

    path="./image2.jpg"
    image=plt.imread(path)

    red_chanel=image[:,:,0]

    green_chanel=image[:,:,1]

    blue_chanel=image[:,:,2]

    height,width=image.shape[:2]
    grayscale=np.zeros((height,width),np.int8)

    grayscale=red_chanel*.2989+green_chanel*.5870+blue_chanel*.1140

    #binary convertion

    #height,width=grayscale.shape

    print(height)

    #grayscale=cv2.cvtColor(image,cv2.COLOR_RGB2GRAY)
    print(grayscale)

    binary=np.zeros((height,width),np.int8)

    for i in range(height):
        for j in range(width):

            if(grayscale[i][j]<=55):
                binary[i][j]=0

            else:
                binary[i][j]=255

    '''
    grayscale=cv2.cvtColor(image,cv2.COLOR_RGB2GRAY)
    th,binary=cv2.threshold(grayscale,127,255,cv2.THRESH_BINARY_INV)
    th,binary=cv2.threshold(grayscale,127,255,cv2.THRESH_BINARY)

    '''

    '''

    plt.subplot(4,2,1)
    plt.title("image")
    plt.imshow(image)

    plt.subplot(4,2,2)
    plt.title("grayscale")
    plt.imshow(grayscale,cmap="gray")

    plt.subplot(4,2,3)
    plt.title("binary")
    plt.imshow(binary,cmap="binary")

    plt.subplot(4,2,4)
    plt.title("red")
    plt.imshow(red_chanel,cmap="Reds")

    plt.subplot(4,2,5)
    plt.title("green")
    plt.imshow(green_chanel,cmap="Greens")

    plt.subplot(4,2,6)
    plt.title("red")
    plt.imshow(blue_chanel,cmap="Blues")

    plt.subplot(4,2,7)
    plt.title("histogram")
    plt.hist(binary.ravel(),255,[0,255])

    '''

    title_set=["image","grayscale","binary","red_chanel","green_chanel","blue_chanel"]
    img_set=[image,grayscale,binary,red_chanel,green_chanel,blue_chanel]

    def plot(title_set,img_set):

        n=len(img_set)

        for i in range(n):
            img=img_set[i]
            plt.subplot(4,2,i+1)
            ch=len(img.shape)

            if(ch==2):
                plt.imshow(img,cmap="gray")
            else:
                plt.imshow(img)

        plt.show()

    plot(title_set,img_set)

if __name__ == "__main__":
    main()


import wave
import threading
from time import sleep
import pyaudio
import numpy as np
#import eyeD3
import sys

CHUNK = 1024

if len(sys.argv) < 2:
    print("Plays a wave file.\n\nUsage: %s filename.wav" % sys.argv[0])
    sys.exit(-1)

wf = wave.open(sys.argv[1], 'rb')

p = pyaudio.PyAudio()

stream = p.open(format=p.get_format_from_width(wf.getsampwidth()),
                channels=wf.getnchannels(),
                rate=wf.getframerate(),
                output=True)

data = wf.readframes(CHUNK)

while data != '':
    #stream.write(data)
    data = wf.readframes(CHUNK)
    eq_data = np.fromstring(data,dtype=np.int16)
    peak=np.average(np.abs(eq_data))*2
    bars="#"*int(50*peak/2**16)
    print("%05d %s"%(peak,bars))

stream.stop_stream()
stream.close()

p.terminate()


from mmseg.datasets.pipelines.transforms import *
from mmseg.datasets.builder import PIPELINES
from scipy.ndimage.morphology import distance_transform_edt
import cv2
import numpy as np
from collections import OrderedDict

"""
This file contains the implementation of various data augmentation pipelines for the DoubleImageDataset class.
These pipelines include:
- DoubleImageResize: Resizes the input images to a given size.
- DoubleImageRandomFlip: Randomly flips the input images horizontally or vertically.
- DoubleImagePad: Pads the input images to a given size.
- DoubleImageNormalize: Normalizes the input images.
- DoubleImageRandomCrop: Randomly crops the input images.
- DoubleImagePhotoMetricDistortion: Applies photometric distortion to the input images.
- DoubleImageLoadAsBinaryLabel: Converts the input segmentation maps to binary labels.
- BinaryLabelDilate: Dilates the binary labels.
- OmmitSmallRegion: Omits small regions from the binary labels.
- DoubleImageRandomRgbBgr: Randomly swaps the color channels of the input images.
...

"""

@PIPELINES.register_module()
class DoubleImageResize(Resize):
    """
    Resize the input images to a given size.
    Required keys are "img1" and "img2", added or modified keys are "img1", "img2", "img_shape", "pad_shape", "scale_factor", "keep_ratio", "scale", and "scale_idx".
    Args:
        img_scale (tuple[int]): (w, h) of resized image. Default: None.
        multiscale_mode (str): Either "range" or "value". Default: "range".
        ratio_range (tuple[float]): (min_ratio, max_ratio) of the resized image. Default: None.
        keep_ratio (bool): Whether to keep the aspect ratio when resizing the image. Default: True.
    """

    def __init__(self, **kwargs):
        super(DoubleImageResize, self).__init__(**kwargs)

    def _random_scale(self, results):
        """
        Randomly sample a scale from the given scale range or list.
        Args:
            results (dict): A result dict contains the img1 and img2.
        """
        if self.ratio_range is not None:
            if self.img_scale is None:
                h, w = results["img1"].shape[:2]
                scale, scale_idx = self.random_sample_ratio((w, h), self.ratio_range)
            else:
                scale, scale_idx = self.random_sample_ratio(
                    self.img_scale[0], self.ratio_range
                )
        elif len(self.img_scale) == 1:
            scale, scale_idx = self.img_scale[0], 0
        elif self.multiscale_mode == "range":
            scale, scale_idx = self.random_sample(self.img_scale)
        elif self.multiscale_mode == "value":
            scale, scale_idx = self.random_select(self.img_scale)
        else:
            raise NotImplementedError

        results["scale"] = scale
        results["scale_idx"] = scale_idx

    def _resize_img(self, results):
        """
        Resize images with ``results['scale']``.
        Args:
            results (dict): A result dict contains the img1 and img2.
        """
        if self.keep_ratio:
            img1, scale_factor = mmcv.imrescale(
                results["img1"], results["scale"], return_scale=True
            )
            img2, scale_factor = mmcv.imrescale(
                results["img2"], results["scale"], return_scale=True
            )
            # the w_scale and h_scale has minor difference
            # a real fix should be done in the mmcv.imrescale in the future
            new_h, new_w = img1.shape[:2]
            h, w = results["img1"].shape[:2]
            w_scale = new_w / w
            h_scale = new_h / h
        else:
            img1, w_scale, h_scale = mmcv.imresize(
                results["img1"], results["scale"], return_scale=True
            )
            img2, w_scale, h_scale = mmcv.imresize(
                results["img2"], results["scale"], return_scale=True
            )
        scale_factor = np.array([w_scale, h_scale, w_scale, h_scale], dtype=np.float32)
        results["img1"] = img1
        results["img2"] = img2
        results["img_shape"] = img1.shape
        results["pad_shape"] = img1.shape  # in case that there is no padding
        results["scale_factor"] = scale_factor
        results["keep_ratio"] = self.keep_ratio

@PIPELINES.register_module()
class DoubleImageRandomFlip(RandomFlip):
    """
    Randomly flip the input images horizontally or vertically.
    Required keys are "img1" and "img2", added or modified keys are "img1", "img2", "flip", and "flip_direction".
    Args:
        prob (float): Probability of flipping. Default: 0.5.
        direction (str): The flipping direction. Options are "horizontal" and "vertical". Default: "horizontal".
    """

    def __init__(self, **kwargs):
        super(DoubleImageRandomFlip, self).__init__(**kwargs)

    def __call__(self, results):
        """
        Call function to flip the input images.
        Args:
            results (dict): A result dict contains the img1 and img2.
        Returns:
            results (dict): A result dict contains the img1, img2, flip, and flip_direction.
        """
        if "flip" not in results:
            flip = True if np.random.rand() < self.prob else False
            results["flip"] = flip
        if "flip_direction" not in results:
            results["flip_direction"] = self.direction
        if results["flip"]:
            # flip image
            results["img1"] = mmcv.imflip(
                results["img1"], direction=results["flip_direction"]
            )
            results["img2"] = mmcv.imflip(
                results["img2"], direction=results["flip_direction"]
            )

            # flip segs
            for key in results.get("seg_fields", []):
                # use copy() to make numpy stride positive
                results[key] = mmcv.imflip(
                    results[key], direction=results["flip_direction"]
                ).copy()
        return results

@PIPELINES.register_module()
class DoubleImagePad(Pad):
    """
    Pad the input images to a given size.
    Required keys are "img1" and "img2", added or modified keys are "img1", "img2", "pad_shape", "pad_fixed_size", and "pad_size_divisor".
    Args:
        size (tuple[int]): Fixed size of padding. Default: None.
        size_divisor (int): The divisor of padded size. Default: None.
        pad_val (float | int | tuple[int]): Padding value. Default: 0.
    """

    def __init__(self, **kwargs):
        super(DoubleImagePad, self).__init__(**kwargs)

    def _pad_img(self, results):
        """
        Pad images according to ``self.size`` or ``self.size_divisor``.
        Args:
            results (dict): A result dict contains the img1 and img2.
        """
        if self.size is not None:
            if (
                results["img1"].shape[0] < self.size[0]
                and results["img1"].shape[1] < self.size[1]
            ):
                padded_img1 = mmcv.impad(
                    results["img1"], shape=self.size, pad_val=self.pad_val
                )
                padded_img2 = mmcv.impad(
                    results["img2"], shape=self.size, pad_val=self.pad_val
                )
            else:
                padded_img1 = results["img1"]
                padded_img2 = results["img2"]
        elif self.size_divisor is not None:
            padded_img1 = mmcv.impad_to_multiple(
                results["img1"], self.size_divisor, pad_val=self.pad_val
            )
            padded_img2 = mmcv.impad_to_multiple(
                results["img2"], self.size_divisor, pad_val=self.pad_val
            )
        results["img1"] = padded_img1
        results["img2"] = padded_img2
        results["pad_shape"] = padded_img1.shape
        results["pad_fixed_size"] = self.size
        results["pad_size_divisor"] = self.size_divisor

@PIPELINES.register_module()
class DoubleImageNormalize(Normalize):
    """
    Normalize the input images.
    Required keys are "img1" and "img2", added or modified keys are "img1", "img2", and "img_norm_cfg".
    Args:
        mean (tuple[float]): Mean values of 3 channels.
        std (tuple[float]): Std values of 3 channels.
        to_rgb (bool): Whether to convert the image from BGR to RGB. Default: True.
    """

    def __init__(self, split=None, **kwargs):
        super(DoubleImageNormalize, self).__init__(**kwargs)
        self.split = split

    def __call__(self, results):
        """
        Call function to normalize the input images.
        Args:
            results (dict): A result dict contains the img1 and img2.
        Returns:
            results (dict): A result dict contains the img1, img2, and img_norm_cfg.
        """
        if self.split is not None:
            img1 = results["img1"]
            assert img1.shape[2] == len(self.split)
            img1 = img1[:, :, self.split]
            results["img1"] = img1

            img2 = results["img2"]
            assert img2.shape[2] == len(self.split)
            img2 = img2[:, :, self.split]
            results["img"] = img2

        results["img1"] = mmcv.imnormalize(
            results["img1"], self.mean, self.std, self.to_rgb
        )
        results["img2"] = mmcv.imnormalize(
            results["img2"], self.mean, self.std, self.to_rgb
        )
        results["img_norm_cfg"] = dict(mean=self.mean, std=self.std, to_rgb=self.to_rgb)
        return results

@PIPELINES.register_module()
class DoubleImageRandomCrop(RandomCrop):
    """
    Crop the input images and semantic segmentation maps with random size and aspect ratio.
    Required keys are "img1" and "img2", added or modified keys are "img1", "img2", "img_shape", and "seg_fields".
    Args:
        crop_size (tuple[int]): Expected size after cropping, (h, w).
        cat_max_ratio (float): The maximum ratio between the cropped region and the visible region of a category. Default: 1.0.
        ignore_index (int): The label index to be ignored in semantic segmentation maps. Default: 255.
    """

    def __init__(self, **kwargs):
        super(DoubleImageRandomCrop, self).__init__(**kwargs)

    def __call__(self, results):
        """
        Call function to crop the input images and semantic segmentation maps.
        Args:
            results (dict): A result dict contains the img1 and img2.
        Returns:
            results (dict): A result dict contains the img1, img2, img_shape, and seg_fields.
        """
        img1 = results["img1"]
        img2 = results["img2"]
        crop_bbox = self.get_crop_bbox(img1)
        if self.cat_max_ratio < 1.0:
            # Repeat 10 times
            for _ in range(10):
                seg_temp = self.crop(results["gt_semantic_seg"], crop_bbox)
                labels, cnt = np.unique(seg_temp, return_counts=True)
                cnt = cnt[labels != self.ignore_index]
                if len(cnt) > 1 and np.max(cnt) / np.sum(cnt) < self.cat_max_ratio:
                    break
                crop_bbox = self.get_crop_bbox(img1)

        # crop the image
        img1 = self.crop(img1, crop_bbox)
        img2 = self.crop(img2, crop_bbox)
        img_shape = img1.shape
        results["img1"] = img1
        results["img2"] = img2
        results["img_shape"] = img_shape

        # crop semantic seg
        for key in results.get("seg_fields", []):
            results[key] = self.crop(results[key], crop_bbox)

        return results

@PIPELINES.register_module()
class DoubleImagePhotoMetricDistortion(PhotoMetricDistortion):
    """
    Distort the input images' photo metric.
    Required keys are "img1" and "img2".
    Args:
        brightness_delta (int): Delta of brightness.
        contrast_range (tuple[float]): Range of contrast.
        saturation_range (tuple[float]): Range of saturation.
        hue_delta (int): Delta of hue.
    """

    def __init__(self, **kwargs):
        super(DoubleImagePhotoMetricDistortion, self).__init__(**kwargs)

    def __call__(self, results):
        """
        Call function to distort the input images' photo metric.
        Args:
            results (dict): A result dict contains the img1 and img2.
        Returns:
            results (dict): A result dict contains the img1 and img2.
        """
        for key in ["img1", "img2"]:
            img = results[key]
            # random brightness
            img = self.brightness(img)

            # mode == 0 --> do random contrast first
            # mode == 1 --> do random contrast last
            mode = random.randint(2)
            if mode == 1:
                img = self.contrast(img)

            # random saturation
            img = self.saturation(img)

            # random hue
            img = self.hue(img)

            # random contrast
            if mode == 0:
                img = self.contrast(img)

            results[key] = img
        return results

@PIPELINES.register_module()
class DoubleImageLoadAsBinaryLabel(object):
    """
    Load the input images as binary label.
    Required keys are "gt_semantic_seg", added or modified keys are "gt_semantic_seg".
    Args:
        ignore_labels (list[int]): The label index to be ignored in semantic segmentation maps. Default: [].
        ignore_ind (int): The label index to be ignored in binary label. Default: 255.
    """

    def __init__(self, ignore_labels=[], ignore_ind=255):
        if not isinstance(ignore_labels, list):
            ignore_labels = [ignore_labels]
        self.ignore_labels = ignore_labels
        self.ignore_ind = ignore_ind

    def __call__(self, results):
        """
        Call function to load the input images as binary label.
        Args:
            results (dict): A result dict contains the gt_semantic_seg.
        Returns:
            results (dict): A result dict contains the gt_semantic_seg.
        """
        seg_map = results["gt_semantic_seg"]
        for ignore_label in self.ignore_labels:
            seg_map[seg_map == ignore_label] = 0
        seg_map_res = np.zeros_like(seg_map)
        seg_map_res[seg_map > 0] = 1
        seg_map_res[seg_map == self.ignore_ind] = 255
        results["gt_semantic_seg"] = seg_map_res
        return results

@PIPELINES.register_module()
class BinaryLabelDilate(object):
    """
    Dilate the binary label.
    Required keys are "gt_semantic_seg", added or modified keys are "gt_semantic_seg".
    Args:
        dilate_pixel (int): The pixel to dilate the binary label. Default: 3.
    """

    def __init__(self, dilate_pixel=3):
        self.dilate_pixel = dilate_pixel
        kernel_size = int(dilate_pixel * 2 + 1)
        self.conv_kernel = cv2.getStructuringElement(
            cv2.MORPH_RECT, (kernel_size, kernel_size)
        )

    def __call__(self, results):
        """
        Call function to dilate the binary label.
        Args:
            results (dict): A result dict contains the gt_semantic_seg.
        Returns:
            results (dict): A result dict contains the gt_semantic_seg.
        """
        seg_map = results["gt_semantic_seg"]
        seg_map = cv2.dilate(seg_map, self.conv_kernel)
        results["gt_semantic_seg"] = seg_map
        return results

@PIPELINES.register_module()
class OmmitSmallRegion(object):
    """
    Ommit the small region in the input image.
    Required keys are "gt_semantic_seg", added or modified keys are "gt_semantic_seg".
    Args:
        min_area (int): The minimum area of the region to be kept. Default: 50.
    """

    def __init__(self, min_area=50):
        self.min_area = min_area

    def __call__(self, results):
        """
        Call function to ommit the small region in the input image.
        Args:
            results (dict): A result dict contains the gt_semantic_seg.
        Returns:
            results (dict): A result dict contains the gt_semantic_seg.
        """
        seg_map = results["gt_semantic_seg"]
        seg_map_binary = seg_map.copy()
        seg_map_binary[seg_map_binary > 0] = 1
        contours, _ = cv2.findContours(
            seg_map_binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
        )
        for contour in contours:
            area = cv2.contourArea(contour)
            if area < self.min_area:
                cv2.drawContours(seg_map, [contour], -1, (0, 0, 0), -1)
        results["gt_semantic_seg"] = seg_map
        return results

@PIPELINES.register_module()
class DoubleImageRandomRgbBgr(object):
    """
    Randomly flip the RGB and BGR channels of the input images.
    Required keys are "img1" and "img2", added or modified keys are "img1" and "img2".
    Args:
        prob (float): Probability of flipping the channels. Default: 0.3.
    """

    def __init__(self, prob=0.3):
        """
        Constructor of the class.
        Args:
            prob (float): Probability of flipping the channels. Default: 0.3.
        """
        self.prob = prob

    def __call__(self, results):
        """
        Call function to randomly flip the RGB and BGR channels of the input images.
        Args:
            results (dict): A result dict contains the img1 and img2.
        Returns:
            results (dict): A result dict contains the img1 and img2.
        """
        if np.random.rand() < self.prob:
            results["img1"] = results["img1"][:, :, ::-1]
        if np.random.rand() < self.prob:
            results["img2"] = results["img2"][:, :, ::-1]
        return results

@PIPELINES.register_module()
class LoadAsMultiLabel(object):
    """
    Load the semantic segmentation label as multi-label.
    Required keys are "gt_semantic_seg", added or modified keys are "gt_semantic_seg_sat1" and "gt_semantic_seg_sat2".
    Args:
        ignore_bg (bool): Whether to ignore the background. Default: True.
    """

    def __init__(self, ignore_bg=True):
        CLASSES = [
            "绿地-动土",
            "动土-绿地",
            "硬化空地-动土",
            "动土-硬化空地",
            "绿地-硬化空地",
            "硬化空地-绿地",
            "动土-在建建筑",
            "绿地-在建建筑",
            "水域-在建建筑",
            "硬化空地-在建建筑",
            "绿地-建筑",
            "动土-建筑",
            "水域-建筑",
            "硬化空地-建筑",
            "建筑-建筑",
            "在建建筑-建筑",
            "在建建筑-绿地",
            "在建建筑-动土",
            "在建建筑-水域",
            "在建建筑-硬化空地",
            "建筑-绿地",
            "建筑-动土",
            "建筑-水域",
            "建筑-硬化空地",
            "建筑-在建建筑",
            "绿地-大棚",
            "动土-大棚",
            "水域-大棚",
            "硬化空地-大棚",
            "建筑-大棚",
            "在建建筑-大棚",
            "大棚-绿地",
            "大棚-水域",
            "大棚-动土",
            "大棚-在建建筑",
            "大棚-建筑",
            "大棚-硬化空地",
            "水域-动土",
            "水域-硬化空地",
            "动土-水域",
            "硬化空地-水域",
            "绿地-运动场",
            "水域-运动场",
            "动土-运动场",
            "硬化空地-运动场",
            "大棚-运动场",
            "建筑-运动场",
            "在建建筑-运动场",
            "运动场-绿地",
            "运动场-动土",
            "运动场-水域",
            "运动场-硬化空地",
            "运动场-建筑",
            "运动场-大棚",
            "运动场-在建建筑",
            "自然裸地-动土",
            "动土-自然裸地",
            "建筑-自然裸地",
            "自然裸地-建筑",
            "大棚-自然裸地",
            "自然裸地-大棚",
            "林地-在建建筑",
            "林地-建筑",
            "在建建筑-林地",
            "在建大棚-林地",
            "林地-动土",
            "林地-水域",
            "运动场-林地",
            "动土-林地",
            "地膜-林地",
            "建筑-林地",
            "林地-无作物耕地",
            "林地-在建大棚",
            "有作物耕地-林地",
            "林地-大棚",
            "林地-地膜",
            "硬化空地-林地",
            "无作物耕地-林地",
            "林地-有作物耕地",
            "水域-林地",
            "林地-运动场",
            "自然裸地-林地",
            "林地-硬化空地",
            "林地-自然裸地",
            "大棚-林地",
            "有作物耕地-动土",
            "无作物耕地-在建建筑",
            "无作物耕地-建筑",
            "在建建筑-无作物耕地",
            "有作物耕地-建筑",
            "在建建筑-有作物耕地",
            "硬化空地-无作物耕地",
            "有作物耕地-在建建筑",
            "无作物耕地-动土",
            "地膜-无作物耕地",
            "在建大棚-无作物耕地",
            "建筑-无作物耕地",
            "运动场-无作物耕地",
            "动土-有作物耕地",
            "无作物耕地-在建大棚",
            "有作物耕地-地膜",
            "大棚-无作物耕地",
            "大棚-有作物耕地",
            "硬化空地-有作物耕地",
            "动土-无作物耕地",
            "无作物耕地-大棚",
            "有作物耕地-大棚",
            "有作物耕地-在建大棚",
            "无作物耕地-地膜",
            "有作物耕地-硬化空地",
            "有作物耕地-运动场",
            "在建大棚-有作物耕地",
            "无作物耕地-运动场",
            "地膜-有作物耕地",
            "建筑-有作物耕地",
            "无作物耕地-硬化空地",
            "运动场-有作物耕地",
        ]
        categories = set()
        for cats in CLASSES:
            categories.update(set(cats.split("-")))
        categories = list(categories)
        categories = sorted(categories)
        label_mapping = [[255, 255]]
        for cats in CLASSES:
            cat1, cat2 = cats.split("-")
            label_mapping.append([categories.index(cat1), categories.index(cat2)])
        self.label_mapping = np.array(label_mapping).astype(np.uint8)
        self.ignore_bg = ignore_bg

    def __call__(self, results):
        seg_map = results["gt_semantic_seg"]
        multi_cls_seg_map = self.label_mapping[seg_map]
        gt_semantic_seg_sat1 = multi_cls_seg_map[:, :, 0]
        gt_semantic_seg_sat2 = multi_cls_seg_map[:, :, 1]
        if not self.ignore_bg:
            gt_semantic_seg_sat1 += 1
            gt_semantic_seg_sat2 += 1
        if "seg_fields" in results:
            results["seg_fields"].extend(
                ["gt_semantic_seg_sat1", "gt_semantic_seg_sat2"]
            )
        results["gt_semantic_seg_sat1"] = gt_semantic_seg_sat1
        results["gt_semantic_seg_sat2"] = gt_semantic_seg_sat2
        return results

@PIPELINES.register_module()
class LoadAs19ClassLabel(object):
    """
    LoadAs19ClassLabel pipeline that maps the original 85-class semantic segmentation labels to 19 classes.
    Args:
        ignore_bg (bool): Whether to ignore the background class (class 0) when mapping labels. Default: False.
        raw_classes (list[str]): List of original 85 classes. If not provided, the default list of classes will be used.
        index2change (dict): Dictionary that maps the index of a class in the original 85-class list to the index of the
            corresponding class in the 19-class list. If not provided, the default mapping will be used.
    """

    def __init__(self, ignore_bg=False, raw_classes=None, index2change=None):
        CLASSES = raw_classes
        INDEX2CHANGE = index2change
        if CLASSES is None:
            CLASSES = [
                "绿地-动土",
                "动土-绿地",
                "硬化空地-动土",
                "动土-硬化空地",
                "绿地-硬化空地",
                "硬化空地-绿地",
                "动土-在建建筑",
                "绿地-在建建筑",
                "水域-在建建筑",
                "硬化空地-在建建筑",
                "绿地-建筑",
                "动土-建筑",
                "水域-建筑",
                "硬化空地-建筑",
                "建筑-建筑",
                "在建建筑-建筑",
                "在建建筑-绿地",
                "在建建筑-动土",
                "在建建筑-水域",
                "在建建筑-硬化空地",
                "建筑-绿地",
                "建筑-动土",
                "建筑-水域",
                "建筑-硬化空地",
                "建筑-在建建筑",
                "绿地-大棚",
                "动土-大棚",
                "水域-大棚",
                "硬化空地-大棚",
                "建筑-大棚",
                "在建建筑-大棚",
                "大棚-绿地",
                "大棚-水域",
                "大棚-动土",
                "大棚-在建建筑",
                "大棚-建筑",
                "大棚-硬化空地",
                "水域-动土",
                "水域-硬化空地",
                "动土-水域",
                "硬化空地-水域",
                "绿地-运动场",
                "水域-运动场",
                "动土-运动场",
                "硬化空地-运动场",
                "大棚-运动场",
                "建筑-运动场",
                "在建建筑-运动场",
                "运动场-绿地",
                "运动场-动土",
                "运动场-水域",
                "运动场-硬化空地",
                "运动场-建筑",
                "运动场-大棚",
                "运动场-在建建筑",
                "自然裸地-动土",
                "动土-自然裸地",
                "建筑-自然裸地",
                "自然裸地-建筑",
                "大棚-自然裸地",
                "自然裸地-大棚",
                "林地-在建建筑",
                "林地-建筑",
                "在建建筑-林地",
                "在建大棚-林地",
                "林地-动土",
                "林地-水域",
                "运动场-林地",
                "动土-林地",
                "地膜-林地",
                "建筑-林地",
                "林地-无作物耕地",
                "林地-在建大棚",
                "有作物耕地-林地",
                "林地-大棚",
                "林地-地膜",
                "硬化空地-林地",
                "无作物耕地-林地",
                "林地-有作物耕地",
                "水域-林地",
                "林地-运动场",
                "自然裸地-林地",
                "林地-硬化空地",
                "林地-自然裸地",
                "大棚-林地",
                "有作物耕地-动土",
                "无作物耕地-在建建筑",
                "无作物耕地-建筑",
                "在建建筑-无作物耕地",
                "有作物耕地-建筑",
                "在建建筑-有作物耕地",
                "硬化空地-无作物耕地",
                "有作物耕地-在建建筑",
                "无作物耕地-动土",
                "地膜-无作物耕地",
                "在建大棚-无作物耕地",
                "建筑-无作物耕地",
                "运动场-无作物耕地",
                "动土-有作物耕地",
                "无作物耕地-在建大棚",
                "有作物耕地-地膜",
                "大棚-无作物耕地",
                "大棚-有作物耕地",
                "硬化空地-有作物耕地",
                "动土-无作物耕地",
                "无作物耕地-大棚",
                "有作物耕地-大棚",
                "有作物耕地-在建大棚",
                "无作物耕地-地膜",
                "有作物耕地-硬化空地",
                "有作物耕地-运动场",
                "在建大棚-有作物耕地",
                "无作物耕地-运动场",
                "地膜-有作物耕地",
                "建筑-有作物耕地",
                "无作物耕地-硬化空地",
                "运动场-有作物耕地",
            ]
        if INDEX2CHANGE is None:
            INDEX2CHANGE = {
                0: [
                    "自然裸地-地膜",
                    "自然裸地-有作物耕地",
                    "无作物耕地-有作物耕地",
                    "地膜-建筑",
                    "绿地-地膜",
                    "自然裸地-无作物耕地",
                    "地膜-硬化空地",
                    "水域-地膜",
                    "绿地-有作物耕地",
                    "有作物耕地-地膜",
                    "在建大棚-大棚",
                    "有作物耕地-绿地",
                    "林地-地膜",
                    "自然裸地-绿地",
                    "无作物耕地-地膜",
                    "有作物耕地-无作物耕地",
                    "绿地-无作物耕地",
                    "无作物耕地-自然裸地",
                    "在建建筑-地膜",
                    "地膜-动土",
                    "大棚-地膜",
                    "林地-无作物耕地",
                    "地膜-自然裸地",
                    "无作物耕地-林地",
                    "无作物耕地-绿地",
                    "地膜-绿地",
                    "地膜-无作物耕地",
                    "地膜-在建建筑",
                    "有作物耕地-自然裸地",
                    "自然裸地-林地",
                    "硬化空地-地膜",
                    "绿地-自然裸地",
                    "大棚-在建大棚",
                    "地膜-有作物耕地",
                    "林地-有作物耕地",
                    "林地-自然裸地",
                    "绿地-林地",
                    "地膜-林地",
                    "有作物耕地-林地",
                    "建筑-地膜",
                    "动土-地膜",
                    "运动场-地膜",
                    "在建大棚-地膜",
                    "地膜-运动场",
                    "地膜-水域",
                    "地膜-在建大棚",
                    "地膜-大棚",
                ],
                1: [
                    "动土-在建建筑",
                    "动土-建筑",
                    "运动场-在建建筑",
                    "运动场-建筑",
                    "自然裸地-建筑",
                    "硬化空地-在建建筑",
                    "硬化空地-建筑",
                    "自然裸地-在建建筑",
                    "无作物耕地-建筑",
                    "无作物耕地-在建建筑",
                ],
                2: [
                    "在建建筑-动土",
                    "建筑-动土",
                    "在建建筑-运动场",
                    "建筑-运动场",
                    "建筑-自然裸地",
                    "在建建筑-硬化空地",
                    "建筑-硬化空地",
                    "建筑-无作物耕地",
                    "在建建筑-自然裸地",
                ],
                3: ["绿地-在建建筑", "绿地-建筑", "林地-在建建筑", "有作物耕地-在建建筑", "林地-建筑", "有作物耕地-建筑"],
                4: [
                    "在建建筑-绿地",
                    "建筑-绿地",
                    "在建建筑-林地",
                    "建筑-林地",
                    "建筑-有作物耕地",
                    "在建建筑-有作物耕地",
                    "在建建筑-无作物耕地",
                ],
                5: ["动土-大棚", "硬化空地-大棚", "自然裸地-大棚", "在建建筑-大棚", "建筑-大棚", "水域-大棚"],
                6: ["大棚-动土", "大棚-硬化空地", "大棚-自然裸地", "大棚-在建建筑", "大棚-建筑", "大棚-水域"],
                7: [
                    "绿地-大棚",
                    "运动场-大棚",
                    "无作物耕地-大棚",
                    "在建建筑-在建大棚",
                    "水域-在建大棚",
                    "林地-在建大棚",
                    "建筑-在建大棚",
                    "运动场-在建大棚",
                    "动土-在建大棚",
                    "林地-大棚",
                    "绿地-在建大棚",
                    "自然裸地-在建大棚",
                    "无作物耕地-在建大棚",
                    "有作物耕地-在建大棚",
                    "硬化空地-在建大棚",
                    "有作物耕地-大棚",
                ],
                8: [
                    "大棚-绿地",
                    "大棚-运动场",
                    "在建大棚-绿地",
                    "在建大棚-自然裸地",
                    "大棚-有作物耕地",
                    "在建大棚-有作物耕地",
                    "在建大棚-动土",
                    "大棚-林地",
                    "在建大棚-林地",
                    "在建大棚-硬化空地",
                    "大棚-无作物耕地",
                    "在建大棚-无作物耕地",
                    "在建大棚-建筑",
                    "在建大棚-运动场",
                    "在建大棚-在建建筑",
                    "在建大棚-水域",
                ],
                9: ["动土-水域", "自然裸地-水域", "硬化空地-水域", "运动场-水域", "无作物耕地-水域"],
                10: ["水域-动土", "水域-自然裸地", "水域-硬化空地", "水域-运动场", "水域-无作物耕地"],
                11: ["动土-硬化空地", "动土-自然裸地", "动土-无作物耕地"],
                12: ["硬化空地-动土", "自然裸地-动土", "无作物耕地-动土"],
                13: [
                    "动土-绿地",
                    "硬化空地-绿地",
                    "动土-运动场",
                    "硬化空地-运动场",
                    "运动场-绿地",
                    "硬化空地-无作物耕地",
                    "硬化空地-林地",
                    "动土-林地",
                    "自然裸地-运动场",
                    "硬化空地-有作物耕地",
                    "运动场-有作物耕地",
                    "动土-有作物耕地",
                    "运动场-林地",
                    "硬化空地-自然裸地",
                ],
                14: [
                    "绿地-动土",
                    "绿地-硬化空地",
                    "运动场-动土",
                    "运动场-硬化空地",
                    "绿地-运动场",
                    "有作物耕地-动土",
                    "自然裸地-硬化空地",
                    "运动场-无作物耕地",
                    "无作物耕地-运动场",
                    "林地-硬化空地",
                    "有作物耕地-硬化空地",
                    "运动场-自然裸地",
                    "林地-运动场",
                    "无作物耕地-硬化空地",
                    "林地-动土",
                    "有作物耕地-运动场",
                ],
                15: ["在建建筑-水域", "建筑-水域"],
                16: ["水域-在建建筑", "水域-建筑"],
                17: ["水域-绿地", "水域-有作物耕地", "水域-林地"],
                18: ["绿地-水域", "有作物耕地-水域", "林地-水域"],
                19: ["在建建筑-建筑", "建筑-在建建筑", "建筑-建筑", "在建建筑-在建建筑"],
            }
        if ignore_bg:
            label_mapping = [255]
            for cat in CLASSES:
                for key, value in INDEX2CHANGE.items():
                    if cat in value:
                        label_mapping.append(np.uint8(key - 1))
                        break
        else:
            label_mapping = [0]
            for cat in CLASSES:
                for key, value in INDEX2CHANGE.items():
                    if cat in value:
                        label_mapping.append(key)
                        break

        assert len(label_mapping) == len(CLASSES) + 1
        self.label_mapping = np.array(label_mapping).astype(np.uint8)

    def __call__(self, results):
        seg_map = results["gt_semantic_seg"]
        seg_map_multi = self.label_mapping[seg_map]
        if "seg_fields" in results:
            results["seg_fields"].extend(["gt_semantic_seg_multi"])
        results["gt_semantic_seg_multi"] = seg_map_multi
        return results

@PIPELINES.register_module()
class LoadAsBinaryEdges(object):
    """
    LoadAsBinaryEdges: a pipeline to load binary edges from semantic segmentation.
    """

    def __init__(self, radius=2, dilate_pixel=2, outside=False):
        """
        Args:
            radius (int): radius of the edges.
            dilate_pixel (int): pixel to dilate.
            outside (bool): whether to load edges outside the object.
        """
        self.radius = radius
        self.dilate_pixel = dilate_pixel
        self.outside = outside
        kernel_size = int(dilate_pixel * 2 + 1)
        self.conv_kernel = cv2.getStructuringElement(
            cv2.MORPH_RECT, (kernel_size, kernel_size)
        )

    def __call__(self, results):
        """
        Load binary edges from semantic segmentation.

        Args:
            results (dict): Result dict from loading pipeline.

        Returns:
            dict: The result dict with binary edges.
        """
        seg_map = results["gt_semantic_seg"]
        edges = np.zeros_like(seg_map, dtype=np.uint8)
        num_class = seg_map.max()
        for i_cls in range(1, num_class + 1):
            mask = (seg_map == i_cls).astype(np.uint8)
            if mask.sum() == 0:
                continue
            dist = distance_transform_edt(1 - mask)
            if self.outside:
                dist[dist == 0] = 99999
            else:
                dist += distance_transform_edt(mask)
            # dist[dist <= self.radius] = 1
            # dist[dist > self.radius] = 0
            # edges += dist.astype(np.uint8)
            edges[dist <= self.radius] += 1
        edges[edges > 0] = 1
        # _edges = edges.copy()
        # _edges[seg_map == 0] = 0
        # _edges = cv2.dilate(_edges, self.conv_kernel)
        # edges[seg_map == 0] = 255
        # edges[_edges == 1] = 1

        results["gt_edges"] = edges
        return results

@PIPELINES.register_module()
class RandomReverse(object):
    """
    Randomly reverse the order of two images and their corresponding semantic segmentation maps.

    Args:
        prob (float): Probability of reversing the order.

    Returns:
        dict: The result dict with reversed images and semantic segmentation maps.
    """

    def __init__(self, prob=0.5):
        self.prob = prob

    def __call__(self, results):
        if np.random.rand() < self.prob:
            if "gt_semantic_seg_sat1" in results:
                results["gt_semantic_seg_sat1"], results["gt_semantic_seg_sat2"] = (
                    results["gt_semantic_seg_sat2"],
                    results["gt_semantic_seg_sat1"],
                )
            results["img1"], results["img2"] = results["img2"], results["img1"]
        return results

@PIPELINES.register_module()
class RandomResizeDownUp(object):
    """
    Randomly resize the input images with a probability of 0.6.
    If the image is resized, it will be either resized to 0.5x or 4x of its original size.
    Args:
        None
    Returns:
        dict: The result dict with resized images.
    """

    def __init__(self):
        pass

    def __call__(self, results):
        results["img1"] = self._random_resize(results["img1"])
        results["img2"] = self._random_resize(results["img2"])
        return results

    def _random_resize(self, img):
        h, w, c = img.shape
        if np.random.rand() < 0.6:
            r1 = np.random.rand()
            if r1 < 0.5:
                s1 = mmcv.imrescale(img, 0.5)
                img = mmcv.imrescale(s1, 2.0)
            else:
                s1 = mmcv.imrescale(img, 0.25)
                img = mmcv.imrescale(s1, 4.0)

        return img

@PIPELINES.register_module()
class RandomShift(object):
    """
    Shift the input images randomly with a probability of `prob`.
    If the image is shifted, it will be shifted by a random number of pixels between `-max_pixel` and `max_pixel`.
    Args:
        prob (float): Probability of shifting the image.
        max_pixel (int): Maximum number of pixels to shift the image.
        return_shift (bool): Whether to return the shift amount as a key in the results dictionary.
        method (str): Method to use for generating the shift amount. Can be either 'uniform' or 'normal'.
    Returns:
        dict: The result dict with shifted images and semantic segmentation maps.
    """

    def __init__(self, prob=0.5, max_pixel=5, return_shift=False, method="uniform"):
        self.prob = prob
        self.max_pixel = max_pixel
        self.return_shift = return_shift
        assert method in ["uniform", "normal"]
        self.method = method

    def __call__(self, results):
        if np.random.rand() < self.prob:
            if self.method == "uniform":
                shiftx = np.random.randint(-self.max_pixel, self.max_pixel + 1)
                shifty = np.random.randint(-self.max_pixel, self.max_pixel + 1)
            elif self.method == "normal":
                shiftx = int(
                    np.clip(np.random.normal(scale=0.4), -1, 1) * self.max_pixel
                )
                shifty = int(
                    np.clip(np.random.normal(scale=0.4), -1, 1) * self.max_pixel
                )
            img1 = results["img1"]
            rows, cols = img1.shape[:2]

            M = np.float32([[1, 0, shiftx], [0, 1, shifty]])
            img1 = cv2.warpAffine(img1, M, (cols, rows))
            results["img1"] = img1
        else:
            shiftx, shifty = 0, 0
        if self.return_shift:
            results["shift"] = -np.array([shiftx, shifty], dtype=float)
        seg_map = results.get("gt_semantic_seg", None)
        if not seg_map is None:
            if shiftx > 0:
                seg_map[:, :shiftx] = 0
            if shiftx < 0:
                seg_map[:, shiftx:] = 0
            if shifty > 0:
                seg_map[:shifty, :] = 0
            if shifty < 0:
                seg_map[shifty:, :] = 0
            results["gt_semantic_seg"] = seg_map
        return results

@PIPELINES.register_module()
class Shift(object):
    """
    Shift the input images by a fixed number of pixels.
    Args:
        pixel (int): Number of pixels to shift the image.
        return_shift (bool): Whether to return the shift amount as a key in the results dictionary.
    Returns:
        dict: The result dict with shifted images and semantic segmentation maps.
    """

    def __init__(self, pixel=5, return_shift=False):
        self.pixel = pixel
        self.return_shift = return_shift

    def __call__(self, results):
        shiftx = self.pixel
        shifty = self.pixel
        img1 = results["img1"]
        rows, cols = img1.shape[:2]

        M = np.float32([[1, 0, shiftx], [0, 1, shifty]])
        img1 = cv2.warpAffine(img1, M, (cols, rows))
        results["img1"] = img1
        if self.return_shift:
            results["shift"] = -np.array([shiftx, shifty], dtype=float)

@PIPELINES.register_module()
class ConvertMCDLabel(object):
    def __init__(
        self,
        danqi_label=False,
        annotation_type="v25",
        mcd_class_name=None,
        mcd_class_id=None,
    ):
        if annotation_type == "custom":
            self.get_custom_mapping(mcd_class_name, mcd_class_id)
        elif annotation_type == "yuzhi":
            self.get_yuzhi_mapping()
        else:
            print("Invalid annotation_type")
            return
        self.danqi_label = danqi_label
        # NOTE: share convert MCD here, mapping is very important!!!

    def get_sense_mapping(self):
        """Get Sense mapping."""

        fg_class_num = 6
        label_mapping, pre_mapping, post_mapping = [], [], []
        self.label_mapping = np.arange(1 + fg_class_num * fg_class_num)
        self.pre_mapping = (self.label_mapping - 1) // fg_class_num
        self.post_mapping = (self.label_mapping - 1) % fg_class_num
        self.pre_mapping[0] = 255
        self.post_mapping[0] = 255

    def get_custom_mapping(self, class_name, class_id):
        """Get Sense mapping."""

        basic_class = []
        for item in class_name[1:]:
            pre_name, post_name = item.split("-")
            basic_class += [pre_name, post_name]
        self.BASIC_CLASS = sorted(list(np.unique(basic_class)))
        basic_class = self.BASIC_CLASS
        pre_mapping, post_mapping = [255], [255]
        for item in class_name[1:]:
            pre_name, post_name = item.split("-")
            pre_index = self.BASIC_CLASS.index(pre_name)
            post_index = self.BASIC_CLASS.index(post_name)
            pre_mapping.append(pre_index)
            post_mapping.append(post_index)
        self.pre_mapping = np.array(pre_mapping)
        self.post_mapping = np.array(post_mapping)
        pad_num = 256 - len(self.pre_mapping)
        self.pre_mapping = np.pad(self.pre_mapping, (0, pad_num), constant_values=255)
        self.post_mapping = np.pad(self.post_mapping, (0, pad_num), constant_values=255)
        # import pickle
        # mcd_meta = {'fg_class_num': len(self.BASIC_CLASS)}
        # pickle.dump(open('mcd_meta.pkl', 'wb'), mcd_meta)

        remapping_label = []
        for pre_index, pre_name in enumerate(basic_class):
            for post_index, post_name in enumerate(basic_class):
                combined_name = f"{pre_name}-{post_name}"
                if combined_name in class_name:
                    remapping_label.append(class_name.index(combined_name))
                else:
                    remapping_label.append(0)
        # print(basic_class, class_name)
        # print(remapping_label)
        self.label_mapping = np.array(remapping_label)

    def get_yuzhi_mapping(self):
        """Get Sense mapping."""

        self.BASIC_CLASS = ["耕地", "绿地", "建筑", "动土", "自然裸地", "水体", "大棚", "构筑物"]
        fg_class_num = 8
        self.label_mapping = np.arange(1 + fg_class_num * fg_class_num)
        self.pre_mapping = (self.label_mapping - 1) // fg_class_num
        self.post_mapping = (self.label_mapping - 1) % fg_class_num
        self.pre_mapping[0] = 255
        self.post_mapping[0] = 255
        pad_num = 256 - len(self.label_mapping)
        self.pre_mapping = np.pad(self.pre_mapping, (0, pad_num), constant_values=255)
        self.post_mapping = np.pad(self.post_mapping, (0, pad_num), constant_values=255)

    def __call__(self, results):
        # results['gt_semantic_seg'] = self.label_mapping[seg_map]
        seg_map = results["gt_semantic_seg"]
        if self.danqi_label:
            pre_map = self.pre_mapping[seg_map]
            post_map = self.post_mapping[seg_map]
            results["gt_semantic_seg_sat1"] = pre_map
            results["gt_semantic_seg_sat2"] = post_map
            # NOTE: very important this debug case
            # import numpy as np
            # for item in ['gt_semantic_seg_sat1', 'gt_semantic_seg_sat2', 'gt_semantic_seg']:
            #     temp = results[item]
            #     print(np.unique(temp, return_counts=True))
            #     indexes = temp > 0
            #     temp[indexes] += 125
            #     cv2.imwrite(f'./debug/{item}.png', temp.astype('uint8'))
        return results

@PIPELINES.register_module()
class SplitMCDLabel(object):
    def __init__(self, fg_class_num):
        self.fg_class_num = fg_class_num

    def separate_map(self, combined_map):
        """Separate map."""

        combined_map = combined_map.copy()  # [0, 64]
        mask = combined_map == 0
        combined_map -= 1  # [-1, 63]
        prev = combined_map // self.fg_class_num  # [0, 7]
        post = combined_map % self.fg_class_num  # [0, 7]
        prev += 1  # [1, 8]
        post += 1
        prev[mask] = 0
        post[mask] = 0
        return prev, post

    def __call__(self, results):
        seg_map = results["gt_semantic_seg"].copy()
        # convert label to binary
        gt_semantic_seg = results["gt_semantic_seg"]
        gt_semantic_seg[gt_semantic_seg > 0] = 1
        results["gt_semantic_seg"] = gt_semantic_seg
        prev, post = self.separate_map(seg_map)
        results["gt_semantic_seg_sat1"] = prev
        results["gt_semantic_seg_sat2"] = post

        return results

@PIPELINES.register_module()
class GSDNorm(object):
    """
    Normalize the GSD of the image to a certain value.

    Args:
        gsd_info_file (str): The path of the file that contains the GSD information.
        norm_gsd (float): The normalized GSD value.
        max_ratio (int): The maximum ratio of the GSD value to the normalized GSD value.
    """

    def __init__(self, gsd_info_file, norm_gsd=1.0, max_ratio=3):
        self.norm_gsd = norm_gsd
        self.gsd_map = self.get_gsd_info(gsd_info_file)
        self.max_ratio = max_ratio

    def get_gsd_info(self, gsd_info_file):
        """
        Get the GSD information from the file.

        Args:
            gsd_info_file (str): The path of the file that contains the GSD information.

        Returns:
            dict: A dictionary that maps the base name of the image to its GSD value.
        """
        gsd_list = [0.15, 0.5, 0.8, 1.0, 1.5, 2.0]
        with open(gsd_info_file, "r") as file:
            data = file.readlines()
        gsd_map = {}
        for ele in data:
            gsd, base_name = ele.replace("\n", "").split("\t")
            gsd = float(gsd)
            base_name = base_name.replace("_sat.npy", "").split("_")[0]
            assert gsd in gsd_list
            gsd_map[base_name] = gsd
        return gsd_map

    def __call__(self, results):
        """
        Call function to normalize the GSD of the image.

        Args:
            results (dict): The input data dictionary.

        Returns:
            dict: The output data dictionary.
        """
        base_name = results["filename"][0].split("/")[-1].split("_")[0]
        gsd = self.gsd_map.get(base_name, self.norm_gsd)
        ratio = gsd / self.norm_gsd
        ratio = np.clip(ratio, 1 / self.max_ratio, self.max_ratio)
        h, w = results["img1"].shape[:2]
        scale = int(w * ratio), int(h * ratio)
        img1, scale_factor = mmcv.imrescale(results["img1"], scale, return_scale=True)
        img2, scale_factor = mmcv.imrescale(results["img2"], scale, return_scale=True)

        for key in results.get("seg_fields", []):
            gt_seg = mmcv.imrescale(results[key], scale, interpolation="nearest")
            results[key] = gt_seg
        return results

@PIPELINES.register_module()
class DoubleImageRandomBlur(object):
    """
    Pipeline for randomly blurring images.

    Args:
        prob (float): The probability of applying the blur.
        kernel_size (int): The size of the kernel for the Gaussian blur.
    """

    def __init__(self, prob=0.3, kernel_size=3):
        self.prob = prob
        self.kernel_size = kernel_size

    def __call__(self, results):
        """
        Call function to apply the blur.

        Args:
            results (dict): The input data dictionary.

        Returns:
            dict: The output data dictionary.
        """
        for key in ["img1", "img2"]:
            if np.random.rand() < self.prob:
                img = results[key]
                img = cv2.GaussianBlur(img, (self.kernel_size, self.kernel_size), 0)
                results[key] = img
        return results

@PIPELINES.register_module()
class DoubleImageRandomFog(object):
    """
    Pipeline for randomly adding fog to images.

    Args:
        prob (float): The probability of applying the fog.
        size (int): The size of the fog.
        random_color (bool): Whether to randomly change the color of the fog.
        thin (bool): Whether to make the fog thinner.
    """

    def __init__(self, prob=0.05, size=400, random_color=False, thin=False):
        import imgaug.augmenters as iaa

        self.aug = iaa.Clouds()
        self.prob = prob
        self.size = size
        self.random_color = random_color
        self.thin = thin

    def __call__(self, results):
        if np.random.rand() < self.prob:
            if np.random.rand() < 0.5:
                aug_name = "img1"
            else:
                aug_name = "img2"
            img = results[aug_name]
            h, w, c = img.shape
            size = min(min(h, w), self.size)
            y = np.random.randint(h - size + 1)
            x = np.random.randint(w - size + 1)
            img_aug = self.aug.augment_image(img)
            img_aug = self.aug.augment_image(img_aug)
            if self.thin:
                if np.random.rand() < 0.5:
                    img_aug = self.aug.augment_image(img_aug)
            else:
                img_aug = self.aug.augment_image(img_aug)
            if self.random_color:
                img_aug = img_aug.astype(int)
                img_aug -= ((np.random.rand(3)) * 70).astype(int)
                img_aug = np.clip(img_aug, 0, 255).astype(np.uint8)
            img[y : y + size, x : x + size] = img_aug[y : y + size, x : x + size]
            results[aug_name] = img
            for key in results.get("seg_fields", []):
                gt_seg = results[key]
                gt_seg[y : y + size, x : x + size] = 0
                results[key] = gt_seg

        return results

@PIPELINES.register_module()
class DoubleImageWatermark(object):
    """
    Pipeline for randomly adding watermarks to images.

    Args:
        png_list (list): A list of paths to the watermark images.
        prob (float): The probability of applying the watermark.
        min_size (int): The minimum size of the watermark.
        max_size (int): The maximum size of the watermark.
    """

    def __init__(self, png_list, prob=0.5, min_size=60, max_size=100):
        self.watermarks = []
        for png in png_list:
            self.watermarks.append(cv2.imread(png))
        self.prob = prob
        self.min_size = min_size
        self.max_size = max_size

    def __call__(self, results):
        """
        Call function to apply the watermark.

        Args:
            results (dict): The input data dictionary.

        Returns:
            dict: The output data dictionary.
        """
        if np.random.rand() < self.prob:
            if np.random.rand() < 0.5:
                aug_name = "img1"
            else:
                aug_name = "img2"
            img = results[aug_name]
            watermark = self.watermarks[np.random.randint(len(self.watermarks))]
            size = np.random.randint(self.min_size, self.max_size)
            watermark = cv2.resize(watermark, (size, size))
            h, w, c = img.shape
            y = np.random.randint((h - size))
            x = np.random.randint(w - size)
            img_tmp = img[y : y + size, x : x + size].astype(int)
            if np.random.rand() > 0.5:
                img_tmp += watermark // np.random.randint(2, 6)
            else:
                _watermark = (watermark > 0).astype(int) * 190
                noise = ((np.random.rand(*(_watermark.shape)) - 0.5) * 30).astype(int)
                _watermark[_watermark > 0] = (_watermark + noise)[_watermark > 0]
                img_tmp[_watermark > 0] = _watermark[_watermark > 0]
            img_tmp = np.clip(img_tmp, 0, 255).astype(np.uint8)
            img[y : y + size, x : x + size] = img_tmp
            results[aug_name] = img
        return results

@PIPELINES.register_module()
class SmallTargetAug(object):
    """
    Pipeline for augmenting small targets.

    Args:
        txtfile (str): The path of the file that contains the information of the small targets.
        patch_size (int): The size of the patch.
        prob (float): The probability of applying the augmentation.
        min_area (int): The minimum area of the small targets.
        max_area (int): The maximum area of the small targets.
        aug_cats (list): The categories of the small targets to be augmented.
        sample_method (str): The method of sampling the small targets.
    """

    def __init__(
        self,
        txtfile,
        patch_size=200,
        prob=0.2,
        min_area=0,
        max_area=400,
        aug_cats=None,
        sample_method="random",
    ):
        with open(txtfile, "r") as file:
            files = file.readlines()
        self.samples = []
        for sample in files:
            _, _, _, cat, area = sample.replace("\n", "").split(" ")
            cat = int(cat)
            area = float(area)
            if area < max_area and area > min_area:
                if aug_cats is None or cat in aug_cats:
                    self.samples.append(sample)
        self.patch_size = patch_size
        self.prob = prob
        self.sample_method = sample_method
        if sample_method == "queue":
            np.random.shuffle(self.samples)
            self.index = 0

    def __call__(self, results):
        """
        Call function to augment the small targets.

        Args:
            results (dict): The input data dictionary.

        Returns:
            dict: The output data dictionary.
        """
        if np.random.rand() < self.prob:
            img = results["img1"]
            h, w, c = img.shape
            y = np.random.randint(h - self.patch_size)
            x = np.random.randint(w - self.patch_size)
            if self.sample_method == "queue":
                sample = self.samples[self.index % len(self.samples)]
                self.index += 1
            else:
                sample = np.random.choice(self.samples)
            patch1, patch2, patch_mask, _, _ = sample.replace("\n", "").split(" ")
            patch1 = cv2.imread(patch1)
            patch2 = cv2.imread(patch2)
            patch_mask = cv2.imread(patch_mask, -1)

            img1 = results["img1"]
            img2 = results["img2"]
            img1[y : y + self.patch_size, x : x + self.patch_size, :] = patch1
            img2[y : y + self.patch_size, x : x + self.patch_size, :] = patch2
            results["img1"] = img1
            results["img2"] = img2

            for key in results.get("seg_fields", []):
                gt_seg = results[key]
                gt_seg[y : y + self.patch_size, x : x + self.patch_size] = patch_mask
                results[key] = gt_seg

        return results

@PIPELINES.register_module()
class GridShuffle(object):
    """
    Pipeline for shuffling the grid of the image.

    Args:
        grid_num (int): The number of grids.
        img_size (int): The size of the image.
    """

    def __init__(self, grid_num=3, img_size=896):
        self.grid_size = int(np.ceil(img_size / grid_num))
        self.pad_size = int(self.grid_size * grid_num)
        self.padding = int(self.pad_size - img_size)
        self.img_size = img_size
        self.grid_num = grid_num
        self.num = self.grid_num**2

    def __call__(self, results):
        """
        Call function to shuffle the grid of the image.

        Args:
            results (dict): The input data dictionary.

        Returns:
            dict: The output data dictionary.
        """
        grid_ind = np.arange(self.num)
        np.random.shuffle(grid_ind)

        for key in results.get("seg_fields", []) + ["img1", "img2"]:
            data = results[key]
            if len(data.shape) == 3:
                data = np.pad(data, [[0, self.padding], [0, self.padding], [0, 0]])
            else:
                data = np.pad(data, [[0, self.padding], [0, self.padding]])
            data_new = np.zeros_like(data)
            ind = 0
            for i in range(self.grid_num):
                for j in range(self.grid_num):
                    row = grid_ind[ind] // self.grid_num
                    col = grid_ind[ind] % self.grid_num
                    data_new[
                        row * self.grid_size : (row + 1) * self.grid_size,
                        col * self.grid_size : (col + 1) * self.grid_size,
                    ] = data[
                        i * self.grid_size : (i + 1) * self.grid_size,
                        j * self.grid_size : (j + 1) * self.grid_size,
                    ]
                    ind += 1
            results[key] = data_new[: self.img_size, : self.img_size]
        return results


import os
from slack_bolt import App
from slack_bolt.adapter.socket_mode import SocketModeHandler
from dotenv import load_dotenv
import openai

import sqlalchemy
import sqlalchemy.ext.declarative
import sqlalchemy.orm
import datetime

# 環境変数を読み込む
load_dotenv()
openai.api_key=os.environ.get("OPENAI_API_KEY")
# ボットトークンとソケットモードハンドラーを使ってアプリを初期化します
app = App(token=os.environ.get("SLACK_BOT_TOKEN"))
model_name = "text-davinci-003"

### パラメータを指定して使いまわす関数を定義
def davinciStrictive(prompt):
    completions = openai.Completion.create(
        prompt = prompt.replace("\n", " "),
        engine = model_name,
        max_tokens = 2048,
        temperature = 0.12,
        n = 1, stop = None
    )
    text = completions.choices[0].text
    return text

# エンジンの定義
engine = sqlalchemy.create_engine('sqlite:///database.db', echo=True)

# sqlalchemyでデータベースのテーブルを扱うための宣言
Base = sqlalchemy.ext.declarative.declarative_base()

# テーブルのフィールドを定義
class Session(Base):
    __tablename__ = 'sessions'
    id = sqlalchemy.Column(
        sqlalchemy.Integer, primary_key=True, autoincrement=True
    )
    prompt = sqlalchemy.Column(sqlalchemy.String(500))
    completion = sqlalchemy.Column(sqlalchemy.String(500))
    channel_id = sqlalchemy.Column(sqlalchemy.String(100))
    created_at = sqlalchemy.Column(sqlalchemy.DateTime)

# データベースにテーブルを作成
Base.metadata.create_all(engine)

# データベースに接続するためのセッションを準備
SessionDataBase = sqlalchemy.orm.sessionmaker(bind=engine)

class SessionManager:
    def __init__(self):
        self.session = SessionDataBase()

    def add_record(self, prompt, completion, channel_id):
        # レコードを準備し、セッションを通してデータベースに送る
        s = Session(
            prompt=prompt,
            completion=completion,
            channel_id=channel_id,
            created_at=datetime.datetime.now()
        )
        self.session.add(s)
        self.session.commit()

    def get_pair_list(self, channel_id):
        # channel_idを指定して、promptとcompletionを古い順に取得する
        record_list = self.session.query(Session).filter_by(channel_id=channel_id).order_by(Session.created_at.asc()).all()
        # promptとcompletionのペアの配列を取得する
        pair_list = [(record.prompt, record.completion) for record in record_list]
        return pair_list

    def delete_pair_list(self, channel_id):
        # データベースから指定されたchannel_idのレコードを削除する
        record_list = self.session.query(Session).filter_by(channel_id=channel_id).all()
        for record in record_list:
            self.session.delete(record)
        self.session.commit()

    def get_pair_count(self, channel_id):
        # データベースから指定されたchannel_idのレコードの件数を取得する
        count = self.session.query(Session).filter_by(channel_id=channel_id).count()
        return count

sm = SessionManager()

@app.event("app_mention")
def message_mention(say, event):
    input_prompt = event["text"].replace("<@"+os.environ.get("SLACK_BOT_ID")+">", "").strip()
    template = ""
    for prompt, completion in sm.get_pair_list(event["channel"]):
        template += "User:"+prompt.replace("\n", " ") + "\n"
        template += "Assistant:"+completion.replace("\n", "") + "\n"
    prompt = template + "User: " + input_prompt.replace("\n", "") + "\n" + "Assistant: "
    print(prompt)
    completion = davinciStrictive(prompt)
    sm.add_record(input_prompt, completion, event["channel"])
    say(completion)

@app.command("/gpt-cmd")
def greet_command(ack, say, respond, command):
    # スラッシュコマンドを受け取ったことを通知
    ack()

    # 引数を取得
    cmd = command["text"]
    # ユーザーIDを取得
    user_id = command["user_id"]
    # チャンネルIDを取得
    channel_id = command["channel_id"]

    # 引数を使用してメッセージを送信
    help_text = """
    スラッシュコマンドの後に文字列を入れると、以下のことができます
```
- history 会話の文脈を表示
- forget  このチャンネル上での会話の文脈をリセット
```
"""
    if cmd == "history":
        history = ""
        for pair in sm.get_pair_list(channel_id):
            history += f"<@{user_id}>:" + pair[0]+ "\nAI:" + pair[1] + "\n"
        respond(history)
        say("今までの会話の数は「" + str(sm.get_pair_count(channel_id))+"」")
    elif cmd == "forget":
        sm.delete_pair_list(channel_id)
        say("このチャンネル上の会話をすべてリセットしました")
    else:
        respond(help_text)

    # respond(f"{cmd}, <@{user_id}> in {channel_id}!")

# アプリを起動します
if __name__ == "__main__":
    SocketModeHandler(app, os.environ["SLACK_APP_TOKEN"]).start()


from django.contrib.auth.decorators import login_required
from django.contrib.auth.models import Group
from django.shortcuts import redirect

@login_required
def upgrade_me(request):
    user = request.user
    premium_group = Group.objects.get(name='authors')
    if not request.user.groups.filter(name='authors').exists():
        premium_group.user_set.add(user)
    return redirect('/news/')

@login_required
def downgrade_me(request):
    user = request.user
    premium_group = Group.objects.get(name='authors')
    if request.user.groups.filter(name='authors').exists():
        premium_group.user_set.remove(user)
    return redirect('/news/')