# -*- coding: utf-8 -*-

import os
import time
import re
from re import sub, match
from lxml import html
from datetime import datetime, timedelta
from typing import Optional, cast

# from copy import deepcopy
from dateutil.relativedelta import *
import traceback
from pprint import pprint

from gluon import current
from gluon.storage import Storage
from gluon.tools import Auth
from gluon.html import *
from gluon.template import render
from gluon.contrib.markdown import WIKI
from gluon.contrib.appconfig import AppConfig
from gluon.tools import Mail
from gluon.validators import IS_EMAIL
from models.article import Article
from pydal import DAL

from gluon.custom_import import track_changes
from models.article import Article
from models.recommendation import Recommendation
from models.user import User

track_changes(True)
import socket

from uuid import uuid4
from contextlib import closing
import shutil

from app_modules import common_tools
from app_modules import common_small_html
from app_modules import old_common
from app_modules import helper
from app_modules import emailing_parts
from app_modules.reminders import getReminder

myconf = AppConfig(reload=True)
contact = myconf.take("contacts.managers")
parallelSubmissionAllowed = myconf.get("config.parallel_submission", default=False)

pciRRactivated = myconf.get("config.registered_reports", default=False)
scheduledSubmissionActivated = myconf.get("config.scheduled_submissions", default=False)

MAIL_DELAY = 1.5  # in seconds

# common view for all emails
MAIL_HTML_LAYOUT = os.path.join(os.path.dirname(__file__), "../../views/mail", "mail.html")

CLEANR = re.compile('<.*?>')

######################################################################################################################################################################
# Mailing tools
######################################################################################################################################################################

appName = myconf.take("app.name")

def email_subject_header(articleId):
    return "%s #%s" % (appName, articleId)

def patch_email_subject(subject, articleId):
    return subject.replace(appName, email_subject_header(articleId))

def mkAuthors(article):
    return article.authors if not article.anonymous_submission else current.T("[undisclosed]")

######################################################################################################################################################################
def getMailer(auth):
    mail = auth.settings.mailer
    mail.settings.server = myconf.take("smtp.server")
    mail.settings.sender = myconf.take("smtp.sender")
    mail.settings.login = myconf.take("smtp.login")
    mail.settings.tls = myconf.get("smtp.tls", default=False)
    mail.settings.ssl = myconf.get("smtp.ssl", default=False)
    return mail

######################################################################################################################################################################
def getMailCommonVars():
    return dict(
        scheme=myconf.take("alerts.scheme"),
        host=myconf.take("alerts.host"),
        port=myconf.take("alerts.port", cast=lambda v: common_tools.takePort(v)),
        appDescription=myconf.take("app.description"),
        appName=myconf.take("app.name"),
        appLongName=myconf.take("app.longname"),
        longname=myconf.take("app.longname"),  # DEPRECATED: for compatibility purposes; to be removed after checking
        appThematics=myconf.take("app.thematics"),
        appContactMail=myconf.take("contacts.managers"),
        appContactLink=A(myconf.take("contacts.managers"), _href="mailto:" + myconf.take("contacts.managers")),
        siteUrl=URL(c="default", f="index", scheme=myconf.take("alerts.scheme"), host=myconf.take("alerts.host"), port=myconf.take("alerts.port")),
    )

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

def getMailForRecommenderCommonVars(auth: Auth, db: DAL, sender: User, article: Article, recommendation: Recommendation, recommender: str, new_round: Optional[str] = None):
    scheme = myconf.take("alerts.scheme")
    host = myconf.take("alerts.host")
    port = myconf.take("alerts.port", cast=lambda v: common_tools.takePort(v))
    is_co_recommender = helper.is_co_recommender(auth, db, recommendation.id)

    mail_vars = getMailCommonVars()
    _recomm = common_tools.get_prev_recomm(db, recommendation) if new_round else recommendation
    r2r_url, trackchanges_url = emailing_parts.getAuthorsReplyLinks(auth, db, _recomm.id)
    r2r_url = str(r2r_url) if r2r_url else "(no author's reply)"
    trackchanges_url = str(trackchanges_url) if trackchanges_url else "(no tracking)"
    # use: r2r_url = r2r_url['_href'] if r2r_url else "(no author's reply)"
    # to pass only the url value to the template instead of the full link html;
    # doing this yields invalid url for the link in the template when no doc exists.

    mail_vars["description"] = myconf.take("app.description")
    mail_vars["longname"] = myconf.take("app.longname") # DEPRECATED
    mail_vars["appLongName"] = myconf.take("app.longname")
    mail_vars["appName"] = myconf.take("app.name")
    mail_vars["thematics"] = myconf.take("app.thematics")
    mail_vars["scheme"] = scheme
    mail_vars["host"] = host
    mail_vars["port"] = port
    mail_vars["site_url"] = URL(c="default", f="index", scheme=scheme, host=host, port=port)
    mail_vars["art_authors"] = mkAuthors(article)
    mail_vars["authors"] = mail_vars["art_authors"]
    mail_vars["articleAuthors"] = mail_vars["art_authors"]
    mail_vars["r2r_url"] = r2r_url
    mail_vars["trackchanges_url"] = trackchanges_url
    mail_vars["destPerson"] = '%s %s'%(recommender.first_name, recommender.last_name)
    mail_vars["LastName"] = recommender.last_name

    if recommendation.doi:
        mail_vars["art_doi"] = common_small_html.mkLinkDOI(recommendation.doi)
        mail_vars["articleDoi"] = mail_vars["art_doi"]
    elif article.doi:
        mail_vars["art_doi"] = common_small_html.mkLinkDOI(article.doi)
        mail_vars["articleDoi"] = mail_vars["art_doi"]

    if article.title:
        mail_vars["art_title"] = common_small_html.md_to_html(article.title)
        mail_vars["articleTitle"] = mail_vars["art_title"]

    if auth.user_id == recommendation.recommender_id:
        mail_vars["sender"] = common_small_html.mkUser(auth, db, recommendation.recommender_id).flatten()
        mail_vars["Institution"] = sender.institution
        mail_vars["Department"] = sender.laboratory
        mail_vars["country"] = sender.country

        if sender.first_name and sender.last_name:
            mail_vars["senderName"] = sender.first_name + ' ' + sender.last_name

    elif is_co_recommender:
        sender = cast(User, auth.user)
        mail_vars["sender"] = common_small_html.mkUser(auth, db, auth.user_id).flatten() + "[co-recommender]"
        mail_vars["Institution"] = sender.institution
        mail_vars["Department"] = sender.laboratory
        mail_vars["country"] = sender.country

        if sender.first_name and sender.last_name:
            mail_vars["senderName"] = sender.first_name + ' ' + sender.last_name

    elif auth.has_membership(role="manager"):
        recommender = User.get_by_id(db, recommendation.recommender_id)
        if recommender:
            mail_vars["sender"] = "The Managing Board of " + myconf.get("app.longname") + " on behalf of " + common_small_html.mkUser(auth, db, recommendation.recommender_id).flatten()
            mail_vars["Institution"] = recommender.institution
            mail_vars["Department"] = recommender.laboratory
            mail_vars["country"] = recommender.country

            if recommender.first_name and recommender.last_name:
                mail_vars["senderName"] = recommender.first_name + ' ' + recommender.last_name

    return mail_vars

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

def getMailForReviewerCommonVars(auth: Auth, db: DAL, sender: User, article: Article, recommendation: Recommendation, reviewer_last_name: Optional[str] = None, new_round: Optional[bool] = False) :
    scheme = myconf.take("alerts.scheme")
    host = myconf.take("alerts.host")
    port = myconf.take("alerts.port", cast=lambda v: common_tools.takePort(v))
    is_co_recommender = helper.is_co_recommender(auth, db, recommendation.id)

    mail_vars = getMailCommonVars()
    _recomm = common_tools.get_prev_recomm(db, recommendation) if new_round else recommendation
    r2r_url, trackchanges_url = emailing_parts.getAuthorsReplyLinks(auth, db, _recomm.id)
    r2r_url = str(r2r_url) if r2r_url else "(no author's reply)"
    trackchanges_url = str(trackchanges_url) if trackchanges_url else "(no tracking)"
    # use: r2r_url = r2r_url['_href'] if r2r_url else "(no author's reply)"
    # to pass only the url value to the template instead of the full link html;
    # doing this yields invalid url for the link in the template when no doc exists.

    mail_vars["description"] = myconf.take("app.description")
    mail_vars["longname"] = myconf.take("app.longname") # DEPRECATED
    mail_vars["appLongName"] = myconf.take("app.longname")
    mail_vars["appName"] = myconf.take("app.name")
    mail_vars["thematics"] = myconf.take("app.thematics")
    mail_vars["scheme"] = scheme
    mail_vars["host"] = host
    mail_vars["port"] = port
    mail_vars["site_url"] = URL(c="default", f="index", scheme=scheme, host=host, port=port)
    mail_vars["art_authors"] = mkAuthors(article)
    mail_vars["authors"] = mail_vars["art_authors"]
    mail_vars["articleAuthors"] = mail_vars["art_authors"]
    mail_vars["r2r_url"] = r2r_url
    mail_vars["trackchanges_url"] = trackchanges_url

    if reviewer_last_name:
        mail_vars["LastName"] = reviewer_last_name

    if recommendation.doi:
        mail_vars["art_doi"] = common_small_html.mkLinkDOI(recommendation.doi)
        mail_vars["articleDoi"] = mail_vars["art_doi"]
    elif article.doi:
        mail_vars["art_doi"] = common_small_html.mkLinkDOI(article.doi)
        mail_vars["articleDoi"] = mail_vars["art_doi"]

    if article.title:
        mail_vars["art_title"] = common_small_html.md_to_html(article.title)
        mail_vars["articleTitle"] = mail_vars["art_title"]

    if auth.user_id == recommendation.recommender_id:
        mail_vars["sender"] = common_small_html.mkUser(auth, db, recommendation.recommender_id).flatten()
        mail_vars["Institution"] = sender.institution
        mail_vars["Department"] = sender.laboratory
        mail_vars["country"] = sender.country

        if sender.first_name and sender.last_name:
            mail_vars["senderName"] = sender.first_name + ' ' + sender.last_name

    elif is_co_recommender:
        sender = cast(User, auth.user)
        mail_vars["sender"] = common_small_html.mkUser(auth, db, auth.user_id).flatten() + "[co-recommender]"
        mail_vars["Institution"] = sender.institution
        mail_vars["Department"] = sender.laboratory
        mail_vars["country"] = sender.country

        if sender.first_name and sender.last_name:
            mail_vars["senderName"] = sender.first_name + ' ' + sender.last_name

    elif auth.has_membership(role="manager"):
        recommender = User.get_by_id(db, recommendation.recommender_id)
        if recommender:
            mail_vars["sender"] = "The Managing Board of " + myconf.get("app.longname") + " on behalf of " + common_small_html.mkUser(auth, db, recommendation.recommender_id).flatten()
            mail_vars["Institution"] = recommender.institution
            mail_vars["Department"] = recommender.laboratory
            mail_vars["country"] = recommender.country

            if recommender.first_name and recommender.last_name:
                mail_vars["senderName"] = recommender.first_name + ' ' + recommender.last_name

    return mail_vars

######################################################################################################################################################################
def getCorrectHashtag(hashtag, article=None, force_scheduled=False):
    if pciRRactivated and article is not None:
        if article.art_stage_1_id is not None or article.report_stage == "STAGE 2":
            hashtag += "Stage2"
        else:
            hashtag += "Stage1"

    if scheduledSubmissionActivated and article is not None:
        if (article.scheduled_submission_date is not None) or (article.status.startswith("Scheduled submission")) or (force_scheduled):
            hashtag += "ScheduledSubmission"

    return hashtag

#######################################################################################################################################################################
def to_string_addresses(address_list):
    return str(address_list).replace('[','').replace(']','').replace("'",'').replace('"','')

#######################################################################################################################################################################
def list_addresses(addresses):
    return [x.strip(' ') for x in list(re.split("[,; ]", addresses))] \
                if addresses else []

#######################################################################################################################################################################
def clean_addresses(dirty_string_adresses):
    '''
    creates a string of clean mail addresses, divided by comma
    '''
    if dirty_string_adresses == None: return '', ''

    list_of_contacts = [contact.strip() for contact in list(re.split("[,;]", dirty_string_adresses))]
    contacts = []
    errors = []
    validator = IS_EMAIL()
    for contact in list_of_contacts:
        if len(contact.split(' ')) > 1:
            for word in contact.split(' '):
                if '@' in word:
                    contact = word
        contact = contact.strip('<>')
        value, error = validator(contact)
        if error is None: contacts.append(contact)
        else: errors.append(contact)

    return ','.join(contacts), ', '.join(errors)

######################################################################################################################################################################
def getMailTemplateHashtag(db, hashTag, myLanguage="default"):
    print(hashTag)
    query = (db.mail_templates.hashtag == hashTag) & (db.mail_templates.lang == myLanguage)
    item = db(query).select().first()

    if item:
        return dict(subject=item.subject, content=item.contents)
    else:
        if scheduledSubmissionActivated and pciRRactivated:
            return generateNewMailTemplates(db, hashTag, myLanguage)
        else:
            return dict(error=True, message="hashtag not found")

######################################################################################################################################################################
def generateNewMailTemplates(db, hashTag, myLanguage):
    baseHashtag = hashTag
    baseHashtag = baseHashtag.replace("Stage1", "")
    baseHashtag = baseHashtag.replace("Stage2", "")
    baseHashtag = baseHashtag.replace("ScheduledSubmission", "")

    # Create stage 1 template
    result1 = insertNewTemplateInDB(db, baseHashtag + "Stage1ScheduledSubmission", baseHashtag + "Stage1", myLanguage)

    # Create stage 2 template
    result2 = insertNewTemplateInDB(db, baseHashtag + "Stage2ScheduledSubmission", baseHashtag + "Stage2", myLanguage)

    if "Stage1" in hashTag:
        return result1
    elif "Stage2" in hashTag:
        if 'error' in result2 and result2['error']:
            return result1
        return result2
    else:
        return dict(error=True, message="hashtag not found")

######################################################################################################################################################################
def insertNewTemplateInDB(db, newHashTag, baseHashtag, myLanguage):
    query = (db.mail_templates.hashtag == baseHashtag) & (db.mail_templates.lang == myLanguage)
    item = db(query).select().first()

    if item:
        print("new template created :" + newHashTag)
        db.mail_templates.insert(
            hashtag=newHashTag, subject=item.subject + " - scheduled submission", contents=item.contents, description=item.description + " (for scheduled submission)"
        )
        return dict(subject=item.subject, content=item.contents)
    else:
        return dict(error=True, message="hashtag not found")

######################################################################################################################################################################
def createMailReport(mail_resu, destPerson, reports):
    if mail_resu:
        reports.append(dict(error=False, message="e-mail sent to %s" % destPerson))
    else:
        reports.append(dict(error=True, message="e-mail NOT SENT to %s" % destPerson))
    return reports

######################################################################################################################################################################
def getFlashMessage(session, reports):
    messages = []

    for report in reports:
        if report["error"]:
            session.flash_status = "warning"

        messages.append(report["message"])
        pass

    print("\n".join(messages))
    if session.flash is None:
        session.flash = "; ".join(messages)
    else:
        session.flash += "; " + "; ".join(messages)

######################################################################################################################################################################
def getMailFooter():
    with open(os.path.join(os.path.dirname(__file__), "../../views/mail", "mail_footer.html"), encoding="utf-8") as myfile:
        data = myfile.read()
    return data

######################################################################################################################################################################
# Footer for all mails
def mkFooter(db):
    # init mail_vars with common infos
    mail_vars = getMailCommonVars()

    # add vars to mail_vars
    mail_vars["baseurl"] = URL(c="default", f="index", scheme=mail_vars["scheme"], host=mail_vars["host"], port=mail_vars["port"])
    mail_vars["profileurl"] = URL(
        c="default",
        f="user",
        args=("login"),
        vars=dict(_next=URL(c="default", f="user", args=("profile"))),
        scheme=mail_vars["scheme"],
        host=mail_vars["host"],
        port=mail_vars["port"],
    )

    footer_template = getMailTemplateHashtag(db, "#EmailFooterTemplate")
    if 'content' in footer_template and footer_template["content"]:
        footer_content = replaceMailVars(footer_template["content"], mail_vars)
        return XML(footer_content)
    else:
        return XML(getMailFooter() % mail_vars)

######################################################################################################################################################################
def insertMailInQueue(
    auth,
    db,
    hashtag_template,
    mail_vars,
    recommendation_id=None,
    recommendation=None,
    article_id=None,
    review=None,
    authors_reply=None,
    sugg_recommender_buttons=None,
    reviewer_invitation_buttons=None,
):
    mail = buildMail(
        db,
        hashtag_template,
        mail_vars,
        recommendation=recommendation,
        review=review,
        authors_reply=authors_reply,
        article_id=article_id,
        sugg_recommender_buttons=sugg_recommender_buttons,
        reviewer_invitation_buttons=reviewer_invitation_buttons,
    )

    return db.mail_queue.insert(
        dest_mail_address=mail_vars["destAddress"],
        cc_mail_addresses=mail_vars.get("ccAddresses"),
        replyto_addresses=mail_vars.get("replytoAddresses"),
        bcc_mail_addresses=mail_vars.get("bccAddresses"),
        mail_subject=mail["subject"],
        mail_content=mail["content"],
        user_id=auth.user_id,
        article_id=article_id,
        recommendation_id=recommendation_id,
        mail_template_hashtag=hashtag_template,
    )

######################################################################################################################################################################
def insertReminderMailInQueue(
    auth,
    db,
    hashtag_template,
    mail_vars,
    recommendation_id=None,
    recommendation=None,
    article_id=None,
    review_id=None,
    review=None,
    authors_reply=None,
    sending_date_forced=None,
    base_sending_date=None,
    reviewer_invitation_buttons=None,
    sender_name: Optional[str]=None
):

    reminder = getReminder(hashtag_template, db.t_reviews[review_id])

    ccAddresses = None
    replytoAddresses = None
    if "ccAddresses" in mail_vars:
        ccAddresses = mail_vars["ccAddresses"]
    if "replytoAddresses" in mail_vars:
        replytoAddresses = mail_vars["replytoAddresses"]

    if reminder:
        elapsed_days = reminder["elapsed_days"][0]
        sending_date = datetime.now() if not base_sending_date \
                else base_sending_date - timedelta(days=7)
        sending_date += timedelta(days=elapsed_days)

    if sending_date_forced:
        sending_date = sending_date_forced

    if True:

        mail = buildMail(
            db, hashtag_template, mail_vars, recommendation=recommendation, review=review, authors_reply=authors_reply, reviewer_invitation_buttons=reviewer_invitation_buttons,
            article_id=article_id,
        )

        db.mail_queue.insert(
            sending_status="pending",
            sending_date=sending_date,
            dest_mail_address=mail_vars["destAddress"],
            cc_mail_addresses=ccAddresses,
            replyto_addresses=replytoAddresses,
            mail_subject=mail["subject"],
            mail_content=mail["content"],
            user_id=auth.user_id,
            recommendation_id=recommendation_id,
            article_id=article_id,
            mail_template_hashtag=hashtag_template,
            review_id=review_id,
            sender_name=sender_name
        )

######################################################################################################################################################################
def insertNewsLetterMailInQueue(
    auth,
    db,
    mail_vars,
    hashtag_template,
    newRecommendations=None,
    newRecommendationsCount=0,
    newPreprintSearchingForReviewers=None,
    newPreprintSearchingForReviewersCount=0,
    newPreprintRequiringRecommender=None,
    newPreprintRequiringRecommenderCount=0,
):

    mail = buildNewsLetterMail(
        db,
        mail_vars,
        hashtag_template,
        newRecommendations,
        newRecommendationsCount,
        newPreprintSearchingForReviewers,
        newPreprintSearchingForReviewersCount,
        newPreprintRequiringRecommender,
        newPreprintRequiringRecommenderCount,
    )

    db.mail_queue.insert(
        dest_mail_address=mail_vars["destAddress"], mail_subject=mail["subject"], mail_content=mail["content"], user_id=auth.user_id, mail_template_hashtag=hashtag_template
    )

######################################################################################################################################################################
def buildMail(db, hashtag_template, mail_vars, recommendation=None, review=None, authors_reply=None, sugg_recommender_buttons=None, reviewer_invitation_buttons=None,
        article_id=None,
    ):

    mail_template = getMailTemplateHashtag(db, hashtag_template)

    subject = replaceMailVars(mail_template["subject"], mail_vars)
    content = replaceMailVars(mail_template["content"], mail_vars)

    if article_id is None:
        subject_without_appname = subject.replace("%s: " % mail_vars["appName"] , "")
    else:
        subject = patch_email_subject(subject, article_id)
        appname_with_article_id = email_subject_header(article_id)
        subject_without_appname = subject.replace("%s: " % appname_with_article_id , "")

    applogo = URL("static", "images/small-background.png", scheme=mail_vars["scheme"], host=mail_vars["host"], port=mail_vars["port"])

    content_rendered = render(
        filename=MAIL_HTML_LAYOUT,
        context=dict(
            subject=subject_without_appname,
            applogo=applogo,
            appname=mail_vars["appName"],
            content=XML(content),
            footer=mkFooter(db),
            recommendation=recommendation,
            review=review,
            authors_reply=authors_reply,
            sugg_recommender_buttons=sugg_recommender_buttons,
            reviewer_invitation_buttons=reviewer_invitation_buttons,
        ),
    )

    return dict(content=content_rendered, subject=subject)

######################################################################################################################################################################
def buildNewsLetterMail(
    db,
    mail_vars,
    hashtag_template,
    newRecommendations=None,
    newRecommendationsCount=0,
    newPreprintSearchingForReviewers=None,
    newPreprintSearchingForReviewersCount=0,
    newPreprintRequiringRecommender=None,
    newPreprintRequiringRecommenderCount=0,
):
    mail_template = getMailTemplateHashtag(db, hashtag_template)

    subject = replaceMailVars(mail_template["subject"], mail_vars)
    content = replaceMailVars(mail_template["content"], mail_vars)

    full_link = { var: mail_vars[var] for var in ["scheme", "host", "port"] }

    subject_without_appname = subject.replace("%s: " % mail_vars["appName"], "")
    applogo = URL("static", "images/small-background.png", **full_link)

    allRecommendationsLink = A(
        current.T("See more recommendations..."),
        _href=URL('default','index', **full_link),
        _style="border-radius: 5px; font-weight: bold; padding: 6px 20px; color: #ffffff; background-color: #3e3f3a;",
    )

    content_rendered = render(
        filename=MAIL_HTML_LAYOUT,
        context=dict(
            subject=subject_without_appname,
            applogo=applogo,
            appname=mail_vars["appName"],
            content=XML(content),
            newRecommendations=XML(newRecommendations),
            newRecommendationsCount=newRecommendationsCount,
            allRecommendationsLink=allRecommendationsLink,
            newPreprintSearchingForReviewers=XML(newPreprintSearchingForReviewers),
            newPreprintSearchingForReviewersCount=newPreprintSearchingForReviewersCount,
            newPreprintRequiringRecommender=XML(newPreprintRequiringRecommender),
            newPreprintRequiringRecommenderCount=newPreprintRequiringRecommenderCount,
            pciRRactivated=pciRRactivated,
            footer=mkFooter(db),
            mail_vars=mail_vars,
        ),
    )

    return dict(content=content_rendered, subject=subject)

######################################################################################################################################################################
def replaceMailVars(text, mail_vars):
    mail_vars_list = mail_vars.keys()

    for var in mail_vars_list:
        if text.find("{{" + var + "}}") > -1:
            if isinstance(mail_vars[var], str):
                replacement_var = mail_vars[var]
            elif isinstance(mail_vars[var], int):
                replacement_var = str(mail_vars[var])
            else:
                try:
                    replacement_var = mail_vars[var].flatten()
                except:
                    replacement_var = str(mail_vars[var])

            text = text.replace("{{" + var + "}}", replacement_var)

    return text

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

def replace_mail_vars_set_not_considered_mail(auth: Auth, db: DAL, article: Article, subject: str, message: str):
    form = Storage(subject=subject, message=message, cc=contact)

    mail_vars = getMailCommonVars()
    mail_vars['destPerson'] = common_small_html.mkUser(auth, db, article.user_id)
    mail_vars['articleTitle'] = common_small_html.md_to_html(article.title)
    mail_vars['unconsider_limit_days'] = myconf.get("config.unconsider_limit_days", default=20)

    form.subject = replaceMailVars(form.subject, mail_vars)
    form.message = replaceMailVars(form.message, mail_vars)

    return form


from dotenv import load_dotenv
load_dotenv()

import click
from model import Graph

@click.group()
def cli():
    pass

@cli.command()
@click.argument('name')
@click.argument('age')
def add(name, age):
    added = Graph.add(name, age)
    if added:
        click.echo(f'Created new Person: {name} ({age})')

@cli.command()
@click.argument('name')
def remove(name):
    removed = Graph.remove(name)
    if removed:
        click.echo(f'Removed {name}')

@cli.command()
@click.argument('name')
@click.argument('age')
def edit(name, age):
    edited = Graph.edit(name, age)
    if edited:
        click.echo(f'Edited {name} successfully, now is {age} years old.')

@cli.command()
def info():
    result = Graph.info()
    click.echo(f'Found {result["n"]} Nodes and {result["r"]} Relationships')

@cli.command()
@click.argument('name')
@click.argument('friend')
def befriend(name, friend):
    ok = Graph.befriend(name, friend)
    success_msg = f'{name} and {friend} are now friends!'
    error_msg = 'No relationship created'
    click.echo(success_msg if ok else error_msg)

@cli.command('list')
@click.argument('name')
@click.option('--depth', default=1)
@click.option('--full', default=False, is_flag=True)
def list_friends(name, depth, full):
    friends = Graph.list_friends(name, depth, full)
    for friend in friends:
        age = f' ({friend["age"]})' if friend["age"] else ''
        click.echo(f'{friend["name"]}{age}')

if __name__ == '__main__':
    cli()


# (c) Nelen & Schuurmans.  GPL licensed, see LICENSE.rst.
from django.conf.urls import include
from django.conf.urls import patterns
from django.conf.urls import url
from django.contrib import admin
from django.views.generic import TemplateView

from lizard_ui.urls import debugmode_urlpatterns

from lizard_damage import forms
from lizard_damage import views

admin.autodiscover()

urlpatterns = patterns(
    '',
    url(r'^ui/', include('lizard_ui.urls')),
    url(r'^admin/', include(admin.site.urls)),

    url(
        r'^$',
        views.Wizard.as_view(
            [forms.FormStep0,
             forms.FormStep1,
             forms.FormStep2,
             forms.FormStep3,
             forms.FormStep4,
             forms.FormStep5,
             forms.FormStep6,
             forms.FormBatenKaart,  # '7'
             forms.FormZipResult,  # '8' for batch zip (single file)
             forms.FormZipResult,  # '9' for baten kaart (2 files)
             forms.FormStepUniformLevelsBatch,  # '10' for batch berekening
             ],
            initial_dict={
                '0': {
                    'name': 'Nieuw scenario',
                    },
                '1': {
                    'floodtime': 1,
                    'flooddate': 9,
                    'repairtime_roads': "0.25",
                    'repairtime_buildings': "1",
                    },
                '2': {
                    'floodtime': 1,
                    'flooddate': 9,
                    'repairtime_roads': "0.25",
                    'repairtime_buildings': "1",
                    }
                },
            condition_dict={
                # Step 1, enable for calc_type 0
                '1': views.show_form_condition([0]),
                # Step 2, enable for calc_type 1, etc
                '2': views.show_form_condition([1]),
                '3': views.show_form_condition([2]),
                '4': views.show_form_condition([3]),
                '5': views.show_form_condition([4]),
                '6': views.show_form_condition([5]),
                '7': views.show_form_condition([6]),
                # Check zipfile and show results
                '8': views.show_form_condition([2, 3, 4, 5]),
                '9': views.show_form_condition([6]),  # Batenkaart files
                '10': views.show_form_condition([7]),  # Batch
                }
            ),
        name='lizard_damage_form'
    ),
    url(r'^disclaimer$',
        views.Disclaimer.as_view(
            template_name="lizard_damage/disclaimer.html"),
        name='lizard_damage_disclaimer'),
    url(
        r'^result/(?P<slug>.*)/$',
        views.DamageScenarioResult.as_view(),
        name='lizard_damage_result'
    ),
    url(
        r'^benefit_result/(?P<slug>.*)/$',
        views.BenefitScenarioResult.as_view(),
        name='lizard_damage_benefit_result'
    ),
    url(
        r'^benefit_result/(?P<slug>.*)/$',
        views.BenefitScenarioResult.as_view(),
        name='lizard_damage_benefit_result'
    ),
    url(
        r'^benefit_result/(?P<slug>.*)/kml$',
        views.BenefitScenarioKML.as_view(),
        name='lizard_damage_benefit_kml'
    ),
    url(
        r'^event/(?P<slug>.*)/(?P<result_type>.*)/kml/$',
        views.DamageEventKML.as_view(),
        name='lizard_damage_event_kml'
    ),
    url(
        r'^geoimage/(?P<scenario_type>[db])/(?P<scenario_id>\d+)/' +
        '(?P<slugs>.*)/',
        views.GeoImageKML.as_view(),
        name='lizard_damage_geo_image_kml'
    ),
    url(
        r'^geoimage_nolegend/(?P<scenario_type>[db])/(?P<scenario_id>\d+)/' +
        '(?P<slugs>.*)/',
        views.GeoImageNoLegendKML.as_view(),
        name='lizard_damage_geo_image_no_legend_kml'
    ),
    url(
        r'^geoimage_landuse/(?P<scenario_type>[db])/(?P<scenario_id>\d+)/' +
        '(?P<slugs>.*)/',
        views.GeoImageLandUseKML.as_view(),
        name='lizard_damage_geo_image_landuse_kml'
    ),
    url(
        r'^geoimage_height/(?P<scenario_type>[db])/(?P<scenario_id>\d+)/' +
        '(?P<slugs>.*)/',
        views.GeoImageHeightKML.as_view(),
        name='lizard_damage_geo_image_height_kml'
    ),
    url(
        r'^legend_height_(?P<min_height>.*)_(?P<max_height>.*).png$',
        views.LegendHeight.as_view(),
        name='lizard_damage_legend_height'
    ),
    url(
        r'^thank_you/$',
        views.ThankYou.as_view(),
        name='lizard_damage_thank_you'
        ),
    url(
        r'^max_area_exceeded/$',
        views.MaxAreaExceeded.as_view(),
        name='lizard_damage_max_area_exceeded'
        ),
    url(r'^test$',
        TemplateView.as_view(template_name="lizard_damage/openlayers.html"),
        name='lizard_damage_test'),
)
urlpatterns += debugmode_urlpatterns()


####Create model to predict Q(S,A)
import torch
import numpy as np
import random
from random import randint

'''
In: Input,

+Image
-----> 3 x conv layer
-----> Dense Layer

Out: 3 x Outputs, one for each option
'''

class SnakeModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        #general use
        self.relu = torch.nn.ReLU()
        self.tanh = torch.nn.Tanh()
        ######Convolutional Layers
        self.conv1 = torch.nn.Conv2d(2,16, kernel_size = 2, stride = 1,padding=(0,0)); torch.nn.init.xavier_uniform_(self.conv1.weight) #19
        self.conv2 = torch.nn.Conv2d(16,32, kernel_size = 3, stride = 2,padding=(1,1)); torch.nn.init.xavier_uniform_(self.conv2.weight) #10
        self.conv3 = torch.nn.Conv2d(32,64, kernel_size = 3, stride = 2,padding=(1,1)); torch.nn.init.xavier_uniform_(self.conv3.weight) #5
        self.flatten = torch.nn.Flatten()

        ######Dense Layers
        self.DenseLayer1 = torch.nn.Linear(1600,3); torch.nn.init.xavier_uniform_(self.DenseLayer1.weight)

    def forward(self,image): #direction removed
        #conv layer
        x1 = self.conv1(image); x1=self.relu(x1)
        x2 = self.conv2(x1); x2=self.relu(x2)
        x3 = self.conv3(x2); x3=self.relu(x3)
        x4 = self.flatten(x3)

        #dense
        output = self.DenseLayer1(x4); output = self.tanh(output)
        return output

    def pick_move(self, image, mode='tiny_epsilon',epsilon=0.2):
        '''
        In: 1x2x20x20
        Out: 0,1,2
        '''
        if mode=='tiny_epsilon':
          random_or_not = np.random.uniform(0,1)
          if random_or_not < epsilon:
              return randint(0,2)
          image = image.reshape(1,2,20,20)
          Qvalues = self.forward(image)
          return torch.argmax(Qvalues,dim=1).numpy()[0]


from dotenv import load_dotenv
import pymongo
import os
import time
import json
from pathlib import Path
from pprint import pprint

class DBclient:
    def __init__(self, conn_str):
        self.conn_str = conn_str
        self.connect()

    def connect(self):
        self.client = pymongo.MongoClient(self.conn_str)
        self.db = self.client.websecradar
        print("connected")

    def getHashPairs(self, hash_pairs_n=10, sleepInterval=5, batch_size=10,
     filename='hash_pairs.json', docs_from=0, docs_to=None):
        print("getting pairs")
        url_col = self.db.crawled_data_urls_v0

        hash_pairs = []
        skipped = 0
        total_document_count = 0
        req_count = batch_size

        count = 1956352 # hardkodirano jer dinamičko određivanje predugo traje
        print("document count is: {}".format(count))

        assert (count >= docs_from)
        if docs_to:
            assert (count >= docs_to)
        else:
            docs_to = count

        for document in url_col.find(batch_size=batch_size)[docs_from:docs_to]:

            total_document_count += 1

            to_append = dict()

            last_hash = ""
            second_last_hash = ""
            for check in reversed(document['checks']):
                if check['timestamp'] == document['last_check']:
                    last_hash = check['hash']
                if check['hash'] != last_hash and last_hash != "" and check['hash'] != None:
                    second_last_hash = check['hash']
                    break

            if (second_last_hash == ""):
                skipped += 1
            else:
                if last_hash == None:
                    print("one site's last check doesn't have a hash")
                else:
                    to_append['second_last_hash'] = second_last_hash
                    to_append['last_hash'] = last_hash
                    to_append['url'] = document['url']
                    hash_pairs.append(to_append)

            req_count = req_count - 1

            if req_count <= 0:
                req_count = batch_size
                if hash_pairs_n:
                    if len(hash_pairs) >= hash_pairs_n:
                        break

                print (" --- skipped {} documents".format(skipped))
                print ("hash_pairs: {}".format(len(hash_pairs)))
                print ("total documents viewed: " + str(total_document_count))
                print("sleep {} seconds\n".format(sleepInterval))
                time.sleep(sleepInterval)

                skipped = 0
                continue

        print("Got {} hash pairs".format(len(hash_pairs)))
        print ("writing hash pairs to '{}'".format(filename))
        pprint (hash_pairs)

        old_pairs = list()
        if os.path.isfile(filename):
            with open(filename, 'r')as f:
                old_pairs = json.loads(f.read())

        for pair in old_pairs:
            if pair not in hash_pairs:
                hash_pairs.append(pair)

        written = 0
        with open(filename, 'w') as f:
            json.dump(hash_pairs, f)
            print(" --- writing complete")
            written = 1

        if not written:
            print("write failed")

    def getDocs(self, directory_name, hash_filename):
        print("getting docs")
        doc_col = self.db.crawled_data_pages_v0

        extra_backSlash = "/" if directory_name[-1] != "/" else ""
        directory = Path("./" + directory_name + extra_backSlash)

        if not directory.exists():
            raise "directory doesn't exist"

        count = 0
        stats = dict()
        with open(hash_filename, "r") as h:
            data_str = h.read()
            data = json.loads(data_str)
            for pair in data:
                print (" --- " + str(count) + " / " + str(len(data)))

                print("sleeping 1 sec")
                time.sleep(1)

                old_doc_doc = doc_col.find_one({"hash": pair['second_last_hash']})
                new_doc_doc = doc_col.find_one({"hash": pair['last_hash']})

                if not old_doc_doc or not new_doc_doc:
                    print("couldn't get documents")
                    count += 1
                    continue

                old_doc = old_doc_doc['page']
                new_doc = new_doc_doc['page']

                old_doc_path = Path("documents3") / (str(count) + "_old.html")
                with old_doc_path.open(mode="w", encoding="utf-8") as f:
                    f.write(old_doc)
                    print("old document written")

                new_doc_path = Path("documents3") / (str(count) + "_new.html")
                with new_doc_path.open(mode="w", encoding="utf-8") as f:
                    f.write(new_doc)
                    print("new document written")

                pair['name_of_download'] = str(count)
                stats[str(count)] = pair

                count += 1

        stats_path = Path("logs") / "docName_url.json"
        with stats_path.open(mode="w", encoding="utf-8") as f:
            json.dump(stats, f)
            print("stats written")

def main():

    config = load_dotenv('.env')
    conn_string = os.getenv('DB_CONNECTION_STRING')

    client = DBclient(conn_string)

    batch_size = 10
    sleepInterval = 3
    num_of_hash_pairs_to_get = None # ako je None, broj dokumenata je jednak docs_to - docs_from
    docs_from = 0 # od koje pozicije pocinje dohvacati linkove
    docs_to = 200 # do koje pozicije pocinje dohvacati, ako je None, trazi se do kraja kolekcije

    hash_filename = Path("logs") / "hash_pairs.json"
    directory_name = "documents3"

    """ client.getHashPairs(hash_pairs_n=num_of_hash_pairs_to_get, sleepInterval=sleepInterval, batch_size=batch_size,
        filename=hash_filename, docs_from=docs_from, docs_to=docs_to) """
    client.getDocs(directory_name, hash_filename)

if __name__ == "__main__":
    main()


import sys
import signal
from distance import Distance
from servo import Servo

def exit_handler(signal, frame):
  print("\n>> exit handler")
  dist.tof.stop_ranging()
  sys.stdout.write("\n")
  sys.exit(0)

signal.signal(signal.SIGINT, exit_handler)

servo = Servo()
dist = Distance(1, 0x29)
dist.start()

def carriage_return():
  sys.stdout.write("\r")
  sys.stdout.flush()

def distance_to_pulse(distance):
  """ distance is cm """
  return int((distance / 70) * 1900 + 500)

def get_distance_cm():
  return dist.tof.get_distance() / 10.0

def move_by_distance(distance_cm):
  pulse = distance_to_pulse(distance_cm)

  if pulse < 500:
    pulse = 500
  elif 2400 < pulse:
    pulse = 2400
  else:
    pulse = pulse

  servo.move(pulse)
  return pulse

def current_value_print():
  sys.stdout.write(u"{:04.1f}cm {} pulse".format(distance_cm, pulse))
  sys.stdout.flush()

while True:
  distance_cm = get_distance_cm()
  pulse = move_by_distance(distance_cm)
  carriage_return()
  current_value_print()


from PyQt5.QtCore import Qt
from PyQt5.QtGui import QCursor, QColor

from qgis.gui import ( QgsMapTool,
                       QgsHighlight )
from qgis.core import ( QgsFields,
                        QgsRectangle,
                        QgsGeometry,
                        QgsFeatureRequest,
                        QgsRenderContext)

from qgis_agri import agriConfig, tr
from qgis_agri.log.logger import QgisLogger as logger
from qgis_agri.gui.layer_util import QGISAgriLayers

#
#-----------------------------------------------------------
class DifferenceSuoliTool(QgsMapTool):

    # --------------------------------------
    #
    # --------------------------------------
    def __init__(self,
                 iface,
                 layer,
                 suoliRefLayers=None,
                 suoliMinArea=agriConfig.TOLERANCES.AREA_MIN_VALID_SUOLO,
                 snapTolerance=agriConfig.TOLERANCES.SNAP_TOLERANCE,
                 suoliSnapLayers=None,
                 msgTitle='QGIS Agri'):

        """ Constructor """
        super().__init__( iface.mapCanvas() )
        self.name = 'DifferenceSuoliTool'
        self.iface = iface
        self.canvas = iface.mapCanvas()
        self.layer = layer
        self.suoliRefLayers = suoliRefLayers
        self.suoliMinArea = suoliMinArea
        self.snapTolerance = snapTolerance
        self.suoliSnapLayers = suoliSnapLayers
        self.highlightGeoms = []
        self.cursor = QCursor(Qt.CrossCursor)
        self.msgTitle = msgTitle
        self.initialize( layer, suoliRefLayers )

    # --------------------------------------
    #
    # --------------------------------------
    def initialize(self, layer, suoliRefLayers=None):
        self.layer = layer
        self.suoliRefLayers = suoliRefLayers

    # --------------------------------------
    #
    # --------------------------------------
    def activate(self):
        super().activate()
        self._deselect_all()
        self.canvas.setCursor(self.cursor)

    # --------------------------------------
    #
    # --------------------------------------
    def deactivate(self):
        super().deactivate()
        self._deselect_all()

    # --------------------------------------
    #
    # --------------------------------------
    def canvasPressEvent(self, e):
        """"Manage selection when clicking on the canvas."""

        self.deleteHighlightGeoms()

        ##if e.button() != Qt.LeftButton:
        ##    return

        # get selected feature
        selFeature = self._select_feature_by_point( e.pos() )
        if selFeature is None:
            return

        # update feature geometry by diffeerence
        self._doDifference( selFeature )

    # --------------------------------------
    #
    # --------------------------------------
    def canvasMoveEvent(self, e):
        self.deleteHighlightGeoms()
        # get feature under cursor
        feature = self._select_feature_by_point( e.pos() )
        if feature is None:
            return
        # highlight feature geometry
        self.highlightGeom( feature.geometry() )

    # --------------------------------------
    #
    # --------------------------------------
    def canvasReleaseEvent(self, e):
        pass

    # --------------------------------------
    #
    # --------------------------------------
    def keyPressEvent(self, e):
        pass

    # --------------------------------------
    #
    # --------------------------------------
    def keyReleaseEvent(self, e):
        pass

    # --------------------------------------
    #
    # --------------------------------------
    def isZoomTool(self):
        return False

    # --------------------------------------
    #
    # --------------------------------------
    def isTransient(self):
        return False

    # --------------------------------------
    #
    # --------------------------------------
    def isEditTool(self):
        return True

    # --------------------------------------
    #
    # --------------------------------------
    def highlightGeom(self, geometry, color=None, fillColor=None, width=1):
        # init
        color = color or QColor( 255,0,0,255 )
        fillColor = fillColor or self.canvas.selectionColor() or QColor( 255,0,0,100 )
        fillColor.setAlpha(100)
        # create highlight object
        h = QgsHighlight( self.canvas, geometry, self.layer )
        h.setColor( color )
        h.setWidth( width )
        h.setFillColor( fillColor )
        self.highlightGeoms.append( h )

    # --------------------------------------
    #
    # --------------------------------------
    def deleteHighlightGeoms(self):
        for h in self.highlightGeoms:
            self.canvas.scene().removeItem( h )
        self.highlightGeoms = []

    # --------------------------------------
    #
    # --------------------------------------
    def _deselect_all(self):
        self.deleteHighlightGeoms()
        if self.iface.mapCanvas().isCachingEnabled():
            self.layer.triggerRepaint()
        else:
            self.iface.mapCanvas().refresh()

    # --------------------------------------
    #
    # --------------------------------------
    def _select_feature_by_point(self, po):
        """Get feature by point"""

        # init
        x = po.x()
        y = po.y()

        canvas = self.canvas
        coo_tranf = canvas.getCoordinateTransform()
        w = canvas.mapUnitsPerPixel() * 3
        rect = QgsRectangle(
            coo_tranf.toMapCoordinates(x-w, y-w),
            coo_tranf.toMapCoordinates(x+w, y+w)
        )
        rect.normalize()

        featRequest = QgsFeatureRequest()\
                             .setFilterRect( rect )\
                             .setFlags( QgsFeatureRequest.ExactIntersect | QgsFeatureRequest.NoGeometry )\
                             .setNoAttributes()

        ##if self.selectionFilter is not None:
        ##    featRequest.setFilterExpression( str( self.selectionFilter ) )

        # get features
        selFeature = None
        featuresIter = self.layer.getFeatures( featRequest )
        featuresId = [ feature.id() for feature in featuresIter ]
        featuresIter.close()
        featuresId.reverse()

        # filter not visible features by style
        renderer = self.layer.renderer().clone()
        ctx = QgsRenderContext()
        renderer.startRender( ctx, QgsFields() )

        for fid in featuresId:
            feature = self.layer.getFeature( fid )
            if renderer.willRenderFeature( feature, ctx ):
                selFeature = feature
                break

        renderer.stopRender( ctx )

        return selFeature

    # --------------------------------------
    #
    # --------------------------------------
    def _doDifference(self, inFeature):
        """ Update feature geometry by difference with others """

        # init
        inGeometry = QgsGeometry( inFeature.geometry() )

        # list all difference layers
        cutLayers = [ self.layer ]
        cutLayers.extend(x for x in self.suoliRefLayers if x not in cutLayers)

        # loop over all difference layers
        nFeatFound = 0
        for layer in cutLayers:
            # get features that insersect the input feature
            outFeats_ids = QGISAgriLayers.getFeaturesIdByGeom( layer,
                                                               inGeometry,
                                                               expression=None )

            # loop features for intersection\difference
            for fid in outFeats_ids:
                if fid == inFeature.id():
                    continue

                # get feature geometry
                feat = layer.getFeature( fid )
                featGeom = QgsGeometry( feat.geometry() )

                # check if there are intersections
                if not inGeometry.intersects( featGeom ):
                    continue

                # create difference feature
                inGeometry = inGeometry.difference( featGeom )
                if inGeometry.isEmpty():
                    logger.msgbar( logger.Level.Warning,
                           tr( "Il suolo selezionato tende ad annullarsi: controllare le sovrapposizioni di geometrie" ),
                           title=self.msgTitle )
                    return

                # increment counter
                nFeatFound += 1

        # check if difference found
        if nFeatFound == 0:
            return

        # start edit command
        self.layer.beginEditCommand( 'DifferenceSuoliTool' )
        try:
            # update feature geometry
            inFeature.setGeometry( inGeometry )

            # repair geometry
            QGISAgriLayers.repair_feature(self.layer,
                                           inFeature,
                                           attr_dict=None,
                                           splitMultiParts=True,
                                           autoSnap=False,
                                           suoliMinArea=self.suoliMinArea,
                                           suoliSnapTolerance=self.snapTolerance,
                                           suoliSnapLayerUris=None,
                                           addTopologicalPoints=True)

            # end edit command
            self.layer.endEditCommand()

        except Exception as e:
            self.layer.destroyEditCommand()
            raise e


# the function where all conversions are gathered in one place
def allconversion(number):
    binary = bin(number)
    octal = oct(number)
    hexadecimal = hex(number)
    print("Binary representation: ",binary)
    print("Octal representation: ",octal)
    print("Hexadecimal representation: ",hexadecimal)

print("\nConvert Decimal Number to Binary or Octal or Hexadecimal or All-in-one")

"""
Binary is a numeric system that only uses two digits '0' and '1'

Octal is a numeric system that uses eight digits (0-7)

Hexadecimal is a numeric system that uses the decimal numbers(base 10)
and siz extra symbols(a,b,c,d,e,f).
"""

while True:
    try:
        decimal_number = int(input("Enter Decimal Number: "))

        if decimal_number != float:
            print("Decimal Number Entered Successfully!")
            break;
    except ValueError:
        # We want to from the user enter an integer value
        print("Please provide an integer value")

converted_number = 0
while True:
    try:
        convert_to = int(input("\nConvert into: [2] Binary, [8] Octal, [16] Hexadecimal, [1]All: \n Please enter the corresponding numeric value: "))
        # Want to user enter the numeric values of corresponding the convertion name
        if convert_to == 2:
            converted_number = bin(decimal_number)
            print("Converting Decimal to Binary: ",converted_number)
            break;
        elif convert_to == 8:
            converted_number = oct(decimal_number)
            print("Converting Decimal to Octal: ",converted_number)
            break;
        elif convert_to == 16:
            converted_number = hex(decimal_number)
            print("Converting Decimal to Hexadecimal: ",converted_number)
            break;
        elif convert_to == 1:
            converted_number = allconversion(decimal_number)
            break;
        else:
            print("Please choose your conversion type from the list")
    except:
        continue


import core.hangman as hangman
from core.hangman_storage import HangmanStorage

STORAGE = "./high_scores"

def main():
    intro()
    state = game_loop()
    ending(state)

def intro():
    print("Welcome to Hangman!")
    high_score = read_high_score()
    print(f"Try to beat the high score of {high_score} points!")

def game_loop():
    state = hangman.start_game()
    while not state.is_finished:
        print_state(state)
        guess = prompt_for_guess()
        state = hangman.take_turn(state, guess)
        print_was_guess_correct(state)
    return state

def print_state(state):
    formatted_current_guess = hangman.format_current_known(state.current_known)
    print(f"Lives remaining: {state.lives_left}")
    print(f"Word: {formatted_current_guess}")
    print()

def prompt_for_guess():
    return input("Guess a letter: ")

def print_was_guess_correct(state):
    if state.was_last_guess_correct:
        print("Nice guess!")
    else:
        print("Ah, it's close, but it's not the one.")
        print("You lose a life.")

def ending(state):
    if state.has_won:
        print("Well done!")
        score = hangman.calculate_score(state)
        print(f"You scored {score} points!")
        if update_high_score(score):
            print("That's a new high score!!")
    else:
        print("Sorry, better luck next time.")

def read_high_score():
    with open(STORAGE, "a+") as file:
        storage = HangmanStorage(file)
        return storage.get_high_score()

def update_high_score(new_score):
    with open(STORAGE, "r+") as file:
        storage = HangmanStorage(file)
        return storage.update_high_score_if_lower_than(new_score)

if __name__ == "__main__":
    main()


# Visualizes the Mandelbrot set
# Author: David Niblick
# 06JUN2021

import numpy as np
from matplotlib import pyplot as plt
import cv2
import time

# center_point = (-.1, .93)
center_point = (-.7499, .0314)
row_pixels = 1600
column_pixels = 3840
scale_pixels = .000001  # this is distance betweeen each pixel
# XMIN, XMAX, YMIN, YMAX = (-.12, -.08, .92, .96)
# ROW, COL = (3840, 1600)
NUM_TICKS = 10
ITERATIONS = 250
BAILOUT = 4

def define_box(scale, center_x, center_y, num_rows, num_columns):
    x_left = center_x - (scale * num_columns / 2)
    x_right = center_x + (scale * num_columns / 2)
    y_down = center_y - (scale * num_rows / 2)
    y_up = center_y + (scale * num_rows / 2)
    return x_left, x_right, y_down, y_up

def linear_mapping(location, old_low, old_high, new_low, new_high):
    return (location - old_low) / (old_high - old_low) * (new_high - new_low) + new_low

def sample_mandelbrot(img_array, x_low, x_high, y_low, y_high):
    y_axis_pixels, x_axis_pixels = np.shape(img_array)
    for m in range(0, x_axis_pixels):
        print('\rImage {0:2.0f}% complete'.format(100 * m / x_axis_pixels), end='')
        for n in range(0, y_axis_pixels):
            count = 0
            z = 0
            x_val = linear_mapping(m, 0, x_axis_pixels, x_low, x_high)
            y_val = linear_mapping(n, 0, y_axis_pixels, y_low, y_high)
            c = x_val + y_val*1.j
            while count < ITERATIONS and abs(z) < BAILOUT:
                count += 1
                z = z**2 + c
            img_array[n, m] = count

def display_image(img_array, x_min, x_max, y_min, y_max, num_ticks=10):
    num_rows, num_cols = np.shape(img_array)
    fig = plt.figure()
    ax = fig.add_subplot(1, 1, 1)
    # ax.set_xticks(np.arange(0, num_cols, num_cols / num_ticks))
    # ax.set_yticks(np.arange(0, num_rows, num_rows / num_ticks))
    # xlabels = np.arange(x_min, x_max, (x_max - x_min) / num_ticks)[0:num_ticks]
    # ax.set_xticklabels([str(round(float(xlabel), 5)) for xlabel in xlabels], rotation=45, ha='right')
    # ylabels = np.arange(y_min, y_max, (y_max - y_min) / num_ticks)[0:num_ticks]
    # ax.set_yticklabels([str(round(float(ylabel), 5)) for ylabel in ylabels])
    ax.set_axis_off()
    im = ax.imshow(np.log(img_array), cmap="gray", interpolation="bilinear")
    plt.show()

def img_post_process(img_array):
    hist, discard = np.histogram(img_array, bins=range(256))
    # check and invert to make img dark
    low_color_count = np.sum(hist[:128])
    high_color_count = np.sum(hist[129:])
    if high_color_count > low_color_count:
        img_array = 256 - img_array

    # Drop floor to darkest color
    img_array = img_array - np.min(img_array)
    return img_array.astype(int)

if __name__ == '__main__':
    mandelbrot_img = np.zeros([row_pixels, column_pixels])
    box = define_box(scale_pixels, center_point[0], center_point[1], row_pixels, column_pixels)
    print('box limits (x limits, y limits): ', box)
    sample_mandelbrot(mandelbrot_img, box[0], box[1], box[2], box[3])
    new_mandelbrot_img = img_post_process(mandelbrot_img)
    display_image(new_mandelbrot_img, box[0], box[1], box[2], box[3], num_ticks=NUM_TICKS)
    timestamp = time.time()
    print('img number: ', timestamp)
    cv2.imwrite('cv_output_{}.png'.format(timestamp), new_mandelbrot_img, [cv2.IMWRITE_PNG_COMPRESSION, 0])


#!/usr/bin/python3

from os import path

from setuptools import setup

with open(path.join(path.abspath(path.dirname(__file__)), 'README.txt')) as f:
    readme_contents = f.read()

setup(name='gitconflict',
      version='0.1',
      description='Git merge-conflicts checker',
      py_modules=['git_conflict'],
      license='GNU LGPL 2.1',
      classifiers=['Development Status :: 2 - Pre-Alpha',
                   'Environment :: Console',
                   'Operating System :: POSIX :: Linux',
                   'Programming Language :: Python :: 3 :: Only',
                   'Programming Language :: Python :: 3.2',
                   'Topic :: Software Development :: Version Control',
                   'License :: OSI Approved :: GNU Lesser General Public '
                   'License v2 or later (LGPLv2+)'],
      author="Vasily Makarov",
      entry_points={'console_scripts': [
          'git-conflict = git_conflict:test_and_print_first_conflict']},
      keywords='git merge conflict detector git-aflow',
      author_email="einmalfel@gmail.com",
      url="https://github.com/einmalfel/git-aflow",
      install_requires=["thingitwrapper"],
      provides=["gitconflict"],
      long_description=readme_contents,
      )


import requests
from bs4 import BeautifulSoup
import re
import mysql.connector
from sklearn import tree
myask1 = input('Enter 1 for all information, otherwise press enter')
erfan = mysql.connector.connect(
        host='localhost',
        user='root',
        passwd='',
        database=''
)
my_countery = ['us', 'spain', 'italy', 'India', 'france', 'germany', 'turkey', 'russia', 'iran', 'brazil','peru','chile','israel','japan'
                   ,'morocco','kuwait','finland','malaysia','egypt','australia','panama','belgium','ecuador','sweden','pakistan'
                   ,'chile','singapore','belarus','qatar','romania','ukraine','indonesia','denmark']
x2 = []
y2 = []
x = []
y = []
uq = 1
tedadnoghte = 1
for i in my_countery:
    my_list = []
    my_list3 = []
    my_list2 = []

    m = 0
    get = requests.get('https://www.worldometers.info/coronavirus/country/' + i + "/")
    one = BeautifulSoup(get.text, "html.parser")
    two = one.find_all('h1')
    three = one.find_all('div', attrs={'class': 'maincounter-number'})
    for ii in two:
        m = re.sub('\s+', '', ii.text).strip()
        my_list.append(m)
    for b in three:
        l = re.sub('\s+', '', b.text).strip()
        my_list2.append(l)
    g = len(my_list2)
    Cursor = erfan.cursor()

    a = str(my_list[0])
    b = str(my_list2[0])
    x.append(int(b.replace(',','')))
    c = str(my_list2[1])

    d = str(my_list2[2])
    x.append(int(d.replace(',','')))

    x.append(int(c.replace(',', '')))
    x2.append(x[0:2])
    y2.append(x[2])
    print(i, 'cheak')
    print('.' * tedadnoghte)
    Cursor.execute('INSERT INTO database name VALUE(\'%s\',\'%s\',\'%s\',\'%s\')' % (a, b, c, d))
    erfan.commit()
    my_list.clear()
    my_list2.clear()
    uq = uq + 1
    x.clear()
    tedadnoghte += 1
erfan = mysql.connector.connect(
    host='localhost',
    user='root',
    passwd='',
    database=''
    )
Cursor = erfan.cursor()
query = 'SELECT * FROM database name'
Cursor.execute(query)
for (name, coronaviruscases, deaths, recovered) in Cursor:
    a = str(name)
    b = str(coronaviruscases)
    c = str(deaths)
    d = str(recovered)
    if myask1 == '1':
        print('name = ', a, '/', 'corona virus cases = ', b, '/', 'deaths = ', c, '/', 'recovered = ', d)
cif = tree.DecisionTreeClassifier()
cif = cif.fit(x2,y2)
tedad_case_fayal = input('Enter the number of corona patients : ')
tedad_khob_shodegan = input('Enter the number of people who have fully recovered : ')

new_data = [[tedad_case_fayal,tedad_khob_shodegan]]
answer = cif.predict(new_data)
print(answer,'will probably die')
erfan.close()


from __future__ import annotations
from typing import Iterator, NamedTuple, Iterable, Optional
from difflib import SequenceMatcher
import enum

from ezdxf.lldxf.tags import Tags
from ezdxf.lldxf.types import DXFVertex, DXFTag

__all__ = ["diff_tags", "print_diff", "OpCode"]

# https://docs.python.org/3/library/difflib.html

class OpCode(enum.Enum):
    replace = enum.auto()
    delete = enum.auto()
    insert = enum.auto()
    equal = enum.auto()

class Operation(NamedTuple):
    opcode: OpCode
    i1: int
    i2: int
    j1: int
    j2: int

CONVERT = {
    "replace": OpCode.replace,
    "delete": OpCode.delete,
    "insert": OpCode.insert,
    "equal": OpCode.equal,
}

def convert_opcodes(opcodes) -> Iterator[Operation]:
    for tag, i1, i2, j1, j2 in opcodes:
        yield Operation(CONVERT[tag], i1, i2, j1, j2)

def round_tags(tags: Tags, ndigits: int) -> Iterator[DXFTag]:
    for tag in tags:
        if isinstance(tag, DXFVertex):
            yield DXFVertex(tag.code, (round(d, ndigits) for d in tag.value))
        elif isinstance(tag.value, float):
            yield DXFTag(tag.code, round(tag.value, ndigits))  # type: ignore
        else:
            yield tag

def diff_tags(
    a: Tags, b: Tags, ndigits: Optional[int] = None
) -> Iterator[Operation]:
    if ndigits is not None:
        a = Tags(round_tags(a, ndigits))
        b = Tags(round_tags(b, ndigits))

    sequencer = SequenceMatcher(a=a, b=b)
    return convert_opcodes(sequencer.get_opcodes())

def print_diff(a: Tags, b: Tags, operations: Iterable[Operation]):
    t1: DXFTag
    t2: DXFTag
    print()
    for op in operations:
        if op.opcode == OpCode.insert:
            for t1 in b[op.j1 : op.j2]:
                print(f"insert a[{op.i1}:{op.i2}]: {t1}")
        elif op.opcode == OpCode.replace:
            for index, t1 in enumerate(a[op.i1 : op.i2], op.i1):
                print(f"replace a[{index}]: {t1}")
            for index, t2 in enumerate(b[op.j1 : op.j2], op.j1):
                print(f"     by b[{index}]: {t2}")
        elif op.opcode == OpCode.delete:
            for index, t1 in enumerate(a[op.i1 : op.i2], op.i1):
                print(f"delete a[{index}]: {t1}")


import json
import os.path as osp
import os
import numpy as np
import argparse
import wandb
from datetime import datetime
from math import log10
from copy import deepcopy

uncompressed_net_unit = {"choco":  2,
                        "gt_hsgd":  2 * 2,
                        "docom":  2 * 2,
                        "gnsd":  2 * 2,
                        "detag":  2 * 2,
                        "sgd":  2,
                        "csgd": 1,
                        "beer":  2 * 2}
init_full_batch_size = {"docom_aistats_lenet_femnist": 805263,
                        "docom_tmlr_femnist": 805263,
                        "docom_aistats_1layer_mnist": 60000,
                        "pprox_spda_ff_mnist": 60000,
                        "docom_cifar100_resnet20": 60000}

# grad_unit_batches = {"choco": 1, "gt_hsgd": 2, "docom": 2, "gnsd": 1, "detag": 1, "sgd": 1, "csgd": 1, "beer": 1}

def str2bool(v):
    if v.lower() in ("yes", "true", "t", "y", "1"):
        return True
    elif v.lower() in ("no", "false", "f", "n", "0"):
        return False
    else:
        raise argparse.ArgumentTypeError("Boolean value expected.")

def get_time_diff(start_str, end_str):
    start_time = datetime.strptime(start_str, '%Y-%m-%d %H:%M:%S')
    end_time = datetime.strptime(end_str, '%Y-%m-%d %H:%M:%S')
    delta = end_time - start_time
    return delta.total_seconds()

def ls_cmd(ls):
    names = os.popen(ls).read().split("\n")
    names = names[:-1] if names[-1] == "" else names
    return names

def get_num_batches_train_per_device_per_epoch(results_dir):
    with open(osp.join(results_dir, "0", "log-1.json")) as f:
        logs = json.load(f)
        if "num_batches_train_per_device_per_epoch" in logs[0]:
            x = logs[0]["num_batches_train_per_device_per_epoch"]
        else:
            raise ValueError("check here: no number of batches")
    return x

def filter_logs(results_dir, tags_to_log, hypprms, log_eval=False):
    all_logs = []
    nodes_dir = ls_cmd("ls -d {}/*/".format(results_dir))

    for node_dir in nodes_dir:
        json_results = ls_cmd("ls -d {}*.json".format(node_dir))
        num_jsons = len(json_results)
        fil_log = {k: {m: [] for m in tags_to_log[k]} for k in tags_to_log} # one log for one agent

        if log_eval:
            n_points = 20
            log_cps = sorted(list(set([int(10**( i/n_points * log10(hypprms["it"]))) for i in range(n_points+1)])))
            fil_log["its"] = deepcopy(log_cps)
        else:
            fil_log["its"] = [int(hypprms["it"] / hypprms["epochs"] * i) for i in range(int(hypprms["epochs"]) + 1)]

        n_checkpoints = int(hypprms["epochs"] + 1) if not log_eval else n_points+1
        fil_log["network_usage_in_MBtyes"] = [0.0 for _ in range(n_checkpoints)]
        epoch_duration_in_sec = [0.0 for _ in range(n_checkpoints)]

        cp_ptr = 0
        for i in range(1, num_jsons+1): # keep the order of reading jsons
            fn = "{}log-{}.json".format(node_dir, i)
            if not fn in json_results:
                raise ValueError("file not found: {}".format(fn))
            with open(fn) as json_f:
                logs = json.load(json_f)
            for log in logs:
                if "epoch" in log:
                    if log_eval:
                        if log["epoch"] * hypprms["batch_per_epoch"] > log_cps[0]:
                            log_cps.pop(0)
                            cp_ptr += 1
                    else:
                        cp_ptr = int(log["epoch"])

                if log["measurement"] == "timer":
                    if log["event"] in ["load_data", "forward_pass", "backward_pass", "sync_complete"]:
                        epoch_duration_in_sec[cp_ptr] += log["value"]
                    elif log["event"] == "opt_init":
                        epoch_duration_in_sec[0] += log["value"]
                if log["measurement"] != "runtime" and log["measurement"] != "stat":
                    continue

                if "time" in log:
                    if fil_log.get("start_time", None) is None:
                        fil_log["start_time"] = log["time"]
                    else:
                        fil_log["end_time"] = log["time"]

                if "type" in log and log["type"] in tags_to_log:
                    for metric in tags_to_log[log["type"]]:
                        fil_log[log["type"]][metric].append(log[metric])

                if "n_bits_to_transmit" in log:
                    fil_log["network_usage_in_MBtyes"][cp_ptr] += log["n_bits_to_transmit"]
        if "start_time" in fil_log and "end_time" in fil_log:
            fil_log["total_duration_in_second"] = get_time_diff(fil_log["start_time"], fil_log["end_time"])
        fil_log["training_durations"] = np.cumsum(epoch_duration_in_sec[:-1])
        all_logs.append(fil_log)
    return all_logs

def reduce_logs(logs, reducer, tags_to_log, world_size):
    nps = {k: {} for k in reducer.keys()}
    for k in tags_to_log.keys():
        for tag in tags_to_log[k]:
            agg = [log[k][tag] for log in logs]
            nps[k][tag] = np.array(agg)
            assert len(nps[k][tag]) == world_size, "invalid world size found in logs: {}, {}, {}, {}".format(len(nps[k][tag]), k, tag ,nps[k][tag].shape)
    for k in reducer.keys():
        for tag, op in zip(tags_to_log[k], reducer[k]):
            if op == max:
                nps[k][tag] = np.max(nps[k][tag], axis=0)
            elif op == min:
                nps[k][tag] = np.min(nps[k][tag], axis=0)
            elif op == "avg":
                nps[k][tag] = np.mean(nps[k][tag], axis=0)
            else:
                raise ValueError("unsupported op {}".format(op))
    nps["network_usage_in_MBtyes"] = np.cumsum(np.sum([log["network_usage_in_MBtyes"][:-1] for log in logs], axis=0))
    nps["its"] = logs[0]["its"]
    return nps

def hypprms_parsing(name, results_dir, world_size, model_size, **kwargs):
    batch_per_epoch = get_num_batches_train_per_device_per_epoch(results_dir)
    name = name.split("_warmup_epochs-0")[0]
    key_names_type = [
        ("l2", "weight_decay", float),
        ("lr", "learning_rate", float),
        ("it", "it", int),
        ("epochs", "epochs", float),
        ("batchsize", "batchsize", int),
        ("stepsize", "consensus_stepsize", float),
        ("k", "k", float),
        ("qsgd", "quan_bits", int),
        ("seed", "seed", int),
        ("beta", "momentum_beta", float),
        ("rounds", "gossip_rounds", int),
        ("eta", "gossip_eta", float),
        ("lrsch", "lrschedule", str),
        ("lrdecay", "lrdecay", float),
        ("betasch", "betaschedule", str),
        ("betadecay", "betadecay", float),
        ("gammadecay", "gammadecay", float),
        ("initb", "initbatchnum", int),
        ("rho", "rho", float),
        ("gamma", "gamma", float),
        ("alpha", "alpha", float),
        ("edgefrac", "edgefrac", float),
        ("T", "T", int),
        ("B", "B", int),
        ]
    to_catch_keys = [tup[0] for tup in key_names_type]
    hypprms = {"name": name, "batch_per_epoch": batch_per_epoch, "proj_name": kwargs["proj_name"]}
    entries = name.split("_")
    for en in entries:
        if "-" in en:
            spl = en.split("-")
            k, v = spl[0], "-".join(spl[1:])
            if k in to_catch_keys:
                hypprms[k] = v
    for k, new_k, t in key_names_type:
        if k in hypprms:
            tmp = hypprms[k]
            hypprms.pop(k)
            hypprms[new_k] = t(tmp)
    return hypprms

def extra_props(hypprms, iterations, world_size, model_size):
    comp_k = hypprms.get("k", 1)
    if "choco" in hypprms["name"]:
        if "qsgd" in hypprms:
            net_use = np.cumsum([uncompressed_net_unit["choco"] * world_size * (64 + model_size * (1 + hypprms["quan_bits"])) for _ in range(int(hypprms["epochs"]))])
        else: # top-k
            net_use = np.cumsum([uncompressed_net_unit["choco"] * comp_k * world_size * model_size *64  for _ in range(int(hypprms["epochs"]))])
        sample_use = np.array([world_size * hypprms["batchsize"] * it for it in iterations])
    elif "gt_hsgd" in hypprms["name"]:
        net_use = np.cumsum([uncompressed_net_unit["gt_hsgd"] * world_size * model_size *64 for _ in range(int(hypprms["epochs"]))])
        first_step_grad = init_full_batch_size[hypprms["proj_name"]]
        sample_use = first_step_grad + np.array([world_size * hypprms["batchsize"] * it for it in iterations])
    elif "docom" in hypprms["name"]:
        net_use = None
        # if "qsgd" in hypprms:
        #     net_use = np.cumsum([uncompressed_net_unit["docom"] * world_size * (64 + model_size * (1 + hypprms["quan_bits"])) for _ in range(int(hypprms["epochs"]))])
        # else: # top-k
        #     net_use = np.cumsum([uncompressed_net_unit["docom"] * comp_k * world_size * model_size *64 for _ in range(int(hypprms["epochs"]))])
        first_step_grad = init_full_batch_size[hypprms["proj_name"]]
        sample_use = first_step_grad + np.array([world_size * hypprms["batchsize"] * it for it in iterations])
    elif "beer" in hypprms["name"]:
        if "qsgd" in hypprms:
            net_use = np.cumsum([uncompressed_net_unit["beer"] * world_size * (64 + model_size * (1 + hypprms["quan_bits"])) for _ in range(int(hypprms["epochs"]))])
        else: # top-k
            net_use = np.cumsum([uncompressed_net_unit["beer"] * comp_k * world_size * model_size*64 for _ in range(int(hypprms["epochs"]))])
        first_step_grad = init_full_batch_size[hypprms["proj_name"]]
        sample_use = first_step_grad + np.array([world_size * hypprms["batchsize"] * it for it in iterations])
    elif "gnsd" in hypprms["name"]:
        net_use = np.cumsum([uncompressed_net_unit["gnsd"] * world_size * model_size*64 for _ in range(int(hypprms["epochs"]))])
        sample_use = np.array([world_size * hypprms["batchsize"] * it for it in iterations])
    elif "detag" in hypprms["name"]:
        net_use = np.cumsum([uncompressed_net_unit["detag"] * world_size * model_size*64 for _ in range(int(hypprms["epochs"]))])
        sample_use = np.array([world_size * hypprms["batchsize"] * it for it in iterations])
    elif "sgd" in hypprms["name"] and "complete" in hypprms["name"]:
        net_use = np.cumsum([uncompressed_net_unit["csgd"]*64 for _ in range(int(hypprms["epochs"]))])
        sample_use = np.array([world_size * hypprms["batchsize"] * it for it in iterations])
    elif "sgd" in hypprms["name"] and not "complete" in hypprms["name"]:
        net_use = np.cumsum([uncompressed_net_unit["sgd"] * world_size * model_size*64 for _ in range(int(hypprms["epochs"]))])
        sample_use = np.array([world_size * hypprms["batchsize"] * it for it in iterations])
    elif "spda" in hypprms["name"] or "di_cs" in hypprms["name"] or "diging" in hypprms["name"] or "fsppd" in hypprms["name"]:
        net_use = None
        sample_use = np.array([world_size * hypprms["batchsize"] * it for it in iterations])
    else:
        raise ValueError("Unimplemented extra_props for {}".format(hypprms["name"]))
    if not net_use is None:
        hypprms["net_use"] = net_use
    hypprms["sample_use"] = sample_use

# wandb
def exp_template(exp, seed):
    if "choco" in exp["name"]:
        if "k" in exp:
            comp_name = "top" if "top_k" in exp["name"] else "random"
            name = "CHOCO-{}lr-{}gamma-{}-{}k-{}batchsize-{}l2".format(exp["learning_rate"], exp["consensus_stepsize"], comp_name, exp["k"], exp["batchsize"], exp["weight_decay"])
            hypprm = {"gamma": exp["consensus_stepsize"],
                    "k": exp["k"]}
        else:
            name = "CHOCO-{}lr-{}gamma-{}bit-{}batchsize-{}l2".format(exp["learning_rate"], exp["consensus_stepsize"], exp["quan_bits"], exp["batchsize"], exp["weight_decay"])
            hypprm = {"gamma": exp["consensus_stepsize"],
                    "quan_bits": exp["quan_bits"]}
    elif "gt_hsgd" in exp["name"]:
        name = "GT-HSGD-{}lr-{}beta-{}batchsize-{}initbatch-{}l2".format(exp["learning_rate"], exp["momentum_beta"], exp["batchsize"], exp["initbatchnum"], exp["weight_decay"])
        hypprm = {"beta": exp["momentum_beta"]}
    elif "docom" in exp["name"]:
        if "k" in exp:
            name = "DoCoM-{}lr-{}gamma-{}beta-{}k-{}batchsize-{}l2".format(exp["learning_rate"], exp["consensus_stepsize"], exp["momentum_beta"], exp["k"], exp["batchsize"], exp["weight_decay"])
            hypprm = {"gamma": exp["consensus_stepsize"],
                    "beta": exp["momentum_beta"],
                    "k": exp["k"]}
        else:
            name = "DoCoM-{}lr-{}gamma-{}beta-{}bits-{}batchsize-{}l2".format(exp["learning_rate"], exp["consensus_stepsize"], exp["momentum_beta"], exp["quan_bits"], exp["batchsize"], exp["weight_decay"])
            hypprm = {"gamma": exp["consensus_stepsize"],
                    "beta": exp["momentum_beta"],
                    "quan_bits": exp["quan_bits"]}
    elif "beer" in exp["name"]:
        comp_name = "top" if "top_k" in exp["name"] else "random"
        if "k" in exp:
            name = "BEER-{}lr-{}gamma-{}-{}k-{}batchsize-{}l2".format(exp["learning_rate"], exp["consensus_stepsize"], comp_name, exp["k"], exp["batchsize"], exp["weight_decay"])
            hypprm = {"gamma": exp["consensus_stepsize"],
                        "k": exp["k"]}
        else:
            name = "BEER-{}lr-{}gamma-{}bits-{}batchsize-{}l2".format(exp["learning_rate"], exp["consensus_stepsize"], exp["quan_bits"], exp["batchsize"], exp["weight_decay"])
            hypprm = {"gamma": exp["consensus_stepsize"],
                        "quan_bits": exp["quan_bits"]}
    elif "gnsd" in exp["name"]:
        name = "GNSD-{}lr-{}batchsize".format(exp["learning_rate"], exp["batchsize"])
        hypprm = {}
    elif "diging" in exp["name"]:
        name = "DIGing-{}lr-{}edge".format(exp["learning_rate"], exp["edgefrac"])
        hypprm = {}
    elif "detag" in exp["name"]:
        name = "DeTAG-{}lr-{}R-{}g_eta-{}batchsize".format(exp["learning_rate"], exp["gossip_rounds"], exp["gossip_eta"], exp["batchsize"])
        hypprm = {"gossip_eta": exp["gossip_eta"],
                "gossip_rounds": exp["gossip_rounds"]}
    elif "sgd" in exp["name"]:
        if "complete" in exp["name"]:
            name = "CSGD-{}lr-{}batchsize".format(exp["learning_rate"], exp["batchsize"])
        else:
            name = "DSGD-{}lr-{}batchsize".format(exp["learning_rate"], exp["batchsize"])
        hypprm = {}
    elif "pprox" in exp["name"]:
        name = "PProxSPDA-{}rho-{}gamma-{}alpha-{}edge-{}k".format(exp["rho"], exp["gamma"], exp["alpha"], exp["edgefrac"], exp["k"])
        hypprm = {"rho": exp["rho"], "gamma": exp["gamma"], "alpha": exp["alpha"],
                    "k": exp["k"], "edge_frac": exp["edgefrac"]}
    elif "fully_spda" in exp["name"]:
        name = "FullySPDA-{}gamma-{}alpha-{}edge-{}k".format(exp["gamma"], exp["alpha"], exp["edgefrac"], exp["k"])
        hypprm = {"gamma": exp["gamma"], "alpha": exp["alpha"],
                    "k": exp["k"], "edge_frac": exp["edgefrac"]}
    elif "fsppd" in exp["name"]:
        name = "FSPPD-{}lr-{}gamma-{}batchsize-{}edge-{}k".format(exp["learning_rate"], exp["gamma"], exp["batchsize"], exp["edgefrac"], exp["k"])
        hypprm = {"gamma": exp["gamma"],
                    "k": exp["k"], "edge_frac": exp["edgefrac"]}
    elif "di_cs" in exp["name"]:
        if "di_cs_svrg" in exp["name"]:
            opt_name = "Di-CS-SVRG"
        elif "di_cs_gd" in exp["name"]:
            opt_name = "Di-CS-GD"
        else:
            opt_name = "Di-CS"
        name = opt_name + "-{}alpha-{}gamma-{}T-{}B-{}edge-{}k".format(exp["alpha"], exp["gamma"], exp["T"], exp["B"], exp["edgefrac"], exp["k"])
        hypprm = {"gamma": exp["gamma"], "alpha": exp["alpha"], "T": exp["T"], "B": exp["B"],
                    "k": exp["k"], "edge_frac": exp["edgefrac"]}
    else:
        raise ValueError("Unimplemented template for {}".format(exp["name"]))
    name = "{}-{}seed".format(name, seed)
    if "SAM" in exp["name"]:
        name = "SAM-" + name

    # if "lrschedule" in exp:
    #     name = "{}-{}lrschedule".format(name, exp["lrschedule"])
    # if "lrdecay" in exp:
    #     name = "{}-{}lrdecay".format(name, exp["lrdecay"])
    # if "gammadecay" in exp:
    #     name = "{}-{}gammadecay".format(name, exp["gammadecay"])
    return name, hypprm

def upload_to_wandb(exp, proj_name, entries_count):
    wandb.login()
    exp_name, exp_hypprm = exp_template(exp, exp["seed"])
    hyperparameters = {
                "_name": exp["name"],
                "epochs": exp["epochs"],
                "batch_size": exp["batchsize"],
                "seed": exp["seed"],
                "learning_rate": exp["learning_rate"],
                **exp_hypprm,
                }
    metrics = [("worst_true_grad_train_global", "eval_local_model_on_full_training_data", "grad_norm"),
                ("worst_loss_train_global", "eval_local_model_on_full_training_data", "loss"),
                ("worst_acc_train_global", "eval_local_model_on_full_training_data", "top1"),
                ("worst_loss_test_global", "eval_local_model_on_full_testing_data", "loss"),
                ("worst_acc_test_global", "eval_local_model_on_full_testing_data", "top1"),
                ("consensus_gap", "averaged_model", "distance")]
    log_length = {name: entries_count[a+"_"+b] for name,a,b in metrics}
    max_log_length = max([log_length[n] for n in log_length])
    with wandb.init(project=proj_name, name=exp_name, config=hyperparameters):
        for i in range(max_log_length):
            wandb_log = {
                        "iteration": exp["its"][i],
                        # "worst_true_grad_train_global": exp["eval_local_model_on_full_training_data"]["grad_norm"][i],
                        # "worst_loss_train_global": exp["eval_local_model_on_full_training_data"]["loss"][i],
                        # "worst_acc_train_global": exp["eval_local_model_on_full_training_data"]["top1"][i],
                        # "worst_loss_test_global": exp["eval_local_model_on_full_testing_data"]["loss"][i],
                        # "worst_acc_test_global": exp["eval_local_model_on_full_testing_data"]["top1"][i],
                        # "consensus_gap": exp["averaged_model"]["distance"][i],
                        # "network_bits_transmitted": exp.get("net_use", exp["network_usage_in_MBtyes"] * 8 * 2**20)[i],
                        "network_MBtyes_transmitted": exp["network_usage_in_MBtyes"][i],
                        "sample_accessed":exp["sample_use"][i],
                        "wall_clock_training_time": exp["wall_clock_training_time"][i],
                        "total_program_time": exp["total_program_time"][i]}
            for name, a, b in metrics:
                if i < log_length[name]:
                    wandb_log[name] = exp[a][b][i]
            if "gnsd" in exp["name"] or "detag" in exp["name"] or "docom" in exp["name"] or "gt_hsgd" in exp["name"] or "beer" in exp["name"]:
                wandb_log = {**wandb_log, "grad_tracker_consensus_gap": exp["averaged_gt"]["distance"][i]}
            wandb.log(wandb_log, step=i)

def get_gt_consensus(saved_model_dir, world_size=36):
    gt_cons = [[] for _ in range(world_size)]
    for ag in range(world_size):
        json_results = ls_cmd("ls -d {}/*.json".format(osp.join(saved_model_dir, str(ag))))
        num_jsons = len(json_results)
        for j in range(1, 1+num_jsons):
            with open(osp.join(saved_model_dir, str(ag), "log-{}.json".format(j))) as f:
                logs = json.load(f)
                for log in logs:
                    if "grad_tracker_consensus_dist" in log:
                        gt_cons[ag].append(log["grad_tracker_consensus_dist"])

    gt_cons = np.array(gt_cons)
    gt_cons = np.max(gt_cons, axis=0)
    return gt_cons

def world_size(exp_name):
    return int(exp_name.split("agents_")[-1].split("_")[0])

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--exp_name', type=str, required=True)
    parser.add_argument('--inference_dir', type=str, required=True)
    parser.add_argument("--log_eval", type=str2bool, default=False)
    parser.add_argument('--model_size', type=int, default=60850)
    parser.add_argument("--proj_name", type=str, default="leaf-femnist")
    conf = parser.parse_args()

    print("parsing log of {} and upload to wandb".format(conf.exp_name))
    conf.world_size = world_size(conf.exp_name)

    results_dir = osp.join(conf.inference_dir, conf.exp_name)

    tags_to_log = {"eval_local_model_on_full_training_data": ["loss", "grad_norm", "top1"],
                "eval_local_model_on_full_testing_data": ["loss", "top1"],
                "averaged_model": ["distance"],
                "averaged_gt": ["distance"]}
    reducer = {"eval_local_model_on_full_training_data": [max, max, min],
            "eval_local_model_on_full_testing_data": [max, min],
            "averaged_model": ["avg"],
            "averaged_gt": ["avg"]}
    hypprms = hypprms_parsing(conf.exp_name, results_dir, conf.world_size, conf.model_size, proj_name=conf.proj_name)
    all_logs = filter_logs(results_dir, tags_to_log, hypprms, log_eval=conf.log_eval)
    reduced_logs = reduce_logs(all_logs, reducer, tags_to_log, conf.world_size)
    extra_props(hypprms, reduced_logs["its"], conf.world_size, conf.model_size)
    entries_count = {}
    for k in tags_to_log:
        for metric in tags_to_log[k]:
            print("{} - {} entries: {}".format(k, metric, len(reduced_logs[k][metric])))
            entries_count[k+"_"+metric] = len(reduced_logs[k][metric])

    total_program_time =  [all_logs[0]["total_duration_in_second"] / int(hypprms["epochs"]) * (i+1) for i in range(int(hypprms["epochs"]))]
    reduced_logs = {**reduced_logs, **hypprms, "wall_clock_training_time": all_logs[0]["training_durations"], "total_program_time": total_program_time}
    # reduced_logs
    print("epochs in log: ", len(reduced_logs["eval_local_model_on_full_training_data"]["loss"]))
    upload_to_wandb(reduced_logs, conf.proj_name, entries_count)


import pickle
from random import randint

import torch
from peft import LoraConfig, TaskType, get_peft_model, prepare_model_for_kbit_training
from sklearn.metrics import precision_recall_fscore_support, accuracy_score, auc, precision_recall_curve
from transformers import TrainingArguments, AutoConfig, \
    AutoModelForSequenceClassification, AutoTokenizer, BitsAndBytesConfig, DataCollatorWithPadding

EPOCHS = 6 # TODO: CHANGE THIS
max_length = 4096
output_dir = f"change_me"  # TODO: CHANGE THIS
num_labels = 2
model_id = "meta-llama/Llama-2-13b-hf"
# model_id = "stanford-crfm/BioMedLM"

# MIMIC-IV DATA
# each pickle in the following format: [[[124325,
#   0,
#   ['Mood disorders',
#    'Diabetes mellitus with complications',
#    'Other circulatory disease'],
#   [1, 1, 1]]]
# For more details, see the documentation on GitHub.
train_pickle_file_path = '/sise/home/benshoho/projects/Med-BERT/Fine-Tunning-Tutorials/data/mimic-iv/chronic_kidney_disease_descriptions_train.pickle'  # TODO: CHANGE THIS
validation_pickle_file_path = '/sise/home/benshoho/projects/Med-BERT/Fine-Tunning-Tutorials/data/mimic-iv/chronic_kidney_disease_descriptions_validation.pickle'  # TODO: CHANGE THIS
test_pickle_file_path = '/sise/home/benshoho/projects/Med-BERT/Fine-Tunning-Tutorials/data/mimic-iv/chronic_kidney_disease_descriptions_test.pickle'  # TODO: CHANGE THIS

# set of descriptions for the mimic-iv dataset according the descriptions of the CCS categories.
mimic_iv_description_codes_path = '/sise/home/benshoho/projects/Med-BERT/Pretraining Code/Data Pre-processing Code/mimic_iv_descriptions_set.types'  # TODO: CHANGE THIS
with open(mimic_iv_description_codes_path, 'rb') as file:
    mimic_iv_description_codes = pickle.load(file)

def print_trainable_parameters(model):
    """
    Prints the number of trainable parameters in the model.
    """
    trainable_params = 0
    all_param = 0
    for _, param in model.named_parameters():
        all_param += param.numel()
        if param.requires_grad:
            trainable_params += param.numel()
    print(
        f"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param}"
    )

# Load the tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_id)
tokenizer.pad_token = tokenizer.eos_token

config = LoraConfig(
    r=32,
    lora_alpha=32,
    lora_dropout=0.1,
    bias="none",
    task_type=TaskType.SEQ_CLS
)

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)

model = AutoModelForSequenceClassification.from_pretrained(model_id, quantization_config=bnb_config,
                                                           config=AutoConfig.from_pretrained(model_id,
                                                                                trust_remote_code=True,
                                                                                num_labels=num_labels),
                                                           trust_remote_code=True)

model.gradient_checkpointing_enable()
model = prepare_model_for_kbit_training(model)

print_trainable_parameters(model)
model = get_peft_model(model, config)
print_trainable_parameters(model)

model.resize_token_embeddings(len(tokenizer))
model.config.pad_token_id = model.config.eos_token_id

def compute_metrics(p):
    p.predictions = torch.from_numpy(p.predictions)
    p.predictions = torch.softmax(p.predictions, dim=1)

    preds = np.argmax(p.predictions.cpu().numpy(), axis=1)
    labels = p.label_ids
    accuracy = accuracy_score(labels, preds)

    precision, recall, f1, _ = precision_recall_fscore_support(labels, preds, average=None)
    precision1, recall1, thresholds = precision_recall_curve(labels, p.predictions[:, 1], pos_label=1)
    auc_precision_recall = auc(recall1, precision1)
    positive_confidences = p.predictions[:, 1]

    return {
        'accuracy': accuracy,
        'precision_class_0': precision[0],
        'precision_class_1': precision[1],
        'recall_class_0': recall[0],
        'recall_class_1': recall[1],
        'f1_class_0': f1[0],
        'f1_class_1': f1[1],
        'aucpr': auc_precision_recall,
        'pos_scores': positive_confidences.tolist(),
        'true_labels': labels.tolist(),
    }

import pickle
from datasets import Dataset

class CustomDataset:
    def __init__(self, pickle_file_path):
        with open(pickle_file_path, "rb") as file:
            data = pickle.load(file)

        # Create empty lists to store the dataset
        patient_ids = []
        labels = []
        diagnoses = []
        visits_num = []
        # Extract data from the loaded pickle
        for item in data:
            patient_ids.append(item[0])
            labels.append(item[1])
            diagnoses_list = item[2]
            diagnoses.append(diagnoses_list)
            visits_num.append(item[3])

        self.data_dict = {
            "patient_id": patient_ids,
            "classification_label": labels,
            "diagnoses": diagnoses,
            "visits_num": visits_num
        }

    def to_dict(self):
        return self.data_dict

dataset_type = 'mimic-iv' if 'mimic' in train_pickle_file_path else 'eicu-crd'
print(f'dataset_type={dataset_type}')

train_dataset = CustomDataset(train_pickle_file_path)
test_dataset = CustomDataset(test_pickle_file_path)
validation_dataset = CustomDataset(validation_pickle_file_path)

train_dict = train_dataset.to_dict()
test_dict = test_dataset.to_dict()
validation_dict = validation_dataset.to_dict()

train_dataset = Dataset.from_dict(train_dict)
test_dataset = Dataset.from_dict(test_dict)
validation_dataset = Dataset.from_dict(validation_dict)

def add_mimic_new_tokens_from_diagnosis_strings():
    new_tokens = set()
    for diagnosis_str in mimic_iv_description_codes:
        new_tokens.add(diagnosis_str)

    tokens_to_add = list(new_tokens - set(tokenizer.get_vocab()))
    prev_num_of_tokens = len(tokenizer)
    tokenizer.add_tokens(tokens_to_add)
    tokenizer.add_tokens(['empty_pad'], special_tokens=True)
    new_num_of_tokens = len(tokenizer)
    model.resize_token_embeddings(len(tokenizer))
    print(f'first 20 added tokens are= {list(new_tokens)[:20]}')
    print(f'prev tokens number= {prev_num_of_tokens}, new tokens number= {new_num_of_tokens}')

# If you want to use eicu-crd and add tokens to the vocab of pretrained tokenizer, use this:

# def add_eicu_crd_new_tokens_from_diagnosis_strings():
#     new_tokens = set()
#     for diagnosis_str in icd_code_to_icd_string_dict.values():
#         diagnoses_part_of_diagnosis_str = diagnosis_str.split('|')
#         # for words:
#         # for diagnoses_part in diagnoses_part_of_diagnosis_str:
#         #     new_tokens = new_tokens.union(diagnoses_part.split(" "))
#         for diagnoses_part in diagnoses_part_of_diagnosis_str:
#             new_tokens.add(diagnoses_part)
#     # consider to remove special tokens such ( etc..

#     tokens_to_add = list(new_tokens - set(tokenizer.get_vocab()))
#     prev_num_of_tokens = len(tokenizer)
#     tokenizer.add_tokens(tokens_to_add)
#     tokenizer.add_tokens(['empty_pad'], special_tokens=True)
#     new_num_of_tokens = len(tokenizer)
#     model.resize_token_embeddings(len(tokenizer))
#     print(f'first 20 added tokens are= {list(new_tokens)[:20]}')
#     print(f'prev tokens number= {prev_num_of_tokens}, new tokens number= {new_num_of_tokens}')

add_mimic_new_tokens_from_diagnosis_strings()
# add_eicu_crd_new_tokens_from_diagnosis_strings()

prompt_template = """
Your task is to determine whether a patient is likely to have a specific Chronic kidney disease based on their diagnosis descriptions provided below.
Each diagnosis description is separated by a comma.

**Patient Diagnosis Descriptions:**

{diagnoses}
"""

def template_dataset(sample):
    sample["text"] = prompt_template.format(diagnoses=sample["diagnoses"],
                                            eos_token=tokenizer.eos_token)
    return sample

# apply prompt template per sample
train_dataset = train_dataset.map(template_dataset)
print(f'example sample from train:\n {train_dataset[0]}')
validation_dataset = validation_dataset.map(template_dataset)

print(train_dataset[randint(0, len(train_dataset))]["diagnoses"])

# apply prompt template per sample
test_dataset = test_dataset.map(template_dataset)
lm_train_dataset = train_dataset.map(
    lambda sample: {
        "input_ids": tokenizer(sample["text"], truncation=True, max_length=max_length).input_ids,
        "attention_mask": tokenizer(sample["text"], truncation=True, max_length=max_length).attention_mask,
        "labels": sample["classification_label"]
    },
    batched=True,
    batch_size=64,
    remove_columns=list(train_dataset.features)
)
lm_validation_dataset = validation_dataset.map(
    lambda sample: {
        "input_ids": tokenizer(sample["text"], truncation=True, max_length=max_length).input_ids,
        "attention_mask": tokenizer(sample["text"], truncation=True, max_length=max_length).attention_mask,
        "labels": sample["classification_label"]
    },
    batched=True,
    batch_size=64,
    remove_columns=list(validation_dataset.features)
)
lm_test_dataset = test_dataset.map(
    lambda sample: {
        "input_ids": tokenizer(sample["text"], truncation=True, max_length=max_length).input_ids,
        "attention_mask": tokenizer(sample["text"], truncation=True, max_length=max_length).attention_mask,
        "labels": sample["classification_label"]
    },
    batched=True,
    batch_size=64,
    remove_columns=list(test_dataset.features)
)

from sklearn.utils.class_weight import compute_class_weight
import numpy as np

labels = lm_train_dataset['labels']

# Calculate class weights for an imbalanced dataset
class_weights = compute_class_weight('balanced', classes=np.unique(labels), y=labels)
class_weights = torch.tensor(class_weights, dtype=torch.float)
print(f'class_weights= {class_weights}')

print(f"Train dataset size: {len(train_dataset)}")
print(f"Val dataset size: {len(validation_dataset)}")
print(f"Test dataset size: {len(test_dataset)}")

from transformers import Trainer

training_args = TrainingArguments(
    output_dir=output_dir,
    evaluation_strategy='steps',
    save_strategy='steps',
    eval_steps=1000,
    save_steps=1000,
    num_train_epochs=EPOCHS,
    per_device_train_batch_size=8,
    per_device_eval_batch_size=4,
    auto_find_batch_size=True,
    logging_steps=30,
    learning_rate=2e-5,
    optim="adamw_torch",
    save_total_limit=20,
    logging_dir='./logs',
    load_best_model_at_end=True,
    metric_for_best_model="eval_aucpr",
    greater_is_better=True,
    dataloader_num_workers=8,
)

model.config.use_cache = False
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=lm_train_dataset,
    eval_dataset=lm_validation_dataset,
    compute_metrics=compute_metrics,
    data_collator=DataCollatorWithPadding(tokenizer=tokenizer)
)

trainer.train()
trainer.save_model(output_dir)

test_results = trainer.evaluate(eval_dataset=lm_test_dataset)

print(f'see outputs in= {output_dir}')
trainer.save_metrics("test", test_results)