import pathlib

import numpy as np
import pandas as pd

import simscale_eba.ResultProcessing as pp

result = pp.directional_result()
result.find_project("Greenlite UK")
result.find_simulation("Facade Pressure Coefficients")
result.find_run("Run 2")
result.query_results()
options = result.return_result_options()

name = "Face4Floor2"
category = "PROBE_POINT_PLOT_STATISTICAL_DATA"
# result.download_result(category=category, name=name)

download_dict = result.download_dict
items = result._find_item(category, name)

direction = float(270)

for item in items:
    if item.direction == direction:
        result_set = item

url = result_set.download.url
data_responce = result.api_client.rest_client.GET(
    url=url,
    headers={result.api_key_header: result.api_key},
    _preload_content=False
)

probe_point_plot_statistical_data_csv = data_responce.data.decode('utf-8')

path = pathlib.Path.cwd()
output_path = path / "{}_{}_{}.csv".format(category, name, result.name)

with open(output_path, 'w', newline="") as file:
    file.write(probe_point_plot_statistical_data_csv)
file.close()

df = pd.read_csv(output_path)
avg_pressure = dict(list(df.groupby("VARIABLE")))["p"]
avg_pressure.index = avg_pressure["ITEM NAME"]
avg_pressure = avg_pressure["AVG"]
avg_cp = avg_pressure / (0.5 * 1.2 * 10 ** 2)

scenario_pressure = avg_cp * (0.5 * 1.2 * 4.62 ** 2)

class window():

    def __init__(self,
                 name,
                 _type="TOP_HUNG",
                 width=1,
                 height=1,
                 angle=10,
                 density=1.2):

        self.name = name
        self.width = width
        self.height = height
        self.opening_angle = float(angle)
        self.type = _type
        self.rho = density

        self.flow_rate = np.array([0, 0.25, 0.5, 1, 2, 4]).reshape((-1, 1))
        self.dp = None

        self.stroke_length = None
        self.length_ratio = None
        self.gradient = None
        self.max_cd = None
        self.cd0 = 0.62  # Orifice Dichcarge Coefficient
        self.a_eq = None  # equivilent_area
        self.a_eff = None  # effective area
        self.a_free = None  # free area

        self.cd = None
        self.dp = None

        self.df = None

    def set_top_hinged(self):
        self.length_ratio = self.width / self.height
        self.stroke_length = np.sin((self.opening_angle * np.pi) / (2 * 180))
        self.a_free = self.width * self.height

        _dict = {
            "1": {"lower bound": 0, "upper bound": 0.5, "gradient": 0.06, "cdmax": 0.612},
            "2": {"lower bound": 0.5, "upper bound": 1, "gradient": 0.048, "cdmax": 0.589},
            "3": {"lower bound": 1, "upper bound": 2, "gradient": 0.04, "cdmax": 0.563},
            "4": {"lower bound": 2, "upper bound": np.inf, "gradient": 0.038, "cdmax": 0.548},
        }

        for key in _dict.keys():
            if (self.length_ratio >= _dict[key]["lower bound"]
                    and self.length_ratio < _dict[key]["upper bound"]):
                self.gradient = _dict[key]["gradient"]
                self.max_cd = _dict[key]["cdmax"]

        self.cd = self.max_cd * (1 - np.exp(-self.gradient * self.opening_angle))

    def get_pressure_drop(self):
        self.dp = (self.rho / 2) * (self.flow_rate / (self.cd * self.a_free))
        return self.dp

    def to_dataframe(self):
        arr = np.concatenate((self.flow_rate, self.dp), axis=1)
        columns = ["V_dot (m3/s)", "dp (Pa)"]
        df = pd.DataFrame(arr, columns=columns)

        self.df = df

    def to_csv(self, path=pathlib.Path.cwd()):
        self.to_dataframe()
        path = path / (self.name + "{}".format(self.opening_angle) + '.csv')
        self.df.to_csv(path)

win = window("top_hung_1", width=0.9, height=1.6)
win.set_top_hinged()
dp = win.get_pressure_drop()

win.to_csv()


#! /usr/bin/env python
"""Preparation for dedupe"""
import concurrent.futures
import os
import re
import time
import typing
import unicodedata
import urllib.parse
from datetime import datetime
from pathlib import Path

import numpy as np
import pandas as pd
from number_parser import parse

from bib_dedupe.constants import ENTRYTYPES
from bib_dedupe.constants import Fields

# TODO : upworker https://en.wikipedia.org/wiki/Category:Multilingual_journals
# Better: https://www.freecodecamp.org/news/scraping-wikipedia-articles-with-python/
# replace in prep (to ensure blocking! - this would also be more robust when the container-prep changes.)

current_dir = Path(__file__).parent
journal_translations_path = current_dir / "wikipedia_journal_translations.csv"
journal_translations = pd.read_csv(journal_translations_path)
JOURNAL_TRANSLATIONS_DICT = dict(
    zip(journal_translations.iloc[:, 1], journal_translations.iloc[:, 0])
)
JOURNAL_TRANSLATIONS_DICT["Genetika"] = "Russian Journal of Genetics"
JOURNAL_TRANSLATIONS_DICT[
    "Zhejiang da xue xue bao. Yi xue ban"
] = "Journal of Zhejiang University. Medical sciences"
JOURNAL_TRANSLATIONS_DICT[
    "Zhonghua yi xue za zhi"
] = "National Medical Journal of China"
JOURNAL_TRANSLATIONS_DICT[
    "Zhong Nan Da Xue Xue Bao Yi Xue Ban"
] = "Journal of Central South University"
JOURNAL_TRANSLATIONS_DICT[
    "Revista Brasileira de Fisioterapia"
] = "Brazilian Journal of Physical Therapy"
JOURNAL_TRANSLATIONS_DICT["Zhongguo fei ai za zhi"] = "Chinese journal of lung cancer"
# https://portal.issn.org/resource/ISSN/1347-8397
JOURNAL_TRANSLATIONS_DICT["Nihon Yakurigaku zasshi"] = "Folia pharmacologia Japonica"
# https://www.journalguide.com/journals/nihon-shinkei-seishin-yakurigaku-zasshi-japanese-journal-of-psychopharmacology
JOURNAL_TRANSLATIONS_DICT[
    "Nihon Shinkei Seishin Yakurigaku Zasshi"
] = "Japanese Journal of Psychopharmacology"
# https://www.scimagojr.com/journalsearch.php?q=21427&tip=sid&clean=0
JOURNAL_TRANSLATIONS_DICT[
    "Zhejiang da xue xue bao. Yi xue ban"
] = "Journal of Zhejiang University. Medical sciences"
JOURNAL_TRANSLATIONS_DICT[
    "Yakubutsu Seishin Kodo"
] = "Japanese Journal of Psychopharmacology"
JOURNAL_TRANSLATIONS_DICT["Folia pharmacologia Japonica"] = "Nihon Yakurigaku Zasshi"
JOURNAL_TRANSLATIONS_DICT["Sheng li xue bao"] = "Acta physiologica Sinica"
JOURNAL_TRANSLATIONS_DICT[
    "Seishin Shinkeigaku Zasshi"
] = "Psychiatria et Neurologia Japonica"
JOURNAL_TRANSLATIONS_DICT[
    "Guang Pu Xue Yu Guang Pu Fen Xi"
] = "Spectroscopy and Spectral Analysis"
JOURNAL_TRANSLATIONS_DICT[
    "Revista Brasileira de Psiquiatria"
] = "Brazilian Journal of Psychiatry"
JOURNAL_TRANSLATIONS_DICT[
    "Jinko Chino Gakkai Ronbunshi"
] = "Transactions of the Japanese Society for Artificial Intelligence"
JOURNAL_TRANSLATIONS_DICT[
    "Nihon Jibiinkoka Gakkai Kaiho"
] = "Journal of Otolaryngology of Japan"
JOURNAL_TRANSLATIONS_DICT["Aizheng"] = "Chinese Journal of Cancer"
JOURNAL_TRANSLATIONS_DICT["Rinsho Shinkeigaku"] = "Clinical neurology"

JOURNAL_TRANSLATIONS_DICT = {
    k.lower(): v.lower() for k, v in JOURNAL_TRANSLATIONS_DICT.items()
}

# Note: for re.sub, "van der" must be checked before "van"
NAME_PREFIXES_LOWER = [
    "van der",
    "van",
    "von der",
    "von",
    "vom",
    "le",
    "den",
    "der",
    "ter",
    "de",
    "da",
    "di",
]

def set_container_title(records_df: pd.DataFrame) -> None:
    # Set container title
    records_df.loc[
        records_df[Fields.ENTRYTYPE] == ENTRYTYPES.ARTICLE, Fields.CONTAINER_TITLE
    ] = records_df[Fields.JOURNAL]
    records_df.loc[
        records_df[Fields.ENTRYTYPE].isin(
            [ENTRYTYPES.INPROCEEDINGS, ENTRYTYPES.PROCEEDINGS, ENTRYTYPES.INBOOK]
        ),
        Fields.CONTAINER_TITLE,
    ] = records_df[Fields.BOOKTITLE]
    records_df.loc[
        records_df[Fields.ENTRYTYPE] == ENTRYTYPES.BOOK, Fields.CONTAINER_TITLE
    ] = records_df[Fields.TITLE]
    records_df[Fields.CONTAINER_TITLE].fillna("", inplace=True)

def get_author_format_case(authors: typing.List[str], original_string: str) -> str:
    # TODO : non-latin, empty (after replacing special characters)
    if authors == [""]:
        return "empty"

    if any(
        keyword in original_string.lower() for keyword in ["group", "agency", "council"]
    ):
        return "organization"

    if (" and " in original_string and ", " in original_string) or (
        " and " not in original_string
        and ", " in original_string
        and len(original_string) < 50
    ):
        return "proper_format"

    if len(authors) < 4 and not any("," in a for a in authors):
        return "single_author_missing_comma"

    if " and " not in authors and "," not in authors and len(original_string) > 5:
        if (
            sum(
                1
                for author in authors
                if author.isupper() and len(re.sub(r"[ .-]", "", author)) < 3
            )
            / len(authors)
            >= 0.4
        ):
            return "abbreviated_initials"

        if (
            sum(1 for match in re.findall(r"[A-Z][a-z\.]+[A-Z][a-z]+", original_string))
            / len(authors)
            >= 0.1
        ):
            if " " in original_string:
                return "missing_spaces_between_words"
            else:
                return "no_spaces_at_al"

    return "special_case"

def get_authors_split(authors: str) -> list:
    # single author
    if len(authors) < 15:
        if " " not in authors and re.search(r"[a-z]{3}[A-Z]", authors):
            return re.split(r"(?<=[a-z])(?=[A-Z])", authors)

        if authors.count(" ") <= 2:
            return authors.split(" ")

    authors_list = re.split(r"(?=[A-Z])", authors)  # .replace(".", ""))
    for i in range(len(authors_list) - 1):
        if (
            authors_list[i].endswith("-")
            or authors_list[i] in ["Mc", "Mac"]
            or (
                len(authors_list[i]) == 1
                and authors_list[i].isupper()
                and len(authors_list[i + 1]) == 1
                and authors_list[i + 1].isupper()
            )
        ):
            authors_list[i + 1] = authors_list[i] + authors_list[i + 1]
            authors_list[i] = ""
    return [author.rstrip() for author in authors_list if author != ""]

def remove_accents(*, input_str: str) -> str:
    """Replace the accents in a string"""

    def rmdiacritics(char: str) -> str:
        """
        Return the base character of char, by "removing" any
        diacritics like accents or curls and strokes and the like.
        """
        try:
            desc = unicodedata.name(char)
            cutoff = desc.find(" WITH ")
            if cutoff != -1:
                desc = desc[:cutoff]
                char = unicodedata.lookup(desc)
        except (KeyError, ValueError):
            pass  # removing "WITH ..." produced an invalid name
        return char

    try:
        nfkd_form = unicodedata.normalize("NFKD", input_str)
        wo_ac_list = [
            rmdiacritics(c) for c in nfkd_form if not unicodedata.combining(c)
        ]
        wo_ac = "".join(wo_ac_list)
    except ValueError:
        wo_ac = input_str
    return wo_ac

def prep_authors(authors_array: np.array, *, debug: bool = False) -> np.array:
    def preprocess_author(authors: str, *, debug: bool) -> str:
        authors = str(authors)

        if authors in ["Anonymous"]:
            return ""

        # Many databases do not handle accents properly...
        authors = (
            remove_accents(input_str=authors)
            .replace("ue", "u")
            .replace("oe", "o")
            .replace("ae", "a")
        )

        authors = re.sub(r"\d", "", authors)
        if authors.lower() in ["nan", "anonymous"]:
            return ""

        if ";" in authors:
            authors = authors.replace(";", " and ")

        # Capitalize von/van/... and add "-" to connect last-names
        authors = re.sub(
            r"([A-Z])(" + "|".join(NAME_PREFIXES_LOWER) + r") (\w+)",
            lambda match: match.group(1)
            + match.group(2).title().replace(" ", "-")
            + "-"
            + match.group(3),
            authors,
        )
        authors = re.sub(
            r"(^| |\.|-)(" + "|".join(NAME_PREFIXES_LOWER) + r") (\w+)",
            lambda match: match.group(1)
            + match.group(2).title().replace(" ", "-")
            + "-"
            + match.group(3),
            authors,
            flags=re.IGNORECASE,
        )

        original_string = authors

        authors_list = get_authors_split(authors)

        format_case = get_author_format_case(authors_list, original_string)

        if debug:
            print(original_string)
            print(f"format_case: {format_case}")
            print(f"authors_list: {authors_list}")

        if format_case == "proper_format":
            authors_str = authors
        elif format_case == "organization":
            authors_str = authors

        elif format_case == "empty":
            authors_str = ""

        elif format_case == "single_author_missing_comma":
            if authors_list[0].isupper():
                authors_list[0] = authors_list[0].title()

            authors_str = authors_list[0] + ", " + " ".join(authors_list[1:])

        # Broadley K.Burton A. C.Avgar T.Boutin S.
        elif format_case == "abbreviated_initials":
            temp_authors_list = []
            next_author: typing.List[str] = []
            for author_fragment in authors_list:
                if re.match(r"[A-Z][a-z]+", author_fragment):
                    temp_authors_list.append(" ".join(next_author))
                    next_author = [author_fragment]
                else:
                    next_author.append(author_fragment)
            temp_authors_list.append(" ".join(next_author))
            temp_authors_list = [author for author in temp_authors_list if author != ""]

            for i in range(len(temp_authors_list)):
                words = temp_authors_list[i].split()
                for j in range(len(words) - 1, -1, -1):
                    if words[j].isupper() and not words[j - 1].isupper():
                        words[j - 1] = words[j - 1] + ","
                        break
                temp_authors_list[i] = " ".join(words)

            for i in range(len(temp_authors_list)):
                if i == len(temp_authors_list) - 1:
                    temp_authors_list[i] = temp_authors_list[i]
                elif ", " in temp_authors_list[i]:
                    temp_authors_list[i] = temp_authors_list[i] + " and "
                else:
                    temp_authors_list[i] = temp_authors_list[i] + " "

            authors_str = "".join(temp_authors_list)

        # PayenJ.-L.IzopetJ.Galindo-MigeotV.Lauwers-Cances
        elif format_case == "no_spaces_at_al":
            authors_str = ""
            for element in authors_list:
                if re.match(r"^[A-Z][a-z]+", element):
                    authors_str += element + " "
                else:
                    authors_str += ", " + element + " and "
            authors_str = authors_str.rstrip(" and ")

        # Vernia FilippoDi Ruscio MirkoStefanelli GianpieroViscido AngeloFrieri GiuseppeLatella Giovanni
        elif format_case == "missing_spaces_between_words":
            new_authors = re.sub(
                r"([A-Z][a-z.]+)([A-Z])", r"\1 SPLIT\2", original_string
            ).split("SPLIT")
            for i, element in enumerate(new_authors):
                words = [
                    w.replace(".", "").rstrip()
                    for w in element.split()
                    if w.lower() not in NAME_PREFIXES_LOWER
                ]
                if len(words) > 1:
                    middle_index = len(words) // 2
                    words.insert(middle_index, ",")
                    new_authors[i] = " ".join(words)
            authors_str = " and ".join(new_authors)

        else:
            # For debugging:
            # print(f"{format_case}: {original_string}")
            authors_str = " and ".join(authors_list)

        authors_str = authors_str.replace(" ,", ",")
        authors_str = authors_str.replace("anonymous", "").replace("jr", "")
        authors_str = re.sub(r"[^A-Za-z0-9, ]+", "", authors_str)
        return authors_str.lower()

    return np.array(
        [preprocess_author(author, debug=debug) for author in authors_array]
    )

def select_authors(authors_array: np.array) -> np.array:
    def select_author(authors: str) -> str:
        """Select first author"""

        authors_list = authors.split(" and ")
        authors_str = " ".join(
            [
                re.sub(
                    r"(^| )(" + "|".join(NAME_PREFIXES_LOWER) + r") ",
                    lambda match: match.group(1) + match.group(2).replace(" ", "-") + "-",  # type: ignore
                    author.split(",")[0],
                    flags=re.IGNORECASE,
                ).replace(" ", "")
                for author in authors_list
            ][:8]
        )
        authors_str = authors_str.replace("anonymous", "").replace("jr", "")
        authors_str = re.sub(r"[^A-Za-z0-9, ]+", "", authors_str)
        return authors_str

    return np.array([select_author(author) for author in authors_array])

def remove_erratum_suffix(title: str) -> str:
    erratum_phrases = ["erratum appears in ", "erratum in "]
    for phrase in erratum_phrases:
        if phrase in title.lower():
            title = title[: title.lower().rfind(phrase) - 2]

    title = re.sub(r" review \d+ refs$", "", title)

    return title

TITLE_STOPWORDS = [
    "a",
    "an",
    "the",
    "in",
    "of",
    "on",
    "for",
    "from",
    "does",
    "do",
    "and",
    "are",
    "on",
    "with",
    "to",
    "or",
    "as",
    "by",
    "their",
    "the",
]

# flake8: noqa: E501
# pylint: disable=line-too-long
def prep_title(title_array: np.array) -> np.array:
    # Remove language and contents after if "Russian" or "Chinese" followed by a newline
    title_array = np.array(
        [
            re.sub(
                r"\. (Russian|Chinese|Spanish|Czech|Italian|Polish|Dutch|Ukrainian|German|French|Japanese|Slovak|Hungarian|Portuguese English|Turkish|Norwegian|Portuguese)(\r?\n)?.*$",
                "",
                title,
                flags=re.IGNORECASE,
            )
            if ". " in title
            else title
            for title in title_array
        ]
    )

    title_array = np.array(
        [
            title.replace("-like", "like")
            .replace("co-", "co")
            .replace("post-", "post")
            .replace("three-dimensional", "threedimensional")
            for title in title_array
        ]
    )

    # Replace 'withdrawn' at the beginning and '(review') at the end (case insensitive)
    title_array = np.array(
        [
            re.sub(
                r"^(withdrawn[.:] )|^(proceedings: )|^(reprint)|( \(review\))$|( \(vol \d+.*\))",
                "",
                title,
                flags=re.IGNORECASE,
            )
            for title in title_array
        ]
    )

    # Replace roman numbers
    title_array = np.array(
        [
            re.sub(
                r" iv([ |-|.]|$)",
                " 4 ",
                re.sub(
                    r" iii([ |-|.]|$)",
                    " 3 ",
                    re.sub(
                        r" ii([ |-|.]|$)",
                        " 2 ",
                        re.sub(r" i([ |-|.]|$)", " 1 ", title, flags=re.IGNORECASE),
                        flags=re.IGNORECASE,
                    ),
                    flags=re.IGNORECASE,
                ),
                flags=re.IGNORECASE,
            )
            for title in title_array
        ]
    )

    # Remove html tags
    title_array = np.array([re.sub(r"<.*?>", " ", title) for title in title_array])

    # Remove language tags added by some databases (at the end)
    title_array = np.array(
        [re.sub(r"\. \[[A-Z][a-z]*\]$", "", title) for title in title_array]
    )

    # Remove trailing "1" if it is not preceded by "part"
    title_array = np.array(
        [
            re.sub(r"1$", "", title) if "part" not in title[-10:].lower() else title
            for title in title_array
        ]
    )
    # Remove erratum suffix
    # https://www.nlm.nih.gov/bsd/policy/errata.html
    title_array = np.array([remove_erratum_suffix(title) for title in title_array])

    # Replace words in parentheses at the end
    title_array = np.array(
        [re.sub(r"\s*\([^)]*\)\s*$", "", value) for value in title_array]
    )

    # Remove '[Review] [33 refs]' and ' [abstract no: 134]' ignoring case
    title_array = np.array(
        [
            re.sub(
                r"\[Review\] \[\d+ refs\]| \[abstract no: \d+\]",
                "",
                title,
                flags=re.IGNORECASE,
            )
            for title in title_array
        ]
    )

    # Replace brackets (often used in formulae)
    title_array = np.array(
        [re.sub(r"([A-Za-z])\(([0-9]*)\)", r"\1\2", title) for title in title_array]
    )

    # Replace special characters
    title_array = np.array(
        [re.sub(r"[^A-Za-z0-9,\[\]]+", " ", title.lower()) for title in title_array]
    )

    # Remove common stopwords
    title_array = np.array(
        [
            " ".join(word for word in title.split() if word not in TITLE_STOPWORDS)
            for title in title_array
        ]
    )

    # Apply parse function to replace numbers
    title_array = np.array([parse(title) for title in title_array])

    # Replace spaces between digits
    title_array = np.array(
        [
            re.sub(r"(\d) (\d)", r"\1\2", title).rstrip(" ].").lstrip("[ ")
            for title in title_array
        ]
    )

    # Replace multiple spaces with a single space
    title_array = np.array(
        [re.sub(r"\s+", " ", title).rstrip().lstrip() for title in title_array]
    )
    return title_array

JOURNAL_STOPWORDS = [
    "of",
    "for",
    "the",
    "and",
    "de",
    "d",
    "et",
    "in",
    "i",
    "&",
    "to",
    "on",
    "die",
    "part",
    "annual",
    "und",
    "für",
    "fur",
    "der",
]

JOURNAL_ABBREV = {
    "amer": "am",
    "soci": "soc",
    "expe": "exp",
    "mole": "mol",
    "scie": "sci",
    "brit": "br",
    "bole": "bol",
    "inte": "int",
    "arqu": "arq",
    "polo": "pol",
    "vete": "vet",
    "desi": "des",
    "mede": "med",
    "tera": "ter",
    "huma": "hum",
    "revu": "rev",
    "natu": "nat",
    "move": "mov",
    "cana": "can",
    "euro": "eur",
    "adva": "adv",
    "medi": "med",
    "anna": "ann",
    "revi": "rev",
    "rese": "res",
}

def prep_container_title(ct_array: np.array) -> np.array:
    def get_abbrev(ct: str) -> str:
        # Use abbreviated versions
        # journal of infection and chemotherapy
        # j infect chemother
        ct = str(ct)

        # Replace trailing "the" (ignore case)
        ct = re.sub(r"\sthe\s*$", "", ct, flags=re.IGNORECASE)

        # Replace translated journals
        ct_lower = ct.lower().rstrip()
        if ct_lower in JOURNAL_TRANSLATIONS_DICT:
            ct = JOURNAL_TRANSLATIONS_DICT[ct_lower]

        ct = (
            ct.lower()
            .replace(" neuro ", " neuro")
            .replace("-", "")
            .replace("journal", "j")
        )

        if "plos one" not in ct:
            ct = parse(ct)  # replace numbers

        ct = " ".join(word[:4] for word in ct.split() if word not in JOURNAL_STOPWORDS)
        # use the same abbreviations (container_title is used for blocking)
        for original, replacement in JOURNAL_ABBREV.items():
            ct = ct.replace(original, replacement)

        # Replace trailing "supp" or leading "proc"
        ct = re.sub(r"^proc\s|\ssupp$", "", ct)
        return ct

    ct_array = np.array(
        [
            re.sub("proceedings of the", "", value.split(".")[0], flags=re.IGNORECASE)
            if "date of publication" in value.lower()
            or "conference start" in value.lower()
            else re.sub("proceedings of the", "", value, flags=re.IGNORECASE)
            for value in ct_array
        ]
    )

    ct_array = np.array(
        [
            # |(\. )
            re.split(r"(\.\d+)|(: )|( - )", value)[0]
            if re.search(r"(\.\d+)|(\. )|(: )|( - )", value)
            else value
            for value in ct_array
        ]
    )

    # If spaces were accidentally removed, restore them (looking for capital letters)
    # ct_array = np.array(
    #     [
    #         re.sub(r"(?<=\w)([A-Z])", r" \1", value)
    #         if "plos" not in value.lower()
    #         else value
    #         for value in ct_array
    #     ]
    # )
    # Note : distinguish confernces before (see digital_work)

    # Replace words in parentheses at the end
    ct_array = np.array(
        [re.sub(r"\s*\([^)]*\)\s*$|('s)", "", value) for value in ct_array]
    )

    # ct_array = np.array([value.replace(".", " ") for value in ct_array])
    ct_array = np.array(
        [re.sub(r"^the\s|^(l')|", "", value, flags=re.IGNORECASE) for value in ct_array]
    )
    ct_array = np.array([re.sub(r"[^A-Za-z ]+", " ", value) for value in ct_array])
    ct_array = np.array(
        [
            re.sub(r"^\s*(st|nd|rd|th) ", "", value, flags=re.IGNORECASE)
            for value in ct_array
        ]
    )

    return np.array([get_abbrev(ct) for ct in ct_array])

def prep_year(year_array: np.array) -> np.array:
    def process_year(value: str) -> str:
        try:
            int_value = int(float(value))

            if not 1900 < int_value < 2100:
                return ""

        except ValueError:
            return ""
        return str(int_value)

    return np.array([process_year(year) for year in year_array])

def prep_volume(volume_array: np.array) -> np.array:
    volume_array = np.array(
        [
            re.search(r"(\d+) \(.*\)", volume).group(1)  # type: ignore
            if re.search(r"(\d+) \(.*\)", volume) is not None
            # pages included: "6 51-6"
            else re.search(r"(\d+) .*", volume).group(1)  # type: ignore
            if re.search(r"(\d+) \d+-\d+", volume) is not None
            else volume
            for volume in volume_array
        ]
    )
    volume_array = np.array(
        [
            re.search(r"(\d+) suppl \d+", volume.lower()).group(1)  # type: ignore
            if re.search(r"(\d+) suppl \d+", volume.lower()) is not None
            else volume
            for volume in volume_array
        ]
    )
    volume_array = np.array(
        [
            re.sub(r"[^\d\(\)]", "", volume)
            # if len(volume) < 4 and any(char.isdigit() for char in volume)
            # else volume
            for volume in volume_array
        ]
    )

    volume_array = np.array(
        [
            re.search(r"(\d+)", volume).group(0)  # type: ignore
            if re.search(r"(\d+)", volume) is not None
            else volume.replace("(", "").replace(")", "")
            for volume in volume_array
        ]
    )
    return np.array(
        [
            "" if volume == "nan" or len(volume) > 100 else volume
            for volume in volume_array
        ]
    )

def prep_number(number_array: np.array) -> np.array:
    number_array = np.array(
        [re.sub(r"[A-Za-z.]*", "", number) for number in number_array]
    )
    number_array = np.array(
        [
            # pages included: "6 51-6"
            re.search(r"(\d+) .*", number).group(1)  # type: ignore
            if re.search(r"(\d+) \d+-\d+", number) is not None
            else number
            for number in number_array
        ]
    )

    number_array = np.array(
        [
            number.replace(" ", "").replace("(", "").replace(")", "")
            for number in number_array
        ]
    )

    return np.array(
        ["" if number in ["nan", "var.pagings"] else number for number in number_array]
    )

month_dict = {
    "jan": "01",
    "feb": "02",
    "mar": "03",
    "apr": "04",
    "may": "05",
    "jun": "06",
    "jul": "07",
    "aug": "08",
    "sep": "09",
    "oct": "10",
    "nov": "11",
    "dec": "12",
}

def prep_pages(pages_array: np.array) -> np.array:
    def process_page(value: str) -> str:
        if value.isalpha():
            return ""

        # Fix Excel errors (e.g., "08-11" is converted to "08-Nov")
        for month, num in month_dict.items():
            if month in value.lower():
                value = value.lower().replace(month, num)
                break

        def roman_to_int(s: str) -> int:
            rom_val = {"I": 1, "V": 5, "X": 10, "L": 50, "C": 100, "D": 500, "M": 1000}
            int_val = 0
            for i in range(len(s)):
                if i > 0 and rom_val[s[i]] > rom_val[s[i - 1]]:
                    int_val += rom_val[s[i]] - 2 * rom_val[s[i - 1]]
                else:
                    int_val += rom_val[s[i]]
            return int_val

        # Check if the value is a roman numeral range
        roman_numeral_range = re.match(r"([IVXLCDM]+)-([IVXLCDM]+)", value, re.I)
        if roman_numeral_range:
            # Convert the roman numerals to integers
            from_page = roman_to_int(roman_numeral_range.group(1).upper())
            to_page = roman_to_int(roman_numeral_range.group(2).upper())
            return f"{from_page}-{to_page}"

        # Remove leading zeros in groups of digits
        value = re.sub(r"\b0+([0-9]+)", r"\1", value)

        value = re.sub(r"[A-Za-z. ]*", "", value)
        if re.match(r"^\d+\s*-?-\s*\d+$", value):
            from_page_str, to_page_str = re.findall(r"(\d+)", value)
            if from_page_str == to_page_str:
                return from_page_str
            if len(from_page_str) > len(to_page_str):
                return (
                    f"{from_page_str}-{from_page_str[:-len(to_page_str)]}{to_page_str}"
                )
            else:
                return f"{from_page_str}-{to_page_str}"
        if value in [
            " ",
            None,
            "nan",
            "na",
            "no pages",
            "no pagination",
            "var.pagings",
        ]:
            return ""
        # re sub everything except numbers and dashes
        return re.sub(r"[^0-9-]", "", value).lstrip("-").rstrip("-")

    return np.vectorize(process_page)(pages_array)

def prep_abstract(abstract_array: np.array) -> np.array:
    abstract_array = np.array(
        [re.sub(r"<.*?>", " ", abstract.lower()) for abstract in abstract_array]
    )

    abstract_array = np.array(
        [
            re.sub(r"^aims\s|^objectives|^background", "", abstract)
            for abstract in abstract_array
        ]
    )

    abstract_array = np.array(
        [
            abstract[: abstract.rfind(". copyright")]
            if ". copyright" in abstract[-300:]
            else abstract[: abstract.rfind("©")]
            if "©" in abstract[-200:]
            else re.sub(r"(\s*\d{4}\s*)?the authors[.?]$", "", abstract)
            if "the authors" in abstract[-100:]
            else abstract[: abstract.rfind("springer-verlag")]
            if "springer-verlag" in abstract[-100:]
            else re.sub(r"\s*\d{4}.*$", "", abstract)
            if re.search(r"\.\s*\d{4}.*$", abstract)
            else re.sub(r" \(c\) \d{4}.*\.$", "", abstract)
            if re.search(r"\. \(c\) \d{4}.*\.$", abstract)
            else re.sub(r"\.\(abstract truncated at 400 words\)$", "", abstract)
            if ".(abstract truncated at 400 words)" in abstract[-80:]
            else abstract
            for abstract in abstract_array
        ]
    )

    # Remove "abstract " at the beginning
    abstract_array = np.array(
        [re.sub(r"^abstract ", "", abstract) for abstract in abstract_array]
    )

    # Remove trailing date
    abstract_array = np.array(
        [re.sub(r"\s*\(\d{4}\)$", "", abstract) for abstract in abstract_array]
    )

    abstract_array = np.array(
        [re.sub(r"[^A-Za-z0-9 .,]", "", abstract) for abstract in abstract_array]
    )
    abstract_array = np.array(
        [re.sub(r"\s+", " ", abstract) for abstract in abstract_array]
    )
    return np.array(
        [
            "" if abstract == "nan" else abstract.lower().rstrip(" .").lstrip(" .")
            for abstract in abstract_array
        ]
    )

def prep_doi(doi_array: np.array) -> np.array:
    doi_array = np.array(
        [re.sub(r"http://dx.doi.org/", "", doi.lower()) for doi in doi_array]
    )
    doi_array = np.array([re.sub(r"-", "_", doi) for doi in doi_array])
    doi_array = np.array([re.sub(r"\[doi\]", "", doi) for doi in doi_array])
    doi_array = np.array([re.sub(r"[\r\n]+", " ; ", doi) for doi in doi_array])

    doi_array = np.array(
        [
            doi.split(";")[1].lstrip()
            if ";" in doi and doi.split(";")[1].lstrip().startswith("10.")
            else doi.split(";")[0].lstrip()
            if ";" in doi and doi.split(";")[0].lstrip().startswith("10.")
            else doi
            for doi in doi_array
        ]
    )
    doi_array = np.array(
        [doi.split("[pii];")[1] if "[pii];" in doi else doi for doi in doi_array]
    )
    doi_array = np.array([urllib.parse.unquote(doi) for doi in doi_array])
    doi_array = np.array([doi if doi.startswith("10.") else "" for doi in doi_array])

    return np.array(
        ["" if doi == "nan" else doi.replace(".", "").rstrip() for doi in doi_array]
    )

def prep_isbn(isbn_array: np.array) -> np.array:
    isbn_array = np.array([re.sub(r"[\r\n]+", ";", isbn) for isbn in isbn_array])
    isbn_array = np.array(
        [
            re.search(r"\b\d{4}-?\d{3}(\d|X)\b", isbn).group()  # type: ignore
            if re.search(r"\b\d{4}-?\d{3}(\d|X)\b", isbn)
            else isbn
            for isbn in isbn_array
        ]
    )
    return np.array(["" if isbn == "nan" else isbn.lower() for isbn in isbn_array])

function_mapping = {
    Fields.AUTHOR: prep_authors,
    Fields.TITLE: prep_title,
    Fields.CONTAINER_TITLE: prep_container_title,
    Fields.YEAR: prep_year,
    Fields.VOLUME: prep_volume,
    Fields.NUMBER: prep_number,
    Fields.PAGES: prep_pages,
    Fields.ABSTRACT: prep_abstract,
    Fields.DOI: prep_doi,
    Fields.ISBN: prep_isbn,
}

def process_df_split(split_df: pd.DataFrame) -> pd.DataFrame:
    split_df.replace(
        to_replace={
            "UNKNOWN": "",
            "&": "and",
            " & ": " and ",
            " + ": " and ",
        },
        inplace=True,
    )

    set_container_title(split_df)

    split_df["author_full"] = split_df[Fields.AUTHOR]

    for field, function in function_mapping.items():
        split_df[field] = function(split_df[field].values)  # type: ignore

    # TODO : integrate into prep_author if the author_full (minimally processed) is effective
    split_df[Fields.AUTHOR] = select_authors(split_df[Fields.AUTHOR].values)

    split_df.loc[split_df[Fields.PAGES] == split_df[Fields.YEAR], Fields.PAGES] = ""

    if Fields.BOOKTITLE in split_df.columns:
        split_df = split_df.drop(columns=[Fields.BOOKTITLE])
    if Fields.JOURNAL in split_df.columns:
        split_df = split_df.drop(columns=[Fields.JOURNAL])

    split_df = split_df.fillna("")

    return split_df

def get_records_for_dedupe(records_df: pd.DataFrame, *, cpu: int = -1) -> pd.DataFrame:
    """Prepare records for dedupe"""

    print(f"{records_df.shape[0]} records")

    print("Prep started at", datetime.now().strftime("%Y-%m-%d %H:%M:%S"))

    start_time = time.time()

    if 0 == records_df.shape[0]:
        return {}

    assert all(
        f in records_df.columns
        for f in [Fields.ENTRYTYPE, Fields.TITLE, Fields.AUTHOR, Fields.YEAR]
    )

    if Fields.STATUS not in records_df:
        records_df[Fields.STATUS] = "md_processed"
        print("setting missing status")

    records_df = records_df[
        ~(
            records_df[Fields.STATUS].isin(
                [
                    "md_imported",
                    "md_needs_manual_preparation",
                ]
            )
        )
    ]

    records_df[Fields.TITLE] = records_df[Fields.TITLE].replace(
        ["#NAME?", "UNKNOWN", ""], np.nan
    )
    if records_df[Fields.TITLE].isnull().any():
        print(
            "Warning: Some records have empty title field. These records will be dropped."
        )
        records_df = records_df.dropna(subset=[Fields.TITLE])

    # TODO : print a warning
    records_df = records_df[records_df[Fields.ENTRYTYPE] != "proceedings"]

    optional_fields = [
        Fields.JOURNAL,
        Fields.BOOKTITLE,
        Fields.SERIES,
        Fields.VOLUME,
        Fields.NUMBER,
        Fields.PAGES,
        Fields.ABSTRACT,
        Fields.ISBN,
        Fields.DOI,
    ]
    for optional_field in optional_fields:
        if optional_field not in records_df:
            records_df[optional_field] = ""

    records_df.drop(
        labels=list(
            records_df.columns.difference(
                [
                    Fields.ID,
                    Fields.ENTRYTYPE,
                    Fields.AUTHOR,
                    Fields.TITLE,
                    Fields.YEAR,
                    Fields.JOURNAL,
                    Fields.CONTAINER_TITLE,
                    Fields.BOOKTITLE,
                    Fields.VOLUME,
                    Fields.NUMBER,
                    Fields.PAGES,
                    Fields.STATUS,
                    Fields.ABSTRACT,
                    Fields.URL,
                    Fields.ISBN,
                    Fields.DOI,
                ]
            )
        ),
        axis=1,
        inplace=True,
    )

    records_df[
        [
            Fields.STATUS,
            Fields.AUTHOR,
            Fields.TITLE,
            Fields.JOURNAL,
            Fields.YEAR,
            Fields.VOLUME,
            Fields.NUMBER,
            Fields.PAGES,
            Fields.ABSTRACT,
            Fields.DOI,
            Fields.ISBN,
        ]
    ] = records_df[
        [
            Fields.STATUS,
            Fields.AUTHOR,
            Fields.TITLE,
            Fields.JOURNAL,
            Fields.YEAR,
            Fields.VOLUME,
            Fields.NUMBER,
            Fields.PAGES,
            Fields.ABSTRACT,
            Fields.DOI,
            Fields.ISBN,
        ]
    ].astype(
        str
    )

    # TODO : extract
    if cpu == -1:
        # -2
        cpu = os.cpu_count() if os.cpu_count() else 1  # type: ignore

    if records_df.shape[0] < 100:
        # don't split (otherwise, a vectorizeon size 0 inputs is raised)
        cpu = 1

    df_split = np.array_split(records_df, cpu)

    with concurrent.futures.ProcessPoolExecutor(max_workers=cpu) as executor:
        results = executor.map(process_df_split, df_split)

    records_df = pd.concat(list(results))

    records_df["ID"] = records_df["ID"].astype(str)
    records_df.loc[records_df["title"] == "nan", "title"] = ""
    records_df.loc[records_df["year"] == "nan", "year"] = ""
    records_df.loc[records_df["author_full"] == "nan", "author_full"] = ""
    records_df.loc[records_df["container_title"] == "nan", "container_title"] = ""

    end_time = time.time()
    print(f"Prep completed after: {end_time - start_time:.2f} seconds")

    return records_df


import torch
import seaborn as sns
import matplotlib.pyplot as plt
import torch.nn.functional as F
from omegaconf import OmegaConf
from huggingface_hub.file_download import hf_hub_download

from sample import Generator
from train import preprocess_dataset
from models.reverse_diffusion import Unet
from models.forward_diffusion import ForwardDiffusion

def test_generation(
    gen: Generator,
    cfg: OmegaConf,
    batch_size: int = 4,
):
    noise = torch.randn((batch_size, 2, 1000)).to(cfg.train.device)

    # if intermediate_outputs=True returns dict of intermediate signals else only denoised signal
    fake_signals = gen.sample(noise, intermediate_outputs=True)

    # idx for plotted intermediate image
    idx_to_plot = [255, 127, 63, 30, 10, 0]
    plotted_signals = [fake_signals[i].cpu() for i in idx_to_plot]

    # plot fake
    fig, axes = plt.subplots(batch_size, len(plotted_signals), figsize=(20, 10))

    # iterate over signals
    for i, ax_rows in enumerate(axes):
        # iterate over timesteps
        for j, ax in enumerate(ax_rows):
            s = plotted_signals[j][i]

            sns.lineplot(s[0], ax=ax)
            sns.lineplot(s[1], alpha=0.6, ax=ax)

            ax.set_title(f"Signal {i} at timestep {idx_to_plot[j]}")

    plt.tight_layout()
    plt.show()

def test_denoising_val_set(
    forward_diffusion: ForwardDiffusion, generator: Generator, cfg: OmegaConf, batch_size: int = 8, timesteps: int = 63
):
    # check denoising on val set
    _, val_dataloader, _ = preprocess_dataset("roszcz/ecg-segmentation-ltafdb", 128, 1, overfit_single_batch=False)
    records = next(iter(val_dataloader))
    signals = records["signal"].to(cfg.train.device)
    mask = records["mask"]
    # find signals with heartbeats
    indices = [idx for idx in range(len(mask)) if mask[idx].sum() > 0]

    # noising s
    t = torch.tensor([timesteps for _ in range(len(signals))]).to(cfg.train.device)

    signals_noisy, _ = forward_diffusion(signals, t)
    denoised_signals = generator.sample(signals_noisy, timesteps=timesteps, intermediate_outputs=False)

    signals = signals.cpu()
    denoised_signals = denoised_signals.cpu()

    # get first 0-batch_size signals with heartbeats
    idx = torch.tensor(indices[:batch_size], dtype=torch.long)
    signals = signals[idx]
    denoised_signals = denoised_signals[idx]

    # plot true
    fig, axes = plt.subplots(batch_size, 2, figsize=(20, 10))

    for i, ax in enumerate(axes):
        loss = F.mse_loss(denoised_signals[i], signals[i])
        print(f"Loss for signal {i}: {loss}")

        # signals
        sns.lineplot(signals[i, 0], ax=ax[0])
        sns.lineplot(signals[i, 1], alpha=0.6, ax=ax[0])
        ax[0].set_title("Original signal")

        # denoised signal
        sns.lineplot(denoised_signals[i, 0], ax=ax[1])
        sns.lineplot(denoised_signals[i, 1], alpha=0.6, ax=ax[1])
        ax[1].set_title("Denoised signal")

    plt.tight_layout()
    plt.show()

if __name__ == "__main__":
    checkpoint = torch.load(
        hf_hub_download(repo_id="JasiekKaczmarczyk/ecg-diffusion", filename="ecg-diffusion-2023-08-18-11-00.ckpt")
    )

    cfg = checkpoint["config"]

    # model
    model = Unet(
        in_channels=cfg.models.unet.in_out_channels,
        out_channels=cfg.models.unet.in_out_channels,
        dim=cfg.models.unet.dim,
        dim_mults=cfg.models.unet.dim_mults,
        kernel_size=cfg.models.unet.kernel_size,
        resnet_block_groups=cfg.models.unet.num_resnet_groups,
    ).to(cfg.train.device)

    # forward diffusion
    forward_diffusion = ForwardDiffusion(
        beta_start=cfg.models.forward_diffusion.beta_start,
        beta_end=cfg.models.forward_diffusion.beta_end,
        timesteps=cfg.models.forward_diffusion.timesteps,
        schedule_type=cfg.models.forward_diffusion.schedule_type,
    ).to(cfg.train.device)

    model.load_state_dict(checkpoint["model"])
    forward_diffusion.load_state_dict(checkpoint["forward_diffusion"])

    gen = Generator(model, forward_diffusion)

    # test generation of fake data
    test_generation(gen, cfg, batch_size=4)

    # test if model is able to denoise slightly noisy val set
    test_denoising_val_set(forward_diffusion, gen, cfg, batch_size=4, timesteps=63)


import pyautogui
from time import sleep

n = input("How many message do you want to send?\n")
msg = input("What is your message?\n")

n = int(n)

print("Please select the message fild and wait...")

sleep(10)

print("The program is starting!")

while True:
    pyautogui.typewrite(msg)
    sleep(.250)
    pyautogui.typewrite("\n")

    n = n - 1

    print("{} messages left.".format(n))

    if n == 0:
        break

print("All the messages was sent!")


import codecademylib3_seaborn
from matplotlib import pyplot as plt
import pandas as pd
import seaborn as sns

## World Cup Matches Dataframe ##

df = pd.read_csv('WorldCupMatches.csv')
print(df.head())
# Create a new column and set it equal to the sum of the columns Home Team Goals and Away Team Goals
df['Total Goals'] = df['Home Team Goals'] + df['Away Team Goals']
print(df.head())

## Bar chart ##

sns.set_context('poster', font_scale=0.8)
sns.set_style('whitegrid')
f, ax = plt.subplots(figsize=(20, 7))
ax = sns.barplot(x=df['Year'], y=df['Total Goals'])
ax.set_title('Average Number Of Goals Scored In World Cup Matches By Year')
plt.show()

## Goals Dataframe ##

df_goals = pd.read_csv('goals.csv')
print(df_goals.head())

## Box plot ##

sns.set_context('notebook', font_scale=1.25)
f, ax2 = plt.subplots(figsize=(20, 7))
ax2 = sns.boxplot(data=df_goals, x='year', y='goals', palette='Spectral')
ax2.set_title('Goals Visualization')
plt.show()


import FWCore.ParameterSet.Config as cms

import os
from relValTools import getFilesFromEOS, getFilesFromDAS, runtype_to_sample
from test_files import test_files

######## Parsing options ########
from FWCore.ParameterSet.VarParsing import VarParsing
options = VarParsing('python')
options.register('runtype', '', VarParsing.multiplicity.singleton, VarParsing.varType.string, "choose sample type; choices=['ZTT', 'ZEE', 'ZMM', 'QCD', 'TTbar', 'TTbarTau', 'ZpTT']")
options.register('release', 'CMSSW_9_4_0_pre2', VarParsing.multiplicity.singleton, VarParsing.varType.string, 'Release string')
options.register('globalTag', 'PU25ns_94X_mc2017_realistic_v1-v1', VarParsing.multiplicity.singleton, VarParsing.varType.string, 'Global tag')
options.register('useRecoJets', False, VarParsing.multiplicity.singleton, VarParsing.varType.bool, 'Use RecoJets, 0 - false, 1 - true ')
options.register('storageSite', 'eos', VarParsing.multiplicity.singleton, VarParsing.varType.string, "Choose between samples store on eos or DAS or in private local folder. choices=['eos','das', 'loc']")
options.register('localdir', '', VarParsing.multiplicity.singleton, VarParsing.varType.string, "Local dir where the samples are looked up. \
    Output is always created under the local dir if this option is activated and the outputfile is not")
options.register('inputfile', '', VarParsing.multiplicity.singleton, VarParsing.varType.string, "Single file location for fast checks")
options.register('outputFileName', '', VarParsing.multiplicity.singleton, VarParsing.varType.string, "Single file location for fast checks")
options.register('debug', False, VarParsing.multiplicity.singleton, VarParsing.varType.bool, "Debug option")
options.register('key', '', VarParsing.multiplicity.singleton, VarParsing.varType.string, "Key for the input file")
options.parseArguments()

RelVal = options.release
globalTag = options.globalTag
useRecoJets = options.useRecoJets
storageSite = options.storageSite
localdir = options.localdir
debug = options.debug
runtype = options.runtype

key = options.key
if len(key) == 0:
    key = runtype_to_sample[runtype] + "_" + RelVal + "_" + globalTag

if len(localdir) > 1 and localdir[-1] is not "/": localdir += "/"

######## Input files ########
inputfile = options.inputfile
if inputfile != "": filelist = [inputfile]
else:
    path = '/store/relval/{}/{}/MINIAODSIM/{}'.format(RelVal, runtype_to_sample[runtype], globalTag)

    if   storageSite == "eos": filelist = getFilesFromEOS(path)
    elif storageSite == "das": filelist = getFilesFromDAS(RelVal, runtype_to_sample[runtype], globalTag)
    elif storageSite == 'loc': filelist = getFilesFromEOS(localdir + runtype_to_sample[runtype] + "/" + RelVal + '-' + globalTag + '/', cmseospath=False)

    if len(filelist) == 0:
        print ('Sample', RelVal, runtype, 'does not exist in', path)
        sys.exit(0)

if storageSite == 'loc': filelist = ['file:' + x for x in filelist]

if key not in test_files.keys():
    test_files[key] = { 'file' : filelist }
else:
    print ("!Warning! Such key already in the default list:", key)

print ("filelist:", filelist)

######## Output file ########
outputFileName = options.outputFileName
if len(outputFileName) == 0:
    if storageSite == 'loc':
        outputFileName = localdir + runtype_to_sample[runtype] + "/" + RelVal + '-' + globalTag + '/' + 'Updated/'

        if not os.path.isdir(outputFileName):

            result = subprocess.check_output("mkdir -p {outputFileName}".format(outputFileName=outputFileName), shell=True)

    outputFileName += 'output.root'
else:
    if "/" in outputFileName and outputFileName[0] != "/":
        print ("location of output file has a dir structure but doesn't start with dash")
        sys.exit(0)
    if outputFileName[-5:] != ".root":
        outputFileName += '.root'
        print ("output file should have a root format - added automatically:", outputFileName)

print ("outputFileName:", outputFileName)

#---------------------------------- Parameter Set --------------------------------------------

process = cms.Process("produceTauIdMVATrainingNtupleMiniAOD")
#process = cms.Process('TauID', eras.Run2_2017, eras.run2_nanoAOD_94XMiniAODv2)

######## Import of standard configurations ########
process.load('Configuration.StandardSequences.Services_cff')
process.load('SimGeneral.HepPDTESSource.pythiapdt_cfi')
process.load('Configuration.EventContent.EventContent_cff')
process.load('FWCore/MessageService/MessageLogger_cfi')
process.load('Configuration.StandardSequences.GeometryRecoDB_cff')
process.load('Configuration.StandardSequences.MagneticField_cff')
process.load('Configuration/StandardSequences/FrontierConditions_GlobalTag_cff')
process.load('Configuration.StandardSequences.EndOfProcess_cff')
process.MessageLogger.cerr.FwkReport.reportEvery = 1000
process.MessageLogger.cerr.threshold = cms.untracked.string('INFO')
#process.load('Configuration.StandardSequences.Geometry_cff')
#process.load('Configuration.Geometry.GeometryIdeal_cff')
# process.load("TrackingTools.TransientTrack.TransientTrackBuilder_cfi")

from Configuration.AlCa.GlobalTag import GlobalTag
process.GlobalTag = GlobalTag(process.GlobalTag, '94X_mc2017_realistic_v14', '')# process.GlobalTag = GlobalTag(process.GlobalTag, 'auto:phase1_2017_realistic', '')

process.source = cms.Source("PoolSource",
    fileNames = cms.untracked.vstring(test_files[key]['file']),
    secondaryFileNames = cms.untracked.vstring()
)

process.maxEvents = cms.untracked.PSet( input = cms.untracked.int32(options.maxEvents))
process.options = cms.untracked.PSet(wantSummary = cms.untracked.bool(True))
#--------------------------------------------------------------------------------

from runTauIdMVA import TauIDEmbedder
na = TauIDEmbedder(process, cms,
    debug=False,
    toKeep = ["2017v1", "2017v2", "newDM2017v2", "dR0p32017v2", "2016v1", "newDM2016v1"]
)
na.runTauID()

print (dir(process.loadRecoTauTagMVAsFromPrepDB.toGet))
print (process.loadRecoTauTagMVAsFromPrepDB.toGet[-1])
#--------------------------------------------------------------------------------
# process.p = cms.Path(process.rerunMvaIsolationSequence * process.NewTauIDsEmbedded)

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

# Output definition (MiniAOD + updated taus)
# MiniAOD output
print("output prep")
process.out = cms.OutputModule("PoolOutputModule",
    compressionAlgorithm = cms.untracked.string('LZMA'),
    compressionLevel = cms.untracked.int32(4),
    dataset = cms.untracked.PSet(
        dataTier = cms.untracked.string('MINIAODSIM'),
        filterName = cms.untracked.string('')
    ),
    dropMetaData = cms.untracked.string('ALL'),
    eventAutoFlushCompressedSize = cms.untracked.int32(-900),
    fastCloning = cms.untracked.bool(False),
    fileName = cms.untracked.string(outputFileName),
    outputCommands = process.MINIAODSIMEventContent.outputCommands,
    overrideBranchesSplitLevel = cms.untracked.VPSet(cms.untracked.PSet(
        branch = cms.untracked.string('patPackedCandidates_packedPFCandidates__*'),
        splitLevel = cms.untracked.int32(99)
    ),
    cms.untracked.PSet(
        branch = cms.untracked.string('recoGenParticles_prunedGenParticles__*'),
        splitLevel = cms.untracked.int32(99)
    ),
    cms.untracked.PSet(
        branch = cms.untracked.string('patTriggerObjectStandAlones_slimmedPatTrigger__*'),
        splitLevel = cms.untracked.int32(99)
    ),
    cms.untracked.PSet(
        branch = cms.untracked.string('patPackedGenParticles_packedGenParticles__*'),
        splitLevel = cms.untracked.int32(99)
    ),
    cms.untracked.PSet(
        branch = cms.untracked.string('patJets_slimmedJets__*'),
        splitLevel = cms.untracked.int32(99)
    ),
    cms.untracked.PSet(
        branch = cms.untracked.string('recoVertexs_offlineSlimmedPrimaryVertices__*'),
        splitLevel = cms.untracked.int32(99)
    ),
    cms.untracked.PSet(
        branch = cms.untracked.string('recoCaloClusters_reducedEgamma_reducedESClusters_*'),
        splitLevel = cms.untracked.int32(99)
    ),
    cms.untracked.PSet(
        branch = cms.untracked.string('EcalRecHitsSorted_reducedEgamma_reducedEBRecHits_*'),
        splitLevel = cms.untracked.int32(99)
    ),
    cms.untracked.PSet(
        branch = cms.untracked.string('EcalRecHitsSorted_reducedEgamma_reducedEERecHits_*'),
        splitLevel = cms.untracked.int32(99)
    ),
    cms.untracked.PSet(
        branch = cms.untracked.string('recoGenJets_slimmedGenJets__*'),
        splitLevel = cms.untracked.int32(99)
    ),
    cms.untracked.PSet(
        branch = cms.untracked.string('patJets_slimmedJetsPuppi__*'),
        splitLevel = cms.untracked.int32(99)
    ),
    cms.untracked.PSet(
        branch = cms.untracked.string('EcalRecHitsSorted_reducedEgamma_reducedESRecHits_*'),
        splitLevel = cms.untracked.int32(99)
    )),
    overrideInputFileSplitLevels = cms.untracked.bool(True),
    splitLevel = cms.untracked.int32(0)
)
process.out.outputCommands.append('keep *_NewTauIDsEmbedded_*_*')

# Path and EndPath definitions
process.tauIDUpdate_step = cms.Path(process.rerunMvaIsolationSequence * process.NewTauIDsEmbedded)
process.endjob_step = cms.EndPath(process.endOfProcess)
process.out_step = cms.EndPath(process.out)

# Schedule definition
process.schedule = cms.Schedule(process.tauIDUpdate_step, process.endjob_step, process.out_step)

# Add early deletion of temporary data products to reduce peak memory need
from Configuration.StandardSequences.earlyDeleteSettings_cff import customiseEarlyDelete
process = customiseEarlyDelete(process)
# End adding early deletion

#Customize MessageLogger
print ('No. of events to process:', process.maxEvents.input.value())
if process.maxEvents.input.value() > 10:
     process.MessageLogger.cerr.FwkReport.reportEvery = process.maxEvents.input.value()//10
if process.maxEvents.input.value() > 10000 or process.maxEvents.input.value() < 0:
     process.MessageLogger.cerr.FwkReport.reportEvery = 1000


import dataclasses
import functools
import inspect
import flax
from flax import linen as nn
from flax.core.lift import Out as ScanOut  # pylint: disable=unused-import
from flax.linen import partitioning
import jax
from jax.lax import with_sharding_constraint as jax_pjit_wsc

# TODO: this file contains JAX transform workarounds to fix/move
# upstream, primarily concerning the JAX checkpoint/remat transform and
# workarounds for issues with XLA SPMD and JAX scan transform.

# Workaround a scan(remat(...)) abstraction bug by manually implementing a
# static_argnums behavior for flax remat via closure before applying jax remat.

ScanIn = partitioning.ScanIn  # used in t5_architecture.py

def core_remat_static(fn,
                      variables=True,
                      rngs=True,
                      prevent_cse=True,
                      static_argnums=(),
                      policy=None):
  """Flax functional core remat version with static_argnums."""
  static_argnums = tuple(sorted(static_argnums))

  def _repack_remat_args(dyn_args, static_args):
    """Remake arg list from static and dynamic args given static_argnums."""
    args = []
    s_cnt, d_cnt = 0, 0
    for i in range(len(dyn_args) + len(static_args)):
      if i in static_argnums:
        args.append(static_args[s_cnt])
        s_cnt += 1
      else:
        args.append(dyn_args[d_cnt])
        d_cnt += 1
    return tuple(args)

  def inner(scope_fn, repack_fn, variable_groups, rng_groups, *args):
    static_args = tuple(x for i, x in enumerate(args) if i in static_argnums)
    dyn_args = tuple(x for i, x in enumerate(args) if i not in static_argnums)

    @functools.partial(jax.remat, prevent_cse=prevent_cse, policy=policy)
    @functools.wraps(fn)
    def rematted(variable_groups, rng_groups, *dyn_args):
      args = _repack_remat_args(dyn_args, static_args)
      scope = scope_fn(variable_groups, rng_groups)
      y = fn(scope, *args)
      return y, repack_fn(scope)

    return rematted(variable_groups, rng_groups, *dyn_args)

  return flax.core.lift.pack(
      inner, (variables,), (variables,), (rngs,), name='remat')

def remat(target,
          variables=True,
          rngs=True,
          prevent_cse=True,
          static_argnums=(),
          policy=None,
          methods=None):
  """Flax lifted remat that supports static_argnums."""
  return flax.linen.transforms.lift_transform(
      core_remat_static,
      target,
      variables=variables,
      rngs=rngs,
      prevent_cse=prevent_cse,
      static_argnums=static_argnums,
      policy=policy,
      methods=methods)

# Allow use of scan/remat on factory functions that return module instances.

# Flaxformer uses keyword-only arguments in its methods, which
# aren't natively supported by most JAX transforms.  We use canonicalizing
# method wrappers to present to jax a pure-positional version of the function.

def canonicalize_arguments(orig_fn):
  """Convert function to use positional arguments only.

  Args:
    orig_fn: callable with signature taking positional and keyword arguments,
      but not variadic *args or **kwargs.

  Returns:
    A version of orig_fn taking only positional arguments, and
    a conversion function that takes the original signature, binds
    it to provided mixed arguments, applies defaults, and returns
    a tuple of positional arguments to use with the transformed
    function.
  """
  sig = inspect.signature(orig_fn)
  params = sig.parameters

  def dekwarged_fn(*args):
    new_args = []
    new_kwargs = {}
    if len(args) != len(params):
      raise ValueError(f'Incorrect number of arguments: '
                       f'got {len(args)}, expected {len(params)}.')
    for i, p in enumerate(params):
      param = params[p]
      if param.kind in (param.POSITIONAL_ONLY, param.POSITIONAL_OR_KEYWORD):
        new_args.append(args[i])
      elif param.kind is param.KEYWORD_ONLY:
        new_kwargs[p] = args[i]
      elif param.kind is param.VAR_POSITIONAL:
        new_args.extend(args[i])
      elif param.kind is param.VAR_KEYWORD:
        new_kwargs.update(args[i])
      else:
        raise ValueError('Unknown signature parameter type.')
    return orig_fn(*new_args, **new_kwargs)

  # We don't use functools.wraps because we are changing the signature, but
  # we want function properties preserved.
  dekwarged_fn.__dict__.update(orig_fn.__dict__)

  def convert_to_args(*args, **kwargs):
    bound = sig.bind(*args, **kwargs)
    bound.apply_defaults()
    return tuple(bound.arguments.values())

  return dekwarged_fn, convert_to_args

def canonicalized_class_transform(trafo, clz, *t_args, **t_kwargs):
  """Applies kwarg canonicalization with flax transform to module class clz.

  NB: This function only handles transforming the __call__ method.

  Args:
    trafo: flax lifted transform (e.g. nn.scan, nn.remat)
    clz: nn.Module class to transform
    *t_args: transform arguments
    **t_kwargs: transform keyword-arguments

  Returns:
    A transformed version of clz whose __call__ function has been transformed,
    additionally handling canonicalization of __call__'s signature to a purely
    positional function before applying the transform.
  """
  # Transform postitional only __call__ form of clz.
  dekwarged_fn, convert_to_args = canonicalize_arguments(clz.__call__)
  trafo_fn = trafo(dekwarged_fn, *t_args, **t_kwargs)

  @functools.wraps(clz.__call__)
  def post_fn(self, *args, **kwargs):
    return trafo_fn(*convert_to_args(self, *args, **kwargs))

  return type(trafo.__name__.capitalize() + clz.__name__, (clz,),
              {'__call__': post_fn})

# Flaxformer uses factory functions instead of partial constructors, we need
# to add a explicit handler for dealing with this case as our usual lifting
# API has no way to distinguish a factory function from a class method.

def apply_transform_to_module_factory(trafo, factory, *args, **kwargs):
  """Fix to apply flax transforms to a module factories via re-instantiation."""

  def new_factory():
    # Create the Module instance from the factory in a disconnected dynamic
    # context solely to collect the construction arguments and class.
    nn.module._context.module_stack.append(None)  # pylint: disable=protected-access
    try:
      inst = factory()
      ctor_args = {
          f.name: object.__getattribute__(inst, f.name)
          for f in dataclasses.fields(inst)
          if f.name not in ('parent', 'name')
      }
    finally:
      nn.module._context.module_stack.pop()  # pylint: disable=protected-access
    # Instantiate the transformed module class with gathered construction args
    # in the current dynamic context.
    return canonicalized_class_transform(trafo, inst.__class__, *args,
                                         **kwargs)(**ctor_args)

  return new_factory

factory_remat = functools.partial(apply_transform_to_module_factory,
                                  partitioning.remat)
factory_scan = functools.partial(apply_transform_to_module_factory,
                                 partitioning.scan_with_axes)
factory_vmap = functools.partial(apply_transform_to_module_factory, nn.vmap)

# Scan inner-function SPMD re-annotation.

# The XLA SPMD subsystem currently "loses" annotations on parameter trees that
# pass through an XLA while loop.  We should investigate fixing this at the XLA
# level, but the workaround for now is to re-apply the known sharding
# information for the scanned layer -inside- the functionalized scan body
# function using pjit's with_sharding_constraint.

def global_mesh_defined():
  """Checks if global xmap/pjit mesh resource environment is defined."""
  maps_env = jax.experimental.maps.thread_resources.env
  return maps_env.physical_mesh.devices.shape != ()  # pylint: disable=g-explicit-bool-comparison

def with_sharding_constraint(x, axis_resources):
  """Wrapper for lax.with_sharding_constraint, no-op on cpu or outside pjit."""
  if jax.devices()[0].platform == 'cpu' or not global_mesh_defined():
    return x
  else:
    return jax_pjit_wsc(x, axis_resources)

def inner_scan_spmd(annotation_tree, scan_axis):
  """Workaround to apply a sharding annotation pytree inside a scan body fn.

  This creates a function to be passed to nn.scan's "data_transform" kwarg.

  Args:
    annotation_tree: pytree of PartitionSpecs
    scan_axis: The axis along which layer parameters were scanned.

  Returns:
    A function to be used with nn.scan's data_transform kwarg to apply the
    SPMD annotations on the inner scan body function.
  """
  if annotation_tree is None:
    return None

  # The annotation tree fed through the model is the scan-expanded one,
  # we need to remove the scan axis from these PartitionSpecs.
  def del_axis(x):
    tmp = list(x)
    tmp.pop(scan_axis)
    return type(x)(*tmp)

  annotation_tree = jax.tree_map(del_axis, annotation_tree)

  def annotate_fn(variable_groups, rng_groups):
    broadcast_vars, carry_vars, *scan_variable_groups = variable_groups

    def maybe_annotate_group(x):
      if tuple(x[0].keys()) == ('params',):
        return ({
            'params': with_sharding_constraint(x[0]['params'], annotation_tree)
        },)
      else:
        return x

    scan_variable_groups = tuple(
        map(maybe_annotate_group, scan_variable_groups))
    variable_groups = (broadcast_vars, carry_vars) + scan_variable_groups
    return variable_groups, rng_groups

  return annotate_fn