from __future__ import annotations

import re
import json
import logging

from abc import ABC, abstractmethod
from time import time
from pathlib import Path
from operator import itemgetter
from functools import partial

from PyQt5.QtCore import QCoreApplication, QRegularExpression

from novelwriter.enum import nwItemLayout
from novelwriter.common import formatTimeStamp, numberToRoman, checkInt
from novelwriter.constants import nwConst, nwHeadFmt, nwRegEx, nwUnicode
from novelwriter.core.project import NWProject

logger = logging.getLogger(__name__)

def stripEscape(text) -> str:
    """Helper function to strip escaped Markdown characters from
    paragraph text.
    """
    if "\\" in text:
        # Checking first is slightly slower when there are escaped
        # characters in the text, but significantly faster when not
        return text.replace(r"\*", "*").replace(r"\~", "~").replace(r"\_", "_")
    return text

class Tokenizer(ABC):
    """Core: Text Tokenizer Abstract Base Class

    This is the base class for all document build classes. It parses the
    novelWriter markup format and generates a registry of tokens and
    text that can be further processed into other output formats by
    subclasses.
    """

    # In-Text Format
    FMT_B_B = 1  # Begin bold
    FMT_B_E = 2  # End bold
    FMT_I_B = 3  # Begin italics
    FMT_I_E = 4  # End italics
    FMT_D_B = 5  # Begin strikeout
    FMT_D_E = 6  # End strikeout

    # Block Type
    T_EMPTY    = 1   # Empty line (new paragraph)
    T_SYNOPSIS = 2   # Synopsis comment
    T_COMMENT  = 3   # Comment line
    T_KEYWORD  = 4   # Command line
    T_TITLE    = 5   # Title
    T_UNNUM    = 6   # Unnumbered
    T_HEAD1    = 7   # Header 1
    T_HEAD2    = 8   # Header 2
    T_HEAD3    = 9   # Header 3
    T_HEAD4    = 10  # Header 4
    T_TEXT     = 11  # Text line
    T_SEP      = 12  # Scene separator
    T_SKIP     = 13  # Paragraph break

    # Block Style
    A_NONE     = 0x0000  # No special style
    A_LEFT     = 0x0001  # Left aligned
    A_RIGHT    = 0x0002  # Right aligned
    A_CENTRE   = 0x0004  # Centred
    A_JUSTIFY  = 0x0008  # Justified
    A_PBB      = 0x0010  # Page break before
    A_PBA      = 0x0020  # Page break after
    A_Z_TOPMRG = 0x0040  # Zero top margin
    A_Z_BTMMRG = 0x0080  # Zero bottom margin
    A_IND_L    = 0x0100  # Left indentation
    A_IND_R    = 0x0200  # Right indentation

    def __init__(self, project: NWProject) -> None:

        self._project = project

        # Data Variables
        self._text   = ""    # The raw text to be tokenized
        self._nwItem = None  # The NWItem associated with the handle
        self._tokens = []    # The list of the processed tokens
        self._result = ""    # The result of the last document

        self._keepMarkdown = False  # Whether to keep the markdown text
        self._allMarkdown  = []     # The result novelWriter markdown of all documents

        # User Settings
        self._textFont    = "Serif"  # Output text font
        self._textSize    = 11       # Output text size
        self._textFixed   = False    # Fixed width text
        self._lineHeight  = 1.15     # Line height in units of em
        self._blockIndent = 4.00     # Block indent in units of em
        self._doJustify   = False    # Justify text
        self._doBodyText  = True     # Include body text
        self._doSynopsis  = False    # Also process synopsis comments
        self._doComments  = False    # Also process comments
        self._doKeywords  = False    # Also process keywords like tags and references

        # Margins
        self._marginTitle = (1.000, 0.500)
        self._marginHead1 = (1.000, 0.500)
        self._marginHead2 = (0.834, 0.500)
        self._marginHead3 = (0.584, 0.500)
        self._marginHead4 = (0.584, 0.500)
        self._marginText  = (0.000, 0.584)
        self._marginMeta  = (0.000, 0.584)

        # Title Formats
        self._fmtTitle   = nwHeadFmt.TITLE  # Formatting for titles
        self._fmtChapter = nwHeadFmt.TITLE  # Formatting for numbered chapters
        self._fmtUnNum   = nwHeadFmt.TITLE  # Formatting for unnumbered chapters
        self._fmtScene   = nwHeadFmt.TITLE  # Formatting for scenes
        self._fmtSection = nwHeadFmt.TITLE  # Formatting for sections

        self._hideScene   = False  # Do not include scene headers
        self._hideSection = False  # Do not include section headers

        self._linkHeaders = False  # Add an anchor before headers

        # Instance Variables
        self._hFormatter = HeadingFormatter(self._project)
        self._firstScene = False  # Flag to indicate that the first scene of the chapter

        # This File
        self._isNone  = False  # Document has unknown layout
        self._isNovel = False  # Document is a novel document
        self._isNote  = False  # Document is a project note
        self._isFirst = True   # Document is the first in a set

        # Error Handling
        self._errData = []

        # Function Mapping
        self._localLookup = self._project.localLookup
        self.tr = partial(QCoreApplication.translate, "Tokenizer")

        # Cached Translations
        self._trSynopsis = self.tr("Synopsis")

        return

    ##
    #  Properties
    ##

    @property
    def theResult(self) -> str:
        """The result of the build process."""
        return self._result

    @property
    def theMarkdown(self) -> list:
        """The combined novelWriter Markdown text."""
        return self._allMarkdown

    @property
    def errData(self) -> list:
        """The error data."""
        return self._errData

    ##
    #  Setters
    ##

    def setTitleFormat(self, hFormat: str) -> None:
        """Set the title format pattern."""
        self._fmtTitle = hFormat.strip()
        return

    def setChapterFormat(self, hFormat: str) -> None:
        """Set the chapert format pattern."""
        self._fmtChapter = hFormat.strip()
        return

    def setUnNumberedFormat(self, hFormat: str) -> None:
        """Set the unnumbered format pattern."""
        self._fmtUnNum = hFormat.strip()
        return

    def setSceneFormat(self, hFormat: str, hide: bool) -> None:
        """Set the scene format pattern and hidden status."""
        self._fmtScene = hFormat.strip()
        self._hideScene = hide
        return

    def setSectionFormat(self, hFormat: str, hide: bool) -> None:
        """Set the section format pattern and hidden status."""
        self._fmtSection = hFormat.strip()
        self._hideSection = hide
        return

    def setFont(self, family: str, size: int, isFixed: bool = False) -> None:
        """Set the build font."""
        self._textFont = family
        self._textSize = round(int(size))
        self._textFixed = isFixed
        return

    def setLineHeight(self, height: float) -> None:
        """Set the line height between 0.5 and 5.0."""
        self._lineHeight = min(max(float(height), 0.5), 5.0)
        return

    def setBlockIndent(self, indent: float) -> None:
        """Set the block indent between 0.0 and 10.0."""
        self._blockIndent = min(max(float(indent), 0.0), 10.0)
        return

    def setJustify(self, state: bool) -> None:
        """Enable or disable text justification."""
        self._doJustify = state
        return

    def setTitleMargins(self, upper: float, lower: float) -> None:
        """Set the upper and lower title margin."""
        self._marginTitle = (float(upper), float(lower))
        return

    def setHead1Margins(self, upper: float, lower: float) -> None:
        """Set the upper and lower header 1 margin."""
        self._marginHead1 = (float(upper), float(lower))
        return

    def setHead2Margins(self, upper: float, lower: float) -> None:
        """Set the upper and lower header 2 margin."""
        self._marginHead2 = (float(upper), float(lower))
        return

    def setHead3Margins(self, upper: float, lower: float) -> None:
        """Set the upper and lower header 3 margin."""
        self._marginHead3 = (float(upper), float(lower))
        return

    def setHead4Margins(self, upper: float, lower: float) -> None:
        """Set the upper and lower header 4 margin."""
        self._marginHead4 = (float(upper), float(lower))
        return

    def setTextMargins(self, upper: float, lower: float) -> None:
        """Set the upper and lower text margin."""
        self._marginText = (float(upper), float(lower))
        return

    def setMetaMargins(self, upper: float, lower: float) -> None:
        """Set the upper and lower meta text margin."""
        self._marginMeta = (float(upper), float(lower))
        return

    def setLinkHeaders(self, state: bool) -> None:
        """Enable or disable adding an anchor before headers."""
        self._linkHeaders = state
        return

    def setBodyText(self, state: bool) -> None:
        """Include body text in build."""
        self._doBodyText = state
        return

    def setSynopsis(self, state: bool) -> None:
        """Include synopsis comments in build."""
        self._doSynopsis = state
        return

    def setComments(self, state: bool) -> None:
        """Include comments in build."""
        self._doComments = state
        return

    def setKeywords(self, state: bool) -> None:
        """Include keywords in build."""
        self._doKeywords = state
        return

    def setKeepMarkdown(self, state: bool) -> None:
        """Keep original markdown during build."""
        self._keepMarkdown = state
        return

    ##
    #  Class Methods
    ##

    @abstractmethod
    def doConvert(self) -> None:
        raise NotImplementedError

    def addRootHeading(self, tHandle: str) -> bool:
        """Add a heading at the start of a new root folder."""
        tItem = self._project.tree[tHandle]
        if not tItem or not tItem.isRootType():
            return False

        if self._isFirst:
            textAlign = self.A_CENTRE
            self._isFirst = False
        else:
            textAlign = self.A_PBB | self.A_CENTRE

        trNotes = self._localLookup("Notes")
        title = f"{trNotes}: {tItem.itemName}"
        self._tokens = []
        self._tokens.append((
            self.T_TITLE, 0, title, None, textAlign
        ))
        if self._keepMarkdown:
            self._allMarkdown.append(f"# {title}\n\n")

        return True

    def setText(self, tHandle: str, text: str | None = None) -> bool:
        """Set the text for the tokenizer from a handle. If theText is
        not set, load it from the file.
        """
        self._nwItem = self._project.tree[tHandle]
        if self._nwItem is None:
            return False

        if text is None:
            text = self._project.storage.getDocument(tHandle).readDocument() or ""

        self._text = text

        docSize = len(self._text)
        if docSize > nwConst.MAX_DOCSIZE:
            errVal = self.tr("Document '{0}' is too big ({1} MB). Skipping.").format(
                self._nwItem.itemName, f"{docSize/1.0e6:.2f}"
            )
            self._text = "# {0}\n\n{1}\n\n".format(self.tr("ERROR"), errVal)
            self._errData.append(errVal)

        self._isNone  = self._nwItem.itemLayout == nwItemLayout.NO_LAYOUT
        self._isNovel = self._nwItem.itemLayout == nwItemLayout.DOCUMENT
        self._isNote  = self._nwItem.itemLayout == nwItemLayout.NOTE

        return True

    def doPreProcessing(self) -> None:
        """Run trough the various replace dictionaries."""
        # Process the user's auto-replace dictionary
        autoReplace = self._project.data.autoReplace
        if len(autoReplace) > 0:
            repDict = {}
            for aKey, aVal in autoReplace.items():
                repDict[f"<{aKey}>"] = aVal
            xRep = re.compile("|".join([re.escape(k) for k in repDict.keys()]), flags=re.DOTALL)
            self._text = xRep.sub(lambda x: repDict[x.group(0)], self._text)

        # Process the character translation map
        trDict = {nwUnicode.U_MAPOSS: nwUnicode.U_RSQUO}
        self._text = self._text.translate(str.maketrans(trDict))

        return

    def tokenizeText(self) -> None:
        """Scan the text for either lines starting with specific
        characters that indicate headers, comments, commands etc, or
        just contain plain text. In the case of plain text, apply the
        same RegExes that the syntax highlighter uses and save the
        locations of these formatting tags into the token array.

        The format of the token list is an entry with a five-tuple for
        each line in the file. The tuple is as follows:
          1: The type of the block, self.T_*
          2: The line in the file where this block occurred
          3: The text content of the block, without leading tags
          4: The internal formatting map of the text, self.FMT_*
          5: The style of the block, self.A_*
        """
        # RegExes for adding formatting tags within text lines
        rxFormats = [
            (QRegularExpression(nwRegEx.FMT_EI), [None, self.FMT_I_B, None, self.FMT_I_E]),
            (QRegularExpression(nwRegEx.FMT_EB), [None, self.FMT_B_B, None, self.FMT_B_E]),
            (QRegularExpression(nwRegEx.FMT_ST), [None, self.FMT_D_B, None, self.FMT_D_E]),
        ]

        self._tokens = []
        tmpMarkdown = []
        nLine = 0
        breakNext = False
        for aLine in self._text.splitlines():
            nLine += 1
            sLine = aLine.strip()

            # Check for blank lines
            if len(sLine) == 0:
                self._tokens.append((
                    self.T_EMPTY, nLine, "", None, self.A_NONE
                ))
                if self._keepMarkdown:
                    tmpMarkdown.append("\n")

                continue

            if breakNext:
                sAlign = self.A_PBB
                breakNext = False
            else:
                sAlign = self.A_NONE

            # Check Line Format
            # =================

            if aLine[0] == "[":
                # Parse special formatting line

                if sLine in ("[NEWPAGE]", "[NEW PAGE]"):
                    breakNext = True
                    continue

                elif sLine == "[VSPACE]":
                    self._tokens.append(
                        (self.T_SKIP, nLine, "", None, sAlign)
                    )
                    continue

                elif sLine.startswith("[VSPACE:") and sLine.endswith("]"):
                    nSkip = checkInt(sLine[8:-1], 0)
                    if nSkip >= 1:
                        self._tokens.append(
                            (self.T_SKIP, nLine, "", None, sAlign)
                        )
                    if nSkip > 1:
                        self._tokens += (nSkip - 1) * [
                            (self.T_SKIP, nLine, "", None, self.A_NONE)
                        ]
                    continue

            elif aLine[0] == "%":
                cLine = aLine[1:].lstrip()
                synTag = cLine[:9].lower()
                if synTag == "synopsis:":
                    self._tokens.append((
                        self.T_SYNOPSIS, nLine, cLine[9:].strip(), None, sAlign
                    ))
                    if self._doSynopsis and self._keepMarkdown:
                        tmpMarkdown.append("%s\n" % aLine)
                else:
                    self._tokens.append((
                        self.T_COMMENT, nLine, aLine[1:].strip(), None, sAlign
                    ))
                    if self._doComments and self._keepMarkdown:
                        tmpMarkdown.append("%s\n" % aLine)

            elif aLine[0] == "@":
                self._tokens.append((
                    self.T_KEYWORD, nLine, aLine[1:].strip(), None, sAlign
                ))
                if self._doKeywords and self._keepMarkdown:
                    tmpMarkdown.append("%s\n" % aLine)

            elif aLine[:2] == "# ":
                if self._isNovel:
                    sAlign |= self.A_CENTRE
                    sAlign |= self.A_PBB

                self._tokens.append((
                    self.T_HEAD1, nLine, aLine[2:].strip(), None, sAlign
                ))
                if self._keepMarkdown:
                    tmpMarkdown.append("%s\n" % aLine)

            elif aLine[:3] == "## ":
                if self._isNovel:
                    sAlign |= self.A_PBB

                self._tokens.append((
                    self.T_HEAD2, nLine, aLine[3:].strip(), None, sAlign
                ))
                if self._keepMarkdown:
                    tmpMarkdown.append("%s\n" % aLine)

            elif aLine[:4] == "### ":
                self._tokens.append((
                    self.T_HEAD3, nLine, aLine[4:].strip(), None, sAlign
                ))
                if self._keepMarkdown:
                    tmpMarkdown.append("%s\n" % aLine)

            elif aLine[:5] == "#### ":
                self._tokens.append((
                    self.T_HEAD4, nLine, aLine[5:].strip(), None, sAlign
                ))
                if self._keepMarkdown:
                    tmpMarkdown.append("%s\n" % aLine)

            elif aLine[:3] == "#! ":
                if self._isNovel:
                    tStyle = self.T_TITLE
                else:
                    tStyle = self.T_HEAD1

                self._tokens.append((
                    tStyle, nLine, aLine[3:].strip(), None, sAlign | self.A_CENTRE
                ))
                if self._keepMarkdown:
                    tmpMarkdown.append("%s\n" % aLine)

            elif aLine[:4] == "##! ":
                if self._isNovel:
                    tStyle = self.T_UNNUM
                    sAlign |= self.A_PBB
                else:
                    tStyle = self.T_HEAD2

                self._tokens.append((
                    tStyle, nLine, aLine[4:].strip(), None, sAlign
                ))
                if self._keepMarkdown:
                    tmpMarkdown.append("%s\n" % aLine)

            else:
                if not self._doBodyText:
                    # Skip all body text
                    continue

                # Check Alignment and Indentation
                alnLeft = False
                alnRight = False
                indLeft = False
                indRight = False
                if aLine.startswith(">>"):
                    alnRight = True
                    aLine = aLine[2:].lstrip(" ")
                elif aLine.startswith(">"):
                    indLeft = True
                    aLine = aLine[1:].lstrip(" ")

                if aLine.endswith("<<"):
                    alnLeft = True
                    aLine = aLine[:-2].rstrip(" ")
                elif aLine.endswith("<"):
                    indRight = True
                    aLine = aLine[:-1].rstrip(" ")

                if alnLeft and alnRight:
                    sAlign |= self.A_CENTRE
                elif alnLeft:
                    sAlign |= self.A_LEFT
                elif alnRight:
                    sAlign |= self.A_RIGHT

                if indLeft:
                    sAlign |= self.A_IND_L
                if indRight:
                    sAlign |= self.A_IND_R

                # Otherwise we use RegEx to find formatting tags within a line of text
                fmtPos = []
                for theRX, theKeys in rxFormats:
                    rxThis = theRX.globalMatch(aLine, 0)
                    while rxThis.hasNext():
                        rxMatch = rxThis.next()
                        for n in range(1, len(theKeys)):
                            if theKeys[n] is not None:
                                xPos = rxMatch.capturedStart(n)
                                xLen = rxMatch.capturedLength(n)
                                fmtPos.append([xPos, xLen, theKeys[n]])

                # Save the line as is, but append the array of formatting locations
                # sorted by position
                fmtPos = sorted(fmtPos, key=itemgetter(0))
                self._tokens.append((
                    self.T_TEXT, nLine, aLine, fmtPos, sAlign
                ))
                if self._keepMarkdown:
                    tmpMarkdown.append("%s\n" % aLine)

        # If we have content, turn off the first page flag
        if self._isFirst and self._tokens:
            self._isFirst = False

            # Make sure the token array doesn't start with a page break
            # on the very first page, adding a blank first page.
            if self._tokens[0][4] & self.A_PBB:
                tToken = self._tokens[0]
                self._tokens[0] = (
                    tToken[0], tToken[1], tToken[2], tToken[3], tToken[4] & ~self.A_PBB
                )

        # Always add an empty line at the end of the file
        self._tokens.append((
            self.T_EMPTY, nLine, "", None, self.A_NONE
        ))
        if self._keepMarkdown:
            tmpMarkdown.append("\n")

        if self._keepMarkdown:
            self._allMarkdown.append("".join(tmpMarkdown))

        # Second Pass
        # ===========
        # Some items need a second pass

        pToken = (self.T_EMPTY, 0, "", None, self.A_NONE)
        nToken = (self.T_EMPTY, 0, "", None, self.A_NONE)
        tCount = len(self._tokens)
        for n, tToken in enumerate(self._tokens):

            if n > 0:
                pToken = self._tokens[n-1]
            if n < tCount - 1:
                nToken = self._tokens[n+1]

            if tToken[0] == self.T_KEYWORD:
                aStyle = tToken[4]
                if pToken[0] == self.T_KEYWORD:
                    aStyle |= self.A_Z_TOPMRG
                if nToken[0] == self.T_KEYWORD:
                    aStyle |= self.A_Z_BTMMRG
                self._tokens[n] = (
                    tToken[0], tToken[1], tToken[2], tToken[3], aStyle
                )

        return

    def doHeaders(self) -> bool:
        """Apply formatting to the text headers for novel files. This
        also applies chapter and scene numbering.
        """
        if not self._isNovel:
            return False

        for n, tToken in enumerate(self._tokens):

            # In case we see text before a scene, we reset the flag
            if tToken[0] == self.T_TEXT:
                self._firstScene = False

            elif tToken[0] == self.T_HEAD1:
                # Partition

                tTemp = self._hFormatter.apply(self._fmtTitle, tToken[2])
                self._tokens[n] = (
                    tToken[0], tToken[1], tTemp, None, tToken[4]
                )

            elif tToken[0] in (self.T_HEAD2, self.T_UNNUM):
                # Chapter

                # Numbered or Unnumbered
                if tToken[0] == self.T_UNNUM:
                    tTemp = self._hFormatter.apply(self._fmtUnNum, tToken[2])
                else:
                    self._hFormatter.incChapter()
                    tTemp = self._hFormatter.apply(self._fmtChapter, tToken[2])

                # Format the chapter header
                self._tokens[n] = (
                    tToken[0], tToken[1], tTemp, None, tToken[4]
                )

                # Set scene variables
                self._firstScene = True
                self._hFormatter.resetScene()

            elif tToken[0] == self.T_HEAD3:
                # Scene

                self._hFormatter.incScene()

                tTemp = self._hFormatter.apply(self._fmtScene, tToken[2])
                if tTemp == "" and self._hideScene:
                    self._tokens[n] = (
                        self.T_EMPTY, tToken[1], "", None, self.A_NONE
                    )
                elif tTemp == "" and not self._hideScene:
                    if self._firstScene:
                        self._tokens[n] = (
                            self.T_EMPTY, tToken[1], "", None, self.A_NONE
                        )
                    else:
                        self._tokens[n] = (
                            self.T_SKIP, tToken[1], "", None, tToken[4]
                        )
                elif tTemp == self._fmtScene:
                    if self._firstScene:
                        self._tokens[n] = (
                            self.T_EMPTY, tToken[1], "", None, self.A_NONE
                        )
                    else:
                        self._tokens[n] = (
                            self.T_SEP, tToken[1], tTemp, None, tToken[4] | self.A_CENTRE
                        )
                else:
                    self._tokens[n] = (
                        tToken[0], tToken[1], tTemp, None, tToken[4]
                    )

                self._firstScene = False

            elif tToken[0] == self.T_HEAD4:
                # Section

                tTemp = self._hFormatter.apply(self._fmtSection, tToken[2])
                if tTemp == "" and self._hideSection:
                    self._tokens[n] = (
                        self.T_EMPTY, tToken[1], "", None, self.A_NONE
                    )
                elif tTemp == "" and not self._hideSection:
                    self._tokens[n] = (
                        self.T_SKIP, tToken[1], "", None, tToken[4]
                    )
                elif tTemp == self._fmtSection:
                    self._tokens[n] = (
                        self.T_SEP, tToken[1], tTemp, None, tToken[4] | self.A_CENTRE
                    )
                else:
                    self._tokens[n] = (
                        tToken[0], tToken[1], tTemp, None, tToken[4]
                    )

        return True

    def saveRawMarkdown(self, path: str | Path) -> None:
        """Save the raw text to a plain text file."""
        with open(path, mode="w", encoding="utf-8") as outFile:
            for nwdPage in self._allMarkdown:
                outFile.write(nwdPage)
        return

    def saveRawMarkdownJSON(self, path: str | Path) -> None:
        """Save the raw text to a JSON file."""
        timeStamp = time()
        data = {
            "meta": {
                "projectName": self._project.data.name,
                "novelTitle": self._project.data.title,
                "novelAuthor": self._project.data.author,
                "buildTime": int(timeStamp),
                "buildTimeStr": formatTimeStamp(timeStamp),
            },
            "text": {
                "nwd": [page.rstrip("\n").split("\n") for page in self._allMarkdown],
            }
        }
        with open(path, mode="w", encoding="utf-8") as fObj:
            json.dump(data, fObj, indent=2)
        return

# END Class Tokenizer

class HeadingFormatter:

    def __init__(self, project: NWProject) -> None:
        self._project = project
        self._chCount = 0
        self._scChCount = 0
        self._scAbsCount = 0
        return

    def incChapter(self) -> None:
        """Increment the chapter counter."""
        self._chCount += 1
        return

    def incScene(self) -> None:
        """Increment the scene counters."""
        self._scChCount += 1
        self._scAbsCount += 1
        return

    def resetScene(self) -> None:
        """Reset the chapter scene counter."""
        self._scChCount = 0
        return

    def apply(self, hFormat: str, text: str) -> str:
        """Apply formatting to a specific heading."""
        hFormat = hFormat.replace(nwHeadFmt.TITLE, text)
        hFormat = hFormat.replace(nwHeadFmt.CH_NUM, str(self._chCount))
        hFormat = hFormat.replace(nwHeadFmt.SC_NUM, str(self._scChCount))
        hFormat = hFormat.replace(nwHeadFmt.SC_ABS, str(self._scAbsCount))
        if nwHeadFmt.CH_WORD in hFormat:
            chWord = self._project.localLookup(self._chCount)
            hFormat = hFormat.replace(nwHeadFmt.CH_WORD, chWord)
        if nwHeadFmt.CH_ROML in hFormat:
            chRom = numberToRoman(self._chCount, toLower=True)
            hFormat = hFormat.replace(nwHeadFmt.CH_ROML, chRom)
        if nwHeadFmt.CH_ROMU in hFormat:
            chRom = numberToRoman(self._chCount, toLower=False)
            hFormat = hFormat.replace(nwHeadFmt.CH_ROMU, chRom)

        return hFormat

# END Class HeadingFormatter


# Follow the greens mission container Class
# Keeps all information handy. Dispatches intruction to do things.
#
# Cannot use Follow the greens.
# We are sorry. We cannot provide Follow the greens service at this airport.
# Reasons:
# This airport does not have a routing network of taxiway.
#
# Can use Follow the greens, but other issue:
# We are sorry. We cannot provide Follow the greens service now.
# Reasons:
# You are too far from the taxiways.
# We could not find a suitable route to your destination.
#
import logging
import xp

from .aircraft import Aircraft
from .airport import Airport
from .lightstring import LightString
from .ui import UIUtil

logging.basicConfig(level=logging.DEBUG)  # filename=('FTG_log.txt')

class ShowTaxiways:

    # Internal status
    STATUS = {
        "NEW": "NEW",
        "INITIALIZED": "INIT",
        "READY": "READY",
        "ACTIVE": "ACTIVE"
    }

    def __init__(self, pi):
        self.__status = ShowTaxiways.STATUS["NEW"]
        self.pi = pi
        self.airport = None
        self.aircraft = None
        self.lights = None
        self.ui = UIUtil(self)      # Where windows are built

    def start(self):
        logging.info("ShowTaxiways::status: %s, %s.", self.__status, self.ui.mainWindowExists())

        if self.ui.mainWindowExists():
            logging.debug("ShowTaxiways::start: mainWindow exists, changing visibility %s.", self.ui.isMainWindowVisible())
            self.ui.toggleVisibilityMainWindow()
            return 1

        # Info 1
        logging.info("ShowTaxiways::start: starting..")
        mainWindow = self.getAirport()
        logging.debug("ShowTaxiways::start: mainWindow created")
        if mainWindow and not xp.isWidgetVisible(mainWindow):
            xp.showWidget(mainWindow)
            logging.debug("ShowTaxiways::start: mainWindow shown")
        logging.info("ShowTaxiways::start: ..started.")
        return 1  # window displayed

    def getAirport(self):
        # Search for airport or prompt for one.
        # If airport is not equiped, we loop here until we get a suitable airport.
        # When one is given and satisfies the condition for FTG
        # we go to next step: Find the end point of Follow the greens.
        # @todo: We need to guess those from dataref
        # Note: Aircraft should be "created" outside of FollowTheGreen
        # and passed to start or getAirport. That way, we can instanciate
        # individual FollowTheGreen for numerous aircrafts.
        self.aircraft = Aircraft("A321", "D", "OO-123", "PO-123")

        pos = self.aircraft.position()
        if pos is None:
            logging.debug("ShowTaxiways::getAirport: no plane position")
            return self.ui.sorry("We could not locate your plane.")

        if pos[0] == 0 and pos[1] == 0:
            logging.debug("ShowTaxiways::getAirport: no plane position")
            return self.ui.sorry("We could not locate your plane.")

        # Info 2
        logging.info("ShowTaxiways::getAirport: Plane postion %s" % pos)
        airport = self.aircraft.airport(pos)
        if airport is None:
            logging.debug("ShowTaxiways::getAirport: no airport")
            return self.ui.promptForAirport()  # prompt for airport will continue with getDestination(airport)

        if airport.name == "NOT FOUND":
            logging.debug("ShowTaxiways::getAirport: no airport (not found)")
            return self.ui.promptForAirport()  # prompt for airport will continue with getDestination(airport)

        # Info 3
        logging.info("ShowTaxiways::getAirport: At %s" % airport.name)
        return self.showTaxiways(airport.navAidID)

    def showTaxiways(self, airport):
        if not self.airport:
            self.airport = Airport(airport)
            status = self.airport.prepare()  # [ok, errmsg]
            if not status[0]:
                logging.debug("ShowTaxiways::showTaxiways: airport not ready: %s" % (status[1]))
                return self.ui.sorry(status[1])
        else:
            logging.debug("ShowTaxiways::showTaxiways: airport already loaded")

        logging.debug("ShowTaxiways::showTaxiways: airport ready")

        self.lights = LightString()
        self.lights.showAll(self.airport)
        if len(self.lights.lights) == 0:
            logging.debug("ShowTaxiways::showTaxiways: no lights")
            return self.ui.sorry("We could not light taxiways.")

        # Info 13
        logging.info("ShowTaxiways::showTaxiways: Added %d lights.", len(self.lights.lights))

        if self.pi is not None and self.pi.menuIdx is not None and self.pi.menuIdx >= 0:
            xp.checkMenuItem(xp.findPluginsMenu(), self.pi.menuIdx, xp.Menu_Checked)
            logging.debug(f"ShowTaxiways::showTaxiways: menu checked ({self.pi.menuIdx})")
        else:
            logging.debug(f"ShowTaxiways::showTaxiways: menu not checked ({self.pi.menuIdx})")

        return self.ui.enjoy()
        # return self.ui.sorry("Follow the greens is not completed yet.")  # development

    def cancel(self, reason=""):
        if self.lights:
            self.lights.destroy()
            self.lights = None
            if self.pi is not None and self.pi.menuIdx is not None and self.pi.menuIdx >= 0:
                try:
                    xp.checkMenuItem(xp.findPluginsMenu(), self.pi.menuIdx, xp.Menu_Unchecked)
                    logging.debug(f"ShowTaxiways::cancel: menu unchecked ({self.pi.menuIdx})")
                except:
                    logging.debug(f"ShowTaxiways::cancel: menu not unchecked ({self.pi.menuIdx}, {xp.Menu_Unchecked})", exc_info=True)

        if self.ui.mainWindowExists():
            self.ui.destroyMainWindow()
            # self.ui = None

        # Info 16
        logging.info("ShowTaxiways::cancel: cancelled: %s.", reason)
        return [True, ""]

    def disable(self):
        # alias to cancel
        return self.cancel("disabled")

    def stop(self):
        # alias to cancel
        return self.cancel("stopped")


#!/usr/bin/env python3

# This solution only passes the first two test sets
# Run time complexity is O(M*S)

DEBUG = False

def solve(M, R, G):
    # print(M, R, G)

    G.insert(0, 0) # make the indexes work without convertion
    def next_metal(r, visited):
        if r not in R:
            visited.add(r)
            return False

        r1, r2 = R[r]
        if r1 == 1 or r2 == 1:
            return False

        if G[r1] > 0 and G[r2] > 0:
            G[r] += 1
            G[r1] -= 1
            G[r2] -= 1
            visited.add(r)
            return True
        elif G[r1] == 0 and r1 not in visited:
            visited.add(r)
            return next_metal(r1, visited)
        elif G[r2] == 0 and r2 not in visited:
            visited.add(r)
            return next_metal(r2, visited)
        else:
            visited.add(r)
            return False

    while next_metal(1, set([1])):
        pass

    return G[1]

if DEBUG:
    print("----------------------------")
    print(solve(3, {1: [2, 3], 2: [1, 3], 3: [1, 2]}, [5, 2, 3]))
    print("----------------------------")
    print(solve(4, {1: [2, 3], 2: [3, 4], 3: [2, 4], 4: [2, 3]}, [5, 2, 3, 2]))
    print("----------------------------")
    print(solve(5, {1: [3, 4], 2: [3, 4], 3: [4, 5], 4: [3, 5], 5: [1, 3]}, [0, 8, 6, 2, 4]))
    print("----------------------------")
    print(solve(4, {1: [3, 4], 2: [2, 3], 3: [2, 3], 4: [2, 3]}, [0, 1, 1, 0]))
    print("----------------------------")
    print(solve(4, {1: [3, 4], 2: [2, 3], 3: [2, 3], 4: [2, 3]}, [0, 0, 0, 0]))

if __name__ == "__main__":
    # input() reads a string with a line of input, stripping the ' ' (newline) at the end.
    T = int(input())
    for i in range(1, T + 1):
        M = int(input())
        R = {}
        for r in range(1, M + 1):
            R[r] = list(map(int, input().split()))

        G = list(map(int, input().split()))
        result = solve(M, R, G)
        print("Case #{}: {}".format(i, result))


# Demo that takes a random audit log sequence and predicts the next set of actions.
import argparse
import inspect
import os
import pickle
import sys
from collections import defaultdict

import pandas as pd
import scipy.stats
import torch
import yaml
from matplotlib.axes import Axes
from pandas import DataFrame
from torch.utils.data import DataLoader
from tqdm import tqdm
from tabulate import tabulate
from matplotlib import pyplot as plt
from transformers import BeamSearchScorer, StoppingCriteriaList, MaxLengthCriteria, LogitsProcessorList

from model.model import EHRAuditGPT2, EHRAuditRWKV, EHRAuditLlama
from model.modules import EHRAuditPretraining, EHRAuditDataModule
from model.data import timestamp_space_calculation
from model.vocab import EHRVocab, EHRAuditTokenizer, EHRAuditLogitsProcessor
import tikzplotlib
import numpy as np
import evaluate

# Fyi: this is a quick-and-dirty way of id'ing the columns, will need to be changed if the tabularization changes
METRIC_NAME_COL = 0
PAT_ID_COL = 1
ACCESS_TIME_COL = 2

class GenerationExperiment:
    def __init__(
        self,
        config: dict,
        path_prefix: str,
        vocab: EHRVocab,
        model: str,
        *args,
        **kwargs,
    ):
        self.config = config
        self.path_prefix = path_prefix
        self.vocab = vocab
        self.model = model

    def stopping_criteria(self, context_length: int = 0, total_length: int = 0):
        return None

    def eval_generation(self, output_df=None, label_df=None, output_tokens=None, label_tokens=None):
        pass

    def _exp_cache_path(self):
        return os.path.normpath(
            os.path.join(
                self.path_prefix,
                self.config["results_path"],
                f"exp_cache_{self.__class__.__name__}",
            )
        )

    def on_finish(self):
        pass

    def plot(self):
        pass

    # None = all preceding, # = fixed preceding, float = fraction of preceding
    def window_size(self):
        return None

    def min_size(self):
        return 2

    def examples_seen(self):
        return 0

# Next whole-action prediction
# Next METRIC_NAME prediction
class NextActionExperiment(GenerationExperiment):
    def __init__(self, config: dict,
                path_prefix: str,
                vocab: EHRVocab,
                model: str,
                *args,
                **kwargs):
        super().__init__(config, path_prefix, vocab, model, *args, **kwargs)
        self.total_seen = 0
        self.correct_by_field = defaultdict(int)
        self.total_correct = 0

    def window_size(self):
        return 1

    def stopping_criteria(self, context_length: int = 0, total_length: int = 0):
        return StoppingCriteriaList(
            [
                MaxLengthCriteria(max_length=context_length + len(self.vocab.field_names(include_special=False))),
            ]
        )

    def eval_generation(self, output_df=None, label_df=None, output_tokens=None, label_tokens=None):
        # Get the next action
        next_action = label_df.iloc[0]
        # Get the predicted next action
        predicted_next_action = output_df.iloc[0]
        # Check if the predicted next action is correct
        for x in [METRIC_NAME_COL, PAT_ID_COL, ACCESS_TIME_COL]:
            if next_action[x] == predicted_next_action[x]:
                self.correct_by_field[x] += 1
        if np.all(next_action == predicted_next_action):
            self.total_correct += 1
        self.total_seen += 1

    def examples_seen(self):
        return self.total_seen

    def on_finish(self):
        model_type = self.model.replace(os.sep, "_")
        model_path = os.path.join(
            self._exp_cache_path(), f"{model_type}.pkl"
        )
        os.makedirs(os.path.dirname(model_path), exist_ok=True)
        with open(model_path, "wb") as f:
            pickle.dump(
                {
                    "total_seen": self.total_seen,
                    "correct_by_field": self.correct_by_field,
                    "total_correct": self.total_correct,
                },
                f,
            )

    def plot(self):
        model_versions = os.listdir(self._exp_cache_path())
        total_seen_by_model = defaultdict(int)
        metric_correct_by_model = defaultdict(int)
        pat_id_correct_by_model = defaultdict(int)
        access_time_correct_by_model = defaultdict(int)
        total_correct_by_model = defaultdict(int)

        # Convert to a LaTeX table
        rows = []
        for model_version in model_versions:
            model_path = os.path.join(
                self._exp_cache_path(), model_version
            )
            model_fs = model_version.split("_")
            model_t = model_fs[0]
            model_type = model_t + "-" + ".".join(model_fs[1:3])

            with open(model_path, "rb") as f:
                model_data = pickle.load(f)
            total_seen_by_model[model_type] = model_data["total_seen"]
            metric_correct_by_model[model_type] = model_data["correct_by_field"][METRIC_NAME_COL]
            pat_id_correct_by_model[model_type] = model_data["correct_by_field"][PAT_ID_COL]
            access_time_correct_by_model[model_type] = model_data["correct_by_field"][ACCESS_TIME_COL]
            total_correct_by_model[model_type] = model_data["total_correct"]

            rows.append({
                "Model": model_type,
                "Total Seen": model_data["total_seen"],
                "METRIC_NAME Accuracy": metric_correct_by_model[model_type] / total_seen_by_model[model_type],
                "PAT_ID Accuracy": pat_id_correct_by_model[model_type] / total_seen_by_model[model_type],
                "ACCESS_TIME Accuracy": access_time_correct_by_model[model_type] / total_seen_by_model[model_type],
                "Total Accuracy": total_correct_by_model[model_type] / total_seen_by_model[model_type],
            })
        df = pd.DataFrame(rows)

        df = df.sort_values(by=["METRIC_NAME Accuracy"], ascending=True)
        # Drop the Total Seen column if all values are the same
        if len(df["Total Seen"].unique()) == 1:
            df = df.drop(columns=["Total Seen"])

        out = df.to_latex(index=False,
                    float_format="%.3f",
                    caption="Results for predicting the next action in an audit log sequence.",
                    label="tab:next_action_results")

        print(out)

class ScoringExperiment(GenerationExperiment):
    # Gets the ROUGE scores for the generated rows.
    def __init__(self, config: dict,
                path_prefix: str,
                vocab: EHRVocab,
                model: str,
                *args,
                **kwargs):
        super().__init__(config, path_prefix, vocab, model, *args, **kwargs)
        self.rouge_L_scores_by_field = defaultdict(list)
        self.rouge_scores_by_field = defaultdict(list)
        self.total_seen = 0

        self.bleu = evaluate.load("bleu")
        self.rouge = evaluate.load("rouge")

    def window_size(self):
        return 0.5

    def stopping_criteria(self, context_length: int = 0, total_length: int = 0):
        return StoppingCriteriaList([
            MaxLengthCriteria(max_length=total_length),
        ])

    def eval_generation(self, output_df=None, label_df=None, output_tokens=None, label_tokens=None):
        predictions_by_field = defaultdict(list)
        label_by_field = defaultdict(list)
        fn = len(output_df.columns)
        for i in range(fn):
            # Get the strided output and label
            max_idx = (output_df.shape[0] - 1) * output_df.shape[1] + i + 1
            field = output_df.columns[i]
            predictions_by_field[field] = [str(x) for x in output_tokens[i:max_idx:fn].tolist()]
            label_by_field[field] = [str(x) for x in label_tokens[i:max_idx:fn].tolist()]
            rouge_field = self.rouge.compute(predictions=predictions_by_field[field], references=label_by_field[field], rouge_types=["rouge1", "rougeL"], use_aggregator=True)
            self.rouge_L_scores_by_field[field].append(
                rouge_field["rougeL"]
            )
            self.rouge_scores_by_field[field].append(
                rouge_field["rouge1"]
            )
        pred_all = [str(x) for x in output_tokens.tolist()]
        label_all = [str(x) for x in label_tokens[:len(pred_all)].tolist()]
        rouge_all = self.rouge.compute(predictions=pred_all, references=label_all, rouge_types=["rouge1", "rougeL"], use_aggregator=True)
        self.rouge_L_scores_by_field["All"].append(
            rouge_all["rougeL"]
        )
        self.rouge_scores_by_field["All"].append(
            rouge_all["rouge1"]
        )
        self.total_seen += 1

    def examples_seen(self):
        return self.total_seen

    def on_finish(self):
        model_type = self.model.replace(os.sep, "_")
        model_path = os.path.join(
            self._exp_cache_path(), f"{model_type}.pkl"
        )
        os.makedirs(os.path.dirname(model_path), exist_ok=True)
        with open(model_path, "wb") as f:
            pickle.dump(
                {
                    "total_seen": self.total_seen,
                    "rouge_scores_by_field": self.rouge_scores_by_field,
                    "rouge_L_scores_by_field": self.rouge_L_scores_by_field,
                },
                f,
            )

    def plot(self):
        model_versions = os.listdir(self._exp_cache_path())

        rows = []
        for model_version in model_versions:
            model_path = os.path.join(self._exp_cache_path(), model_version)
            with open(model_path, "rb") as f:
                model_data = pickle.load(f)
            model_type = model_version.replace(".pkl", "").split("_")
            model_type = f"{model_type[0]}-{model_type[1]}.{model_type[2]}"
            rows.append({
                ("", "Model"): model_type,
                ("ROUGE-1", "METRIC_NAME"): np.mean(model_data["rouge_scores_by_field"]["METRIC_NAME"]),
                ("ROUGE-1", "PAT_ID"): np.mean(model_data["rouge_scores_by_field"]["PAT_ID"]),
                ("ROUGE-1", "ACCESS_TIME"): np.mean(model_data["rouge_scores_by_field"]["ACCESS_TIME"]),
                ("ROUGE-1", "All"): np.mean(model_data["rouge_scores_by_field"]["All"]),
                ("ROUGE-L", "METRIC_NAME"): np.mean(model_data["rouge_L_scores_by_field"]["METRIC_NAME"]),
                ("ROUGE-L", "PAT_ID"): np.mean(model_data["rouge_L_scores_by_field"]["PAT_ID"]),
                ("ROUGE-L", "ACCESS_TIME"): np.mean(model_data["rouge_L_scores_by_field"]["ACCESS_TIME"]),
                ("ROUGE-L", "All"): np.mean(model_data["rouge_L_scores_by_field"]["All"]),
            })
        df = pd.DataFrame(rows, columns=rows[0].keys())
        df = df.sort_values(by=[("", "Model")], ascending=False)
        out = df.to_latex(index=False,
                    float_format="%.3f",
                    caption="ROUGE-1 and ROUGE-L scores for the generated rows.",
                    label="tab:next_action_results")
        print(out)

if __name__ == "__main__":
    # Get arguments
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--model", type=int, default=None, help="Model to use for pretraining."
    )
    parser.add_argument(
        "--val",
        action="store_true",
        help="Run with the validation dataset instead of the test.",
    )
    parser.add_argument(
        "--start",
        type=int,
        default=0,
        help="Index of the first audit log to use for the demo.",
    )
    parser.add_argument(
        "--count",
        type=int,
        default=0,
        help="Number of audit logs to use for the demo.",
    )
    parser.add_argument(
        "--search",
        "-s",
        type=str,
        default="greedy",
        help="Search method to use for decoding. If beam, :k is the beam size.",
    )
    parser.add_argument(
        "-p",
        "--plot",
        type=str,
        default="no", # other options: "yes", "only"
    )
    parser.add_argument(
        "--exp",
        type=str,
        default="NextActionExperiment",
        help="Experiment to run.",
    )
    args = parser.parse_args()
    # Get the list of models from the config file
    config_path = os.path.normpath(
        os.path.join(os.path.dirname(__file__), "config.yaml")
    )
    with open(config_path, "r") as f:
        config = yaml.safe_load(f)

    path_prefix = ""
    for prefix in config["path_prefix"]:
        if os.path.exists(prefix):
            path_prefix = prefix
            break

    if path_prefix == "":
        raise RuntimeError("No valid drive mounted.")

    model_paths = os.path.normpath(
        os.path.join(path_prefix, config["pretrained_model_path"])
    )
    # Get recursive list of subdirectories
    model_list = []
    for root, dirs, files in os.walk(model_paths):
        # If there's a .bin file, it's a model
        if any([file.endswith(".bin") for file in files]):
            # Append the last three directories to the model list
            model_list.append(os.path.join(*root.split(os.sep)[-3:]))

    if len(model_list) == 0:
        raise ValueError(f"No models found in {format(model_paths)}")

    model_list = sorted(model_list)

    if args.model is None:
        print("Select a model to evaluate:")
        for i, model in enumerate(model_list):
            print(f"{i}: {model}")

        model_idx = int(input("Model index >>>"))
    else:
        model_idx = args.model

    model_name = model_list[model_idx]
    model_path = os.path.normpath(
        os.path.join(path_prefix, config["pretrained_model_path"], model_name)
    )

    # Get the device to use
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # Load the test dataset
    vocab = EHRVocab(
        vocab_path=os.path.normpath(os.path.join(path_prefix, config["vocab_path"]))
    )
    types = {
        "gpt2": EHRAuditGPT2,
        "rwkv": EHRAuditRWKV,
        "llama": EHRAuditLlama,
    }

    model_type = model_list[model_idx].split(os.sep)[0]
    model = types[model_type].from_pretrained(model_path, vocab=vocab)
    model.to(device)

    if isinstance(model, EHRAuditGPT2):
        model.generation_config.pad_token_id = model.generation_config.eos_token_id
    elif isinstance(model, EHRAuditLlama):
        model.generation_config.eos_token_id = 0
        model.generation_config.pad_token_id = model.generation_config.eos_token_id
    elif isinstance(model, EHRAuditRWKV):
        model.generation_config.eos_token_id = 0
        model.generation_config.pad_token_id = model.generation_config.eos_token_id

    dm = EHRAuditDataModule(
        yaml_config_path=config_path,
        vocab=vocab,
        batch_size=1,  # Just one sample at a time
    )
    dm.setup()
    if args.val:
        dl = dm.val_dataloader()
    else:
        dl = dm.test_dataloader()

    tk = EHRAuditTokenizer(vocab, timestamp_spaces_cal=timestamp_space_calculation([
                    config["timestamp_bins"]["spacing"],
                    config["timestamp_bins"]["min"],
                    config["timestamp_bins"]["max"],
                    config["timestamp_bins"]["bins"],
                ]))

    exp = eval(args.exp)(vocab=vocab, model=model_name, config=config, path_prefix=path_prefix)
    if args.plot == "only":
        exp.plot()
        sys.exit(0)

    row_len = len(vocab.field_ids) - 1  # Exclude special fields

    max_samples = min(len(dl), args.count) if args.count > 0 else len(dl)
    pbar = tqdm(total=max_samples, desc=f"Examples Seen For {exp.__class__.__name__}")
    for idx, batch in enumerate(dl):
        if idx < args.start:
            continue

        input_ids, labels = batch
        with torch.no_grad():
            # Find the eos index
            nonzeros = (labels.view(-1) == -100).nonzero(as_tuple=True)
            if len(nonzeros[0]) == 0:
                eos_index = len(labels.view(-1)) - 1
            else:
                eos_index = nonzeros[0][0].item() - 1

            row_count = eos_index // row_len
            min_rows = exp.window_size()
            if 0 < min_rows <= 1:
                min_rows = int(1 / min_rows)

            if row_count <= min_rows or row_count <= 2:  # Not applicable
                continue

            for i in range(min_rows, row_count - 1):
                window_size = i * row_len

                # Copy the labels, and zero out past the window length.
                input_ids_c = torch.zeros((input_ids.shape[0], window_size)).long()
                input_ids_c[:, :window_size] = input_ids[:, :window_size]

                # Stack them to the beam size

                # Stopping criteria
                sc_list = exp.stopping_criteria(context_length=window_size, total_length=row_count * row_len)

                # Process logits so that only tokens in the correct field are allowed
                logits_processor = LogitsProcessorList(
                    [
                        EHRAuditLogitsProcessor(vocab=vocab),
                    ]
                )

                # Generate the outputs from the window.
                if args.search == "greedy":
                    outputs = model.greedy_search(
                        input_ids_c.to(device),
                        stopping_criteria=sc_list,
                        logits_processor=logits_processor,
                    )
                elif args.search == "sample":
                    outputs = model.sample(
                        input_ids_c.to(device),
                        stopping_criteria=sc_list,
                        logits_processor=logits_processor,
                    )
                elif "contrastive" in args.search:
                    opts = args.search.split(":")
                    if len(opts) == 1:
                        opts.append("5")
                    if len(opts) <= 2:
                        opts.append("0.1")
                    outputs = model.contrastive_search(
                        input_ids_c.to(device),
                        stopping_criteria=sc_list,
                        logits_processor=logits_processor,
                        top_k=int(opts[1]),
                        penalty_alpha=float(opts[2]),
                    )
                elif "beam" in args.search:
                    # Errors out with CUBLAS_STATUS_NOT_INITIALIZED, not sure why
                    opts = args.search.split(":")
                    if len(opts) == 1:
                        opts.append("5")
                    beam_size = int(opts[1])
                    input_ids_c = input_ids_c.repeat(beam_size, input_ids.size(1)).to(device)

                    beam_scorer = BeamSearchScorer(
                        batch_size=1,
                        num_beams=beam_size,
                        device=device,
                    )

                    outputs = model.beam_search(
                        input_ids_c.to(device),
                        stopping_criteria=sc_list,
                        logits_processor=logits_processor,
                        beam_scorer=beam_scorer,
                    )

                # Decode the outputs
                output_tokens = outputs[0, window_size:].cpu().numpy()
                predictions: DataFrame = tk.decode(output_tokens)
                label_tokens = input_ids[0, window_size:].cpu().numpy()
                labels: DataFrame = tk.decode(label_tokens)

                exp.eval_generation(output_df=predictions, label_df=labels, output_tokens=output_tokens, label_tokens=label_tokens)
                pbar.n = exp.examples_seen()
                pbar.refresh()
                if isinstance(exp, ScoringExperiment):
                    # Only score one batch at a time
                    break
                """
                # Change the indicies of predictions and full_input to a multi-index
                full_input.columns = pd.MultiIndex.from_tuples(
                    [("Ground-Truth", x) for x in full_input.columns]
                )
                predictions.columns = pd.MultiIndex.from_tuples(
                    [("Prediction", x) for x in predictions.columns]
                )

                # Concatenate the two dataframes so that they're side by side
                full_df = pd.concat([full_input, predictions], axis=1)

                # Add a column for displaying ground-truth or predictions.
                full_df["Type"] = ["Input"] * (window_size // row_len) + ["Prediction"] * args.predict_size

                with pd.option_context("display.max_columns", None, "display.max_rows", None, "display.width", None):
                    print(f"==== Predictions for Example {idx} ====")
                    print(full_df)
                """
        if exp.examples_seen() >= max_samples:
            break

    pbar.close()
    exp.on_finish()
    if args.plot == "yes":
        exp.plot()


from sys import stdin

dy = [[-1, 1, -2, 2, 0, -1, 1, -1, 1],
      [0, 0, 1, 1, 3, 1, 1, 2, 2],
      [-1, 1, -2, 2, 0, -1, 1, -1, 1],
      [0, 0, -1, -1, -3, -1, -1, -2, -2]
      ]

dx = [[0, 0, -1, -1, -3, -1, -1, -2, -2],
      [-1, 1, -2, 2, 0, -1, 1, -1, 1],
      [0, 0, 1, 1, 3, 1, 1, 2, 2],
      [-1, 1, -2, 2, 0, -1, 1, -1, 1]
      ]

per = [0.01, 0.01, 0.02, 0.02, 0.05, 0.07, 0.07, 0.1, 0.1]

di = [[0, -1], [1, 0], [0, 1], [-1, 0]]
cd = 0

sand_out = 0

N = int(stdin.readline().strip())
Map = [list(map(int, stdin.readline().split())) for _ in range(N)]
check = [[False for _ in range(N)] for _ in range(N)]

SP = [N//2, N//2]
check[SP[0]][SP[1]] = True

def tornado():
    global cd, check, sand_out, SP, Map
    y, x = SP
    move_ = 0

    while True:
        move_ += 1
        ny = y + di[cd][0] * move_
        nx = x + di[cd][1] * move_

        by = ny - di[cd][0]
        bx = nx - di[cd][1]

        nny = y + di[cd][0] * (move_ + 1)
        nnx = x + di[cd][1] * (move_ + 1)

        # debug
        # print(move_, ny, by, nx, bx, SP)
        # for t in range(N):
        #     print(Map[t])

        if 0 <= ny < N and 0 <= nx < N:
            check[ny][nx] = True
        else:
            break

        cur_sand = Map[ny][nx]
        total_sand = Map[ny][nx]
        SP = [ny, nx]

        #print(sand_out, cur_sand, total_sand)

        Map[ny][nx] = 0

        for i in range(9):
            to_y = by + dy[cd][i]
            to_x = bx + dx[cd][i]

            to_sand = int(total_sand * per[i])

            if to_sand > 0:
                if 0 <= to_y < N and 0 <= to_x < N:
                    #print(to_sand, i, to_y, to_x)
                    Map[to_y][to_x] += to_sand
                    cur_sand -= to_sand
                else:
                    sand_out += to_sand
                    cur_sand -= to_sand

        if 0 <= nny < N and 0 <= nnx < N:
            Map[nny][nnx] += cur_sand

        else:
            sand_out += cur_sand

        if SP == [0, 0]:
            break

        if check[ny + di[cd-3][0]][nx + di[cd-3][1]] is False:
            cd += 1
            cd %= 4
            break

while not (SP[0] == 0 and SP[1] == 0):
    tornado()

#print(Map)
print(sand_out)

6.연습/마법사상어와토네이도.py

마법사상어와토네이도.py


#!/usr/bin/env python

from math import sin, cos, atan2, sqrt, pi
import json
import urllib
import requests
import time
from flask import Flask, request, make_response

# server settings
HOST = '127.0.0.1'
PORT = 5000
DEBUG_MODE = False

# conversion factor from degrees to radians
DEGREES2RADIANS = pi/180.0

# conversion factor from kilometres to miles
KM2MILES = 6.21371192e-1

# conversion factor from kilometres to yards
KM2YARDS = 1.0936133e3

# mean radius of the Earth [km]
RADIUS_EARTH = 6.371e3

# location of White Bear Yard, 144a Clerkenwell Road, London, EC1R 5DF, UK.
WHITE_BEAR_YARD = {'lat': 51.52219930, 'lng': -0.10976180}

def geodesic_distance(point_A, point_B):
    """Geodesic distance between two points on the Earth."""
    lat_A, lng_A = point_A['lat']*DEGREES2RADIANS, point_A['lng']*DEGREES2RADIANS
    lat_B, lng_B = point_B['lat']*DEGREES2RADIANS, point_B['lng']*DEGREES2RADIANS
    # Vincenty's formula eg. http://en.wikipedia.org/wiki/Great-circle_distance
    dl = abs(lng_B - lng_A)
    a = cos(lat_B)*sin(dl)
    b = cos(lat_A)*sin(lat_B)-sin(lat_A)*cos(lat_B)*cos(dl)
    c = sin(lat_A)*sin(lat_B)+cos(lat_A)*cos(lat_B)*cos(dl)
    return RADIUS_EARTH*atan2(sqrt(a*a+b*b),c)

def geocode_address(address):
    """Latitude and longitude of specified address using Google Maps Geocoding API."""
    url = 'http://maps.googleapis.com/maps/api/geocode/json?'
    parameters = {'address': address, 'sensor': 'false'}
    response = requests.get(url + urllib.urlencode(parameters))
    if (response.status_code==200):
        return json.loads(response.content)
    return None

def get_distance(lat,lng,units=None):
    """Return distance between specified position and White Bear Yard in specified units."""
    distance = geodesic_distance(WHITE_BEAR_YARD,{'lat': lat, 'lng': lng})
    if units=='m':
        return distance*1.0e3
    elif units=='yards':
        return distance*KM2YARDS
    elif units=='miles':
        return distance*KM2MILES
    return distance

def parse_query(args):
  """Parse request query and return parameters."""
  query = args.get('q')
  units = args.get('u')
  return (query, units)

def create_response(code,message,lat,lng,address,distance,units):
    """Create response using specified parameters."""
    return {'status': {'code': code,
                       'message': message},
            'result': {'lat': lat,
                       'lng': lng,
                       'address': address,
                       'distance': distance,
                       'units': units,
                       'elapsed': 0.0}}

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

app = Flask(__name__)

@app.route("/geodesics")
def geodesics():
    t_start = time.time()
    query, units = parse_query(request.args)
    if query:
        position = geocode_address(query)
        if (position) and (len(position[u'results'])>0):
            lat = position[u'results'][0][u'geometry'][u'location'][u'lat']
            lng = position[u'results'][0][u'geometry'][u'location'][u'lng']
            response = create_response(200,
                                       'Successful query',
                                       lat,
                                       lng,
                                       position[u'results'][0][u'formatted_address'],
                                       get_distance(lat,lng,units),
                                       units if units else 'km')
        else:
            response = create_response(400,
                                       'Badly formed query',
                                       '',
                                       '',
                                       '',
                                       '',
                                       '')
    else:
        response = create_response(204,
                                   'No valid query made',
                                   WHITE_BEAR_YARD['lat'],
                                   WHITE_BEAR_YARD['lng'],
                                   '144A Clerkenwell Rd, London, Greater London EC1R 5DF, UK',
                                   0.0,
                                   'km')
    response['result']['elapsed'] = (time.time() - t_start)*1.0e3
    http_response = make_response(json.dumps(response))
    http_response.headers['Content-type'] = 'application/json'
    http_response.headers['Access-Control-Allow-Origin'] = '*'
    return http_response

@app.errorhandler(404)
def page_not_found(error):
    t_start = time.time()
    response = create_response(404,
                               'No such endpoint',
                               '',
                               '',
                               '',
                               '',
                               '')
    response['result']['elapsed'] = (time.time() - t_start)*1.0e3
    http_response = make_response((json.dumps(response),404))
    http_response.headers['Content-type'] = 'application/json'
    http_response.headers['Access-Control-Allow-Origin'] = '*'
    return http_response

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

if __name__=="__main__":
  app.run(debug=DEBUG_MODE, host=HOST, port=PORT)


v=int(input())
vote=list(input())
A,B=0,0
for i in range (v):
    if vote[i]=='A':
        A+=1
    else:
        B+=1

if A>B:
    print('A')
elif B>A:
    print('B')
else:
    print("Tie")


from concurrent.futures import ThreadPoolExecutor

from aloud import models
from aloud.console import console
from aloud.openai import oai

def to_audio(
    speakable: str,
    voice: models.Voice = models.Voice.default,
    voice_model: models.VoiceModel = models.VoiceModel.default,
    voice_response_format: models.VoiceResponseFormat = models.VoiceResponseFormat.default,
) -> bytes:
    chunk_size = 4000  # 4096 is the max, but we need to leave some space for the joined newlines
    article_chunks = [speakable[i : i + chunk_size] for i in range(0, len(speakable), chunk_size)]
    with console.status(f'Converting speakable to audio with {voice}@{voice_model}...'):
        with ThreadPoolExecutor() as executor:
            # noinspection PyTypeChecker
            futures = [
                executor.submit(chunk_to_audio, article_chunk, voice, voice_model, voice_response_format)
                for article_chunk in article_chunks
            ]
        audios = [future.result() for future in futures]

    audio = b''.join(audios)
    console.print('\n[b green]ðŸŽµ Done generating audio!')
    return audio

# noinspection PyShadowingBuiltins
def chunk_to_audio(
    input: str,
    voice: models.Voice = models.Voice.default,
    voice_model: models.VoiceModel = models.VoiceModel.default,
    voice_response_format: models.VoiceResponseFormat = models.VoiceResponseFormat.default,
) -> bytes:
    return oai.audio.speech.create(
        input=input,
        model=voice_model,
        voice=voice,
        response_format=voice_response_format,
    ).content


"""An object for containing the response pagination information."""

from marshmallow import fields

from ...models.base import BaseModel, BaseModelSchema

class Paginated(BaseModel):
    """Class representing the pagination details of a response."""

    class Meta:
        """Paginated metadata."""

        schema_class = "PaginatedSchema"

    def __init__(
        self,
        *,
        start: int = None,
        end: int = None,
        limit: int = None,
        total: int = None,
        **kwargs
    ):
        """
        Initialize a Paginated instance.

        Args:
            start: The first record offset
            end: The last record offset
            limit: Enforced limit on the number of records
            total: Total number of records available

        """
        super(Paginated, self).__init__(**kwargs)
        self.start = start
        self.end = end
        self.limit = limit
        self.total = total

class PaginatedSchema(BaseModelSchema):
    """Paginated schema."""

    class Meta:
        """PaginatedSchema metadata."""

        model_class = "Paginated"

    start = fields.Int(required=False)
    end = fields.Int(required=False)
    limit = fields.Int(required=False)
    total = fields.Int(required=False)


import json
import os
import unittest

from unittest.mock import MagicMock

from pywink.api import get_devices_from_response_dict
from pywink.devices import types as device_types

class HubTests(unittest.TestCase):

    def setUp(self):
        super(HubTests, self).setUp()
        self.api_interface = MagicMock()
        all_devices = os.listdir('{}/api_responses/'.format(os.path.dirname(__file__)))
        self.response_dict = {}
        device_list = []
        for json_file in all_devices:
            if os.path.isfile('{}/api_responses/{}'.format(os.path.dirname(__file__), json_file)):
                _json_file = open('{}/api_responses/{}'.format(os.path.dirname(__file__), json_file))
                device_list.append(json.load(_json_file))
                _json_file.close()
        self.response_dict["data"] = device_list

    def test_unit_should_be_none(self):
        devices = get_devices_from_response_dict(self.response_dict, device_types.HUB)
        for device in devices:
            self.assertIsNone(device.unit())

    def test_kidde_radio_code_should_not_be_none(self):
        devices = get_devices_from_response_dict(self.response_dict, device_types.HUB)
        for device in devices:
            if device.manufacturer_device_model() == "wink_project_one":
                self.assertIsNone(device.kidde_radio_code())
            elif device.manufacturer_device_model() == "philips":
                self.assertIsNone(device.kidde_radio_code())
            else:
                self.assertIsNotNone(device.kidde_radio_code())

    def test_update_needed_is_false(self):
        devices = get_devices_from_response_dict(self.response_dict, device_types.HUB)
        for device in devices:
            self.assertFalse(device.update_needed())

    def test_ip_address_is_not_none(self):
        devices = get_devices_from_response_dict(self.response_dict, device_types.HUB)
        for device in devices:
            if device.manufacturer_device_model() == "philips":
                continue
            else:
                self.assertIsNotNone(device.ip_address())

    def test_firmware_version_is_not_none(self):
        devices = get_devices_from_response_dict(self.response_dict, device_types.HUB)
        for device in devices:
            if device.manufacturer_device_model() == "philips":
                continue
            else:
                self.assertIsNotNone(device.firmware_version())


import sqlite3

conn = sqlite3.connect('count1.sqlite')
cur = conn.cursor()
#execute() just excute one sql command!!!
cur.execute('DROP TABLE IF EXISTS Counts')
cur.execute('CREATE TABLE Counts(org TEXT, count INTEGER)')

fname = input('Enter file name: ')
if (len(fname) < 1): fname = 'mbox-short.txt'
fh = open(fname)
for line in fh:
    if not line.startswith('From: '): continue
    pieces = line.split('@')
#print(pieces[1])
    org = pieces[1]
    #prevent sql injection!!, second one is tuple
    cur.execute('SELECT count FROM Counts WHERE org = ? ', (org,))
    # grab the first one to row
    row = cur.fetchone()
    if row is None:
        cur.execute('INSERT INTO Counts(org, count)VALUES (?, 1)', (org,))
    else:
        cur.execute('UPDATE Counts SET count = count + 1 WHERE org = ?', (org,))
    #write data to disk
conn.commit()
sqlstr = 'SELECT org, count FROM Counts ORDER BY count DESC LIMIT 10'

for row in cur.execute(sqlstr):
    print(str(row[0]), row[1])
cur.close()


numbers = range(2, 7)  # normal call with separate arguments
print(list(numbers)) # [2, 3, 4, 5, 6]
args = [2, 7]
numbers = range(*args)  # call with arguments unpacked from a list
print(numbers)      # [2, 3, 4, 5,6]

countries = ['Finland', 'Sweden', 'Norway', 'Denmark', 'Iceland']
fin, sw, nor, *rest = countries
print(fin, sw, nor, rest)   # Finland Sweden Norway ['Denmark', 'Iceland']
numbers = [1, 2, 3, 4, 5, 6, 7]
one, *middle, last = numbers
print(one, middle, last)      # 1 [2, 3, 4, 5, 6] 7

def unpacking_person_info(name, country, city, age):
    return '{name} lives in {country}, {city}. He is {age} year old.'
dct = {'name': 'Asabeneh', 'country': 'Finland', 'city': 'Helsinki', 'age': 250}
print(unpacking_person_info(**dct))  # Asabeneh lives in Finland, Helsinki. He is 250 years old.

def sum_all(*args):
    s = 0
    for i in args:
        s += i
    return s
print(sum_all(1, 2, 3))             # 6
print(sum_all(1, 2, 3, 4, 5, 6, 7)) # 28

def packing_person_info(**kwargs):
    # check the type of kwargs and it is a dict type
    # print(type(kwargs))
	# Printing dictionary items
    for key in kwargs:
        print("{key} = {kwargs[key]}")
    return kwargs

print(packing_person_info(name="Asabeneh",
      country="Finland", city="Helsinki", age=250))

# Spreading in Python
lst_one = [1, 2, 3]
lst_two = [4, 5, 6, 7]
lst = [0, *lst_one, *lst_two]
print(lst)          # [0, 1, 2, 3, 4, 5, 6, 7]

# Enumerate
for index, i in enumerate(countries):
    print('hi')
    if i == 'Finland':
        print('The country {i} has been found at index {index}')

# ZIP
fruits = ['banana', 'orange', 'mango', 'lemon']
vegetables = ['Tomato', 'Potato', 'Cabbage','Onion', 'Carrot']
fruits_and_veges = []
for f, v in zip(fruits, vegetables):
    fruits_and_veges.append({'fruit':f, 'veg':v})

print(fruits_and_veges)


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

import locale

simulation_mode = False
port_speed = 38400
port_name = ""
port = ""
promode = False
elm = None
log = "ddt"
opt_cfc0 = False
opt_n1c = True
log_all = False
auto_refresh = False
elm_failed = False
# KWP2000 Slow init
opt_si = False
report_data = True
ecus_dir = "ecus/"
graphics_dir = "graphics/"
last_error = ""
main_window = None
ecu_scanner = None
debug = 'DDTDEBUG' in os.environ
cantimeout = 0
refreshrate = 5
mode_edit = False
safe_commands = ["3E", "14", "21", "22", "17", "19", "10"]

def get_last_error():
    global last_error
    err = last_error
    last_error = ""
    return err

def translator(filename):
    try:
        lang, enc = locale.getlocale()
        os.environ['LANG'] = lang
    except:
        # defaul tranlation if err en_US
        os.environ["LANG"] = 'fr_FR'
        pass

    # Set up message catalog access
    t = gettext.translation(filename, 'locale', fallback=True) # not ok in python 3.11.x, codeset="utf-8")
    return t.gettext


# https://leetcode.com/problems/two-sum-ii-input-array-is-sorted/submissions/868161853/
from typing import List

class Solution:
    def twoSum(self, numbers: List[int], target: int) -> List[int]:
        l, r = 0, len(numbers)-1
        while l < r:
            s = numbers[l] + numbers[r]
            if s == target:
                return [l+1, r+1]
            elif s < target:
                l += 1
            else:
                r -= 1

numbers = [2, 7, 11, 15]
target = 26
result = Solution().twoSum(numbers, target)
print(" numbers: ", numbers, "\n target: ", target, "\n result: ", result)