import re
from datetime import datetime, timedelta
from functools import cache, wraps
from typing import Iterable

from cymple.builder import NodeAfterMergeAvailable, NodeAvailable, QueryBuilder

from ...model import (
    MatchStrategy,
    Node,
    Relationship,
    RelationshipIdentityShape,
    RelationshipWithNodes,
    TimeToLiveConfiguration,
)
from ...schema.schema import GraphObjectType
from ..query_executor import OperationOnNodeIdentity, OperationOnRelationshipIdentity
from .query import Query, QueryBatch

PROPERTIES_PARAM_NAME = "properties"
ADDITIONAL_LABELS_PARAM_NAME = "additional_labels"
GENERIC_NODE_REF_NAME = "node"
FROM_NODE_REF_NAME = "from_node"
TO_NODE_REF_NAME = "to_node"
RELATIONSHIP_REF_NAME = "rel"
PARAMETER_CORRECTION_REGEX = re.compile(r"\"(params.__\w+)\"")
DELETE_NODE_QUERY = "MATCH (n) WHERE id(n) = id DETACH DELETE n"
DELETE_REL_QUERY = "MATCH ()-[r]->() WHERE id(r) = id DELETE r"

def correct_parameters(f):
    @wraps(f)
    def wrapper(*args, **kwargs):
        query = f(*args, **kwargs)
        return PARAMETER_CORRECTION_REGEX.sub(r"\1", query)

    return wrapper

def generate_prefixed_param_name(property_name: str, prefix: str) -> str:
    return f"__{prefix}_{property_name}"

def generate_properties_set_with_prefix(properties: Iterable[str], prefix: str):
    return {
        prop: f"params.{generate_prefixed_param_name(prop, prefix)}"
        for prop in properties
    }

def generate_where_set_with_prefix(properties: frozenset, prefix: str):
    return {
        f"{prefix}.{prop}": f"params.{generate_prefixed_param_name(prop, prefix)}"
        for prop in properties
    }

@cache
def _match_node(
    node_operation: OperationOnNodeIdentity, name=GENERIC_NODE_REF_NAME
) -> NodeAvailable:
    op = "=~" if node_operation.match_strategy == MatchStrategy.FUZZY else "="
    identity = node_operation.node_identity
    props = generate_where_set_with_prefix(identity.keys, name)
    return (
        QueryBuilder()
        .match()
        .node(labels=identity.type, ref_name=name)
        .where_multiple(props, comparison_operator=op)
    )

@cache
def _merge_node(
    node_operation: OperationOnNodeIdentity, name=GENERIC_NODE_REF_NAME
) -> NodeAfterMergeAvailable:
    properties = generate_properties_set_with_prefix(
        node_operation.node_identity.keys, name
    )
    return (
        QueryBuilder()
        .merge()
        .node(
            labels=node_operation.node_identity.type,
            ref_name=name,
            properties=properties,
        )
    )

@cache
def _merge_relationship(rel_identity: RelationshipIdentityShape):
    keys = generate_properties_set_with_prefix(rel_identity.keys, RELATIONSHIP_REF_NAME)
    match_rel_query = (
        QueryBuilder()
        .match_optional()
        .node(ref_name=FROM_NODE_REF_NAME)
        .related_to(
            ref_name=RELATIONSHIP_REF_NAME,
            properties=keys,
            label=rel_identity.type,
        )
        .node(ref_name=TO_NODE_REF_NAME)
    )
    create_rel_query = str(match_rel_query).replace("OPTIONAL MATCH ", "CREATE")
    set_properties_query = f"SET {RELATIONSHIP_REF_NAME} += params.{generate_prefixed_param_name(PROPERTIES_PARAM_NAME, RELATIONSHIP_REF_NAME)}"

    return f"""
    {match_rel_query}
    FOREACH (x IN CASE WHEN {RELATIONSHIP_REF_NAME} IS NULL THEN [1] ELSE [] END |
        {create_rel_query} {set_properties_query})
    FOREACH (i in CASE WHEN {RELATIONSHIP_REF_NAME} IS NOT NULL THEN [1] ELSE [] END |
        {set_properties_query})
    """

class Neo4jIngestQueryBuilder:
    @cache
    @correct_parameters
    def generate_update_node_operation_query_statement(
        self,
        operation: OperationOnNodeIdentity,
    ) -> str:
        """Generate a query to update a node in the database given a node type and a match strategy."""

        if operation.match_strategy == MatchStrategy.EAGER:
            query = str(_merge_node(operation))
        else:
            query = str(_match_node(operation))

        result = f"{query} SET {GENERIC_NODE_REF_NAME} += params.{generate_prefixed_param_name(PROPERTIES_PARAM_NAME, GENERIC_NODE_REF_NAME)}"
        if operation.node_identity.additional_types:
            result += f" WITH {GENERIC_NODE_REF_NAME}, params CALL apoc.create.addLabels({GENERIC_NODE_REF_NAME}, params.{generate_prefixed_param_name(ADDITIONAL_LABELS_PARAM_NAME, GENERIC_NODE_REF_NAME)}) yield node RETURN true"
        return result

    def generate_update_node_operation_params(self, node: Node) -> dict:
        """Generate the parameters for a query to update a node in the database."""

        params = self.generate_node_key_params(node)
        params[
            generate_prefixed_param_name(PROPERTIES_PARAM_NAME, GENERIC_NODE_REF_NAME)
        ] = node.properties
        params[
            generate_prefixed_param_name(
                ADDITIONAL_LABELS_PARAM_NAME, GENERIC_NODE_REF_NAME
            )
        ] = node.additional_types

        return params

    def generate_node_key_params(self, node: Node, name=GENERIC_NODE_REF_NAME) -> dict:
        """Generate the parameters for a query to update a node in the database."""

        return {
            generate_prefixed_param_name(k, name): v for k, v in node.key_values.items()
        }

    @cache
    @correct_parameters
    def generate_update_relationship_operation_query_statement(
        self,
        operation: OperationOnRelationshipIdentity,
    ) -> str:
        """Generate a query to update a relationship in the database given a relationship operation."""

        match_from_node_segment = _match_node(operation.from_node, FROM_NODE_REF_NAME)
        match_to_node_segment = _match_node(operation.to_node, TO_NODE_REF_NAME)
        merge_rel_segment = _merge_relationship(operation.relationship_identity)
        return f"{match_from_node_segment} {match_to_node_segment} {merge_rel_segment}"

    def generate_update_rel_params(self, rel: Relationship) -> dict:
        """Generate the parameters for a query to update a relationship in the database."""

        params = {
            generate_prefixed_param_name(k, RELATIONSHIP_REF_NAME): v
            for k, v in rel.key_values.items()
        }
        params[
            generate_prefixed_param_name(PROPERTIES_PARAM_NAME, RELATIONSHIP_REF_NAME)
        ] = rel.properties

        return params

    def generate_update_rel_between_nodes_params(
        self, rel: RelationshipWithNodes
    ) -> dict:
        """Generate the parameters for a query to update a relationship in the database."""

        params = self.generate_update_rel_params(rel.relationship)
        params.update(self.generate_node_key_params(rel.from_node, FROM_NODE_REF_NAME))
        params.update(self.generate_node_key_params(rel.to_node, TO_NODE_REF_NAME))
        return params

    def generate_batch_update_node_operation_batch(
        self,
        operation: OperationOnNodeIdentity,
        nodes: Iterable[Node],
    ) -> QueryBatch:
        """Generate a batch of queries to update nodes in the database in the same way of the same type."""

        query_statement = self.generate_update_node_operation_query_statement(operation)
        params = [self.generate_update_node_operation_params(node) for node in nodes]
        return QueryBatch(query_statement, params)

    def generate_batch_update_relationship_query_batch(
        self,
        operation: OperationOnRelationshipIdentity,
        relationships: Iterable[RelationshipWithNodes],
    ) -> QueryBatch:
        """Generate a batch of queries to update relationships in the database in the same way of the same type."""

        query = self.generate_update_relationship_operation_query_statement(operation)
        params = [
            self.generate_update_rel_between_nodes_params(rel) for rel in relationships
        ]
        return QueryBatch(query, params)

    def generate_ttl_match_query(self, config: TimeToLiveConfiguration) -> Query:
        earliest_allowed_time = datetime.utcnow() - timedelta(
            hours=config.expiry_in_hours
        )
        params = {"earliest_allowed_time": earliest_allowed_time}
        if config.custom_query is not None:
            return Query(config.custom_query, params)

        query_builder = QueryBuilder()
        ref_name = "x"

        if config.graph_object_type == GraphObjectType.NODE:
            query_builder = query_builder.match().node(
                labels=config.object_type, ref_name=ref_name
            )
        else:
            query_builder = (
                query_builder.match()
                .node()
                .related_to(label=config.object_type, ref_name=ref_name)
                .node()
            )

        query_builder = query_builder.where_literal(
            f"{ref_name}.last_ingested_at <= $earliest_allowed_time"
        ).return_literal(f"id({ref_name}) as id")

        return Query(str(query_builder), params)

    def generate_ttl_query_from_configuration(
        self, config: TimeToLiveConfiguration
    ) -> Query:
        ttl_match_query = self.generate_ttl_match_query(config)
        operation = (
            DELETE_NODE_QUERY
            if config.graph_object_type == GraphObjectType.NODE
            else DELETE_REL_QUERY
        )
        return ttl_match_query.feed_batched_query(operation)


import gi.repository
gi.require_version('Budgie', '1.0')
from gi.repository import Budgie, GObject, Gtk, GdkPixbuf, GLib, Gio, Gdk
import time
import os
import svgwrite
import datetime
from createsvg import createSVG
from math import sin, cos, pi

"""
    Analog Clock Applet for the Budgie Panel

    Copyright © 2020 Samuel Lane
    http://github.com/samlane-ma/

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <https://www.gnu.org/licenses/>.
"""

""" Some "constants":
    Clock is drawn at 200x200 px resolution, with a center at 100,100 and
    then scaled to fit panel, but these can be played around with to fine
    tune the clocks appearance.
"""
MAXIMUM_SIZE       = 200  # 200 default
MINIMUM_SIZE       =  22  #  22 default

IMAGE_SIZE         = 200  # 200 default
X_CENTER           = 100  #  50 default
Y_CENTER           = 100  #  50 default
CLOCK_RADIUS       =  92  #  46 default
HOUR_HAND_LENGTH   =  56  #  28 default
MINUTE_HAND_LENGTH =  76  #  38 default
FRAME_THICKNESS    =  10  #   5 default
HAND_THICKNESS     =  10  #   5 default
CENTER_DOT         =   8  #   3 default
MARKING_THICKNESS  =   8  #   4 default
MARKING_LENGTH     =  16  #   8 default

UPDATE_INTERVAL    =   5  #   5 default (in seconds)
FORCE_REDRAW       =  -1  # This is just to clarify why -1 is used in the code

app_settings = Gio.Settings.new("com.github.samlane-ma.budgie-analog-clock")
path = "com.solus-project.budgie-panel"

class BudgieAnalogClock(GObject.GObject, Budgie.Plugin):
    """ This is simply an entry point into your Budgie Applet implementation.
        Note you must always override Object, and implement Plugin.
    """
    # Good manners, make sure we have unique name in GObject type system
    __gtype_name__ = "BudgieAnalogClock"

    def __init__(self):
        """ Initialisation is important.
        """
        GObject.Object.__init__(self)

    def do_get_panel_widget(self, uuid):
        """ This is where the real fun happens. Return a new Budgie.Applet
            instance with the given UUID. The UUID is determined by the
            BudgiePanelManager, and is used for lifetime tracking.
        """
        return BudgieAnalogClockApplet(uuid)

class BudgieAnalogClockSettings(Gtk.Grid):

    def __init__(self, setting):
        super().__init__()

        self.label_text = ["", "Clock Size (px)","Frame Color","Hands Color",
                           "Face Color","Transparent face","","Show hour marks"]
        self.setting_name = ["clock-outline","clock-hands","clock-face"]

        for n in range(8):
            label = Gtk.Label()
            label.set_text(self.label_text[n])
            label.set_halign(Gtk.Align.START)
            label.set_valign(Gtk.Align.CENTER)
            self.attach(label, 0, n, 1, 1)

        spin_clock_size_adj = Gtk.Adjustment(value=app_settings.get_int("clock-size"),
                                  lower=MINIMUM_SIZE, upper=MAXIMUM_SIZE, step_incr=1)
        spin_clock_size = Gtk.SpinButton()
        spin_clock_size.set_adjustment(spin_clock_size_adj)
        spin_clock_size.set_digits(0)
        self.attach(spin_clock_size, 1, 1, 1, 1)

        self.colorbuttons = []
        for n in range(3):
            load_color = app_settings.get_string(self.setting_name[n])
            color = Gdk.RGBA()
            if load_color == "none":
                color.parse("rgba(0,0,0,0)")
            else:
                color.parse(load_color)
            button = Gtk.ColorButton.new_with_rgba(color)
            button.connect("color_set",self.on_color_changed,self.setting_name[n])
            self.colorbuttons.append(button)
            self.attach(self.colorbuttons[n], 1, n+2, 1, 1)

        button_set_transparent = Gtk.Button.new_with_label("Set")
        button_set_transparent.connect("clicked",self.on_set_transparent)
        self.attach(button_set_transparent, 1, 5, 1, 1)
        switch_markings = Gtk.Switch()
        switch_markings.set_halign(Gtk.Align.END)
        self.attach(switch_markings, 1, 7, 1, 1)

        app_settings.bind("clock-size",spin_clock_size,"value",Gio.SettingsBindFlags.DEFAULT)
        app_settings.bind("draw-marks",switch_markings,"active",Gio.SettingsBindFlags.DEFAULT)

        self.show_all()

    def on_set_transparent(self, button):
        self.colorbuttons[2].set_alpha(0)
        app_settings.set_string("clock-face","none")

    def on_color_changed(self, button, clock_part):
        color = button.get_color()
        hex_code = "#{:02x}{:02x}{:02x}".format(int(color.red/256),
                                                int(color.green/256),
                                                int(color.blue/256))
        app_settings.set_string(clock_part, hex_code)

class BudgieAnalogClockApplet(Budgie.Applet):
    """ Budgie.Applet is in fact a Gtk.Bin """
    manager = None

    def __init__(self, uuid):

        Budgie.Applet.__init__(self)

        self.uuid = uuid
        self.keep_running = True

        user = os.environ["USER"]
        self.tmp = os.path.join("/tmp/", user + "_panel_analog_clock.svg")
        self.max_size = MAXIMUM_SIZE

        self.load_settings()

        self.box = Gtk.Box()
        self.add(self.box)
        self.clock_image = Gtk.Image()
        self.box.add(self.clock_image)
        self.show_all()
        self.app_signal_id = app_settings.connect("changed",self.on_settings_change)
        GLib.timeout_add_seconds(UPDATE_INTERVAL, self.update_time)

        # This triggers the applet to shut down when removed from the panel
        # It needs a slight delay to work correctly
        GObject.idle_add(self.watch_applet, self.uuid)

    def find_applet (self, check_uuid, applets):
        for find_uuid in applets:
            if find_uuid == check_uuid:
                return True
        return False

    def watch_applet (self, check_uuid):
        applets = []
        panel_settings = Gio.Settings(path);
        allpanels_list = panel_settings.get_strv("panels")
        for p in allpanels_list:
            panelpath = "/com/solus-project/budgie-panel/panels/{" + p + "}/"
            self.currpanelsubject_settings = Gio.Settings.new_with_path(path + ".panel", panelpath)
            applets = self.currpanelsubject_settings.get_strv("applets")
            if self.find_applet(check_uuid, applets):
                # Need this signal id to disconnect it on quit
                self.panel_signal_id = self.currpanelsubject_settings.connect(
                              "changed::applets", self.is_applet_running, check_uuid)
        return False

    def is_applet_running (self, arg1, arg2, check_uuid):
        applets =self.currpanelsubject_settings.get_strv("applets")
        if not self.find_applet(check_uuid, applets):
            # Disconnect the signals
            self.currpanelsubject_settings.disconnect(self.panel_signal_id)
            app_settings.disconnect(self.app_signal_id)
            self.keep_running = False

    def do_panel_position_changed(self,position):
        if position == Budgie.PanelPosition.TOP or position == Budgie.PanelPosition.BOTTOM:
            self.box.set_orientation(Gtk.Orientation.HORIZONTAL)
        else:
            self.box.set_orientation(Gtk.Orientation.VERTICAL)
        self.update_clock()

    def do_panel_size_changed(self,panel_size,icon_size,small_icon_size):
        # Keeps the clock smaller than panel, but no smaller than MINIMUM_SIZE
        self.max_size = panel_size - 6
        if self.max_size < MINIMUM_SIZE:
            self.max_size = MINIMUM_SIZE
        current_size = app_settings.get_int("clock-size")
        # Don't let the clock get bigger than the Budgie Panel
        if current_size > self.max_size:
            self.clock_scale = self.max_size
        self.update_clock()

    def validate_settings(self):
        """ Reset invalid colors to defaults - "none" is a valid color name.
            Invalid settings (hopefully) should really only ever occur if
            the dconf values are manually set to something wrong
        """
        setting_name = ["clock-hands", "clock-outline", "clock-face"]
        default_color = ["#000000", "#000000", "#FFFFFF"]
        for n in range(3):
            testcolor = Gdk.RGBA()
            colorname = app_settings.get_string(setting_name[n])
            if (colorname != "none") and (not testcolor.parse(colorname)):
                app_settings.set_string(setting_name[n],default_color[n])

    def on_settings_change(self, *args):
        self.load_settings()
        self.update_clock()

    def load_settings(self):
        self.clock_scale = app_settings.get_int("clock-size")
        if self.clock_scale > self.max_size:
            self.clock_scale = self.max_size
        self.hands_color = app_settings.get_string("clock-hands")
        self.line_color = app_settings.get_string("clock-outline")
        self.fill_color = app_settings.get_string("clock-face")
        self.draw_hour_marks = app_settings.get_boolean("draw-marks")

    def update_clock(self):
        self.old_minute = FORCE_REDRAW
        self.validate_settings()
        self.update_time()

    def update_time(self):
        self.current_time = datetime.datetime.now()
        # Don't redraw unless time (minute) has changed
        if self.current_time.minute != self.old_minute:
            """ In the rare instance where the time zone is changed while applet
                is running, tzset will ensure the clock recognizes the change.
                Possibly may also be needed to recognize daylight savings change?
            """
            time.tzset()
            self.old_minute = self.current_time.minute
            self.create_clock_image(self.current_time.hour, self.current_time.minute)
            GObject.idle_add(self.load_new_image)
        return self.keep_running

    def load_new_image(self):
        self.clock_image.set_tooltip_text(self.current_time.strftime("%a %x"))
        self.pixbuf = GdkPixbuf.Pixbuf.new_from_file_at_scale(self.tmp, self.clock_scale,
                                                              self.clock_scale, True)
        self.clock_image.set_from_pixbuf(self.pixbuf)

    def create_clock_image (self, hours, mins):
        # If time is PM
        if hours > 12:
            hours -= 12
        # Treat hour hand like minute hand so it can be between hour markings
        hours = hours * 5 + (mins / 12)
        clock_svg = createSVG(self.tmp, IMAGE_SIZE, IMAGE_SIZE)
        # Draw an outside circle for the clock, and a small circle at the base of the hands
        clock_svg.add_circle(X_CENTER, Y_CENTER, CLOCK_RADIUS, fill=self.fill_color,
                            stroke=self.line_color, stroke_width=FRAME_THICKNESS)
        clock_svg.add_circle(X_CENTER, Y_CENTER, CENTER_DOT, stroke=self.hands_color,
                            fill=self.hands_color, stroke_width=1)

        # We are going to add hour markings around the outside edge of the clock
        if self.draw_hour_marks:
            for markings in range(12):
                if markings % 3 != 0:
                    mark_size = MARKING_THICKNESS * 2 / 3
                    mark_length = MARKING_LENGTH * 2 / 3
                else:
                    mark_size = MARKING_THICKNESS
                    mark_length = MARKING_LENGTH
                mark_x_start, mark_y_start = self.get_clock_hand_xy(markings * 5, CLOCK_RADIUS)
                mark_x_end, mark_y_end = self.get_clock_hand_xy(markings * 5, CLOCK_RADIUS - mark_length)
                clock_svg.add_line(mark_x_start, mark_y_start, mark_x_end, mark_y_end,
                                  stroke=self.line_color, stroke_width=mark_size)
        # Draw the minute and hour hands from the center to the calculated points
        hour_hand_x, hour_hand_y = self.get_clock_hand_xy (hours, HOUR_HAND_LENGTH)
        minute_hand_x, minute_hand_y = self.get_clock_hand_xy (mins, MINUTE_HAND_LENGTH)
        clock_svg.add_line(X_CENTER, Y_CENTER, hour_hand_x, hour_hand_y,
                          stroke=self.hands_color, stroke_width=HAND_THICKNESS)
        clock_svg.add_line(X_CENTER, Y_CENTER, minute_hand_x, minute_hand_y,
                          stroke=self.hands_color, stroke_width=HAND_THICKNESS)
        clock_svg.write_svg()

    def get_clock_hand_xy (self, hand_position,length):
        """ This fixes the issue that 0 degrees on a cirlce is actually 3:00
            on a clock, not 12:00 -essentially rotates the hands 90 degrees
        """
        hand_position -= 15
        if hand_position < 0:
            hand_position += 60
        # And here is how we determine the x and y coordinate to draw to
        radians = (hand_position * (pi * 2) / 60)
        x_position = round (X_CENTER + length * cos(radians))
        y_position = round (Y_CENTER + length * sin(radians))
        return x_position, y_position

    def do_supports_settings(self):
        """Return True if support setting through Budgie Setting,
        False otherwise.
        """
        return True

    def do_get_settings_ui(self):
        """Return the applet settings with given uuid"""
        return BudgieAnalogClockSettings(self.get_applet_settings(self.uuid))


import numpy as np
from Bio import SeqIO
from Bio.Align._aligners import PairwiseAligner
from Bio.Seq import Seq
from Levenshtein import distance
from six import unichr
import matplotlib.pyplot as plt
import numpy as np
#Sequences

filename_sequence_human = 'C:/Users/48512/Desktop/sequence.fasta'
filename_sequence_dog = 'C:/Users/48512/Desktop/sequence_dog.fasta'

sequence_human = SeqIO.read(filename_sequence_human, 'fasta')
sequence_dog = SeqIO.read(filename_sequence_dog, 'fasta')

seq_human = sequence_human.seq
seq_dog = sequence_dog.seq

seq1 = Seq('ACTTAG')
seq2 = Seq('ACTACT')

#print(distance(seq_human,seq_dog))

def delta(x,y):
    return 0 if x == y else 1
def M(seq1,seq2,i,j,k):
    return sum(delta(x,y) for x,y in zip(seq1[i:i+k],seq2[j:j+k]))

def makeMatrix(seq1,seq2,k):
    n = len(seq1)
    m = len(seq2)
    return [[M(seq1,seq2,i,j,k) for j in range(m-k+1)] for i in range(n-k+1)]

def plotMatrix(M,t, seq1, seq2, nonblank = unichr(0x25A0), blank = ' '):
    print(' |' + seq2)
    print('-'*(2 + len(seq2)))
    for label,row in zip(seq1,M):
        line = ''.join(nonblank if s < t else blank for s in row)
        print(label + '|' + line)
def dotplot(seq1,seq2,k = 1,t = 1):
    M = makeMatrix(seq1,seq2,k)
    plotMatrix(M, t, seq1,seq2)

plt.imshow(np.array(makeMatrix(seq_human,seq_dog,1)))
xt=plt.xticks(np.arange(len(list(seq_human))), list(seq_human))
yt = plt.yticks(np.arange(len(list(seq_human))), list(seq_human))
plt.show()


import os, sys
import random
import time
import ntpath
from collections import defaultdict
import ntpath
import pandas as pd  # Windows users: I copy-paste the pandas,dateutil and pyltz folders from anaconda 2!! into the site-packages folder of pymol(only for temporary use, other wise it gets confused with the paths of the packages)
import numpy as np
from pandas import Series
# Biopython
from Bio import SeqRecord, Alphabet, SeqIO
from Bio.Seq import Seq
import Bio.PDB as PDB
from Bio.Seq import MutableSeq
from Bio.PDB.Polypeptide import is_aa
from Bio.SVDSuperimposer import SVDSuperimposer
from pyro.infer.mcmc.api import MCMC
#Pymol
import pymol
from mpl_toolkits.mplot3d import Axes3D
import scipy.stats
# TORCH: "Tensors"
import torch
from torch.distributions import constraints, transform_to
# PYRO
import pyro
import pyro.distributions as dist
from pyro import poutine
from pyro.infer.autoguide import AutoDelta, AutoDiagonalNormal, AutoLowRankMultivariateNormal, init_to_median
from torch.optim import Adam, LBFGS
from pyro.infer import SVI, TraceEnum_ELBO, config_enumerate, Trace_ELBO, TraceGraph_ELBO,JitTrace_ELBO
import seaborn as sns
import matplotlib.pyplot as plt
import matplotlib as mpl
import tqdm
tqdm.monitor_interval = 0
from ignite.handlers import EarlyStopping
from ignite.engine import Engine, Events
from pyro.infer import SVI, EmpiricalMarginal
from pyro.optim import PyroOptim
from pyro.infer.abstract_infer import TracePredictive
from pyro.infer.mcmc import MCMC, NUTS
from pyro.infer.mcmc.util import initialize_model, predictive
import torch.multiprocessing as mp
#mp.set_start_method('spawn') ###For increasing number of chains
#Other details
use_cuda = True
tqdm.monitor_interval = 0
mpl.use('agg') #TkAgg
cuda = torch.device('cuda')
cuda0 = torch.device('cuda:0')
cuda2 = torch.device('cuda:1')
if use_cuda:
    torch.set_default_tensor_type('torch.cuda.FloatTensor') #PyTorch tensors have a dimension limit of 25 integers in GPU and 64 in CPU
    device = torch.device(cuda2)
    print("Using the GPU",torch.cuda.get_device_name(0))

else:
    print("Using CPU detected")
torch.multiprocessing.set_sharing_strategy('file_system')  # Needed to avoid runtime errors
class SVIEngine(Engine):
    def __init__(self, *args, step_args=None, **kwargs):
        self.svi = SVI(*args, **kwargs)
        self._step_args = step_args or {}
        super(SVIEngine, self).__init__(self._update)
    def _update(self, engine, batch):
        return -engine.svi.step(batch, **self._step_args)
PyroOptim.state_dict = lambda self: self.get_state()
def Max_variance(structure):
    '''Calculates the maximum distance to the origin of the structure, this value will define the variance of the prior's distribution'''
    centered = Center_torch(structure)
    mul = centered@torch.t(structure)
    max_var = torch.sqrt(torch.max(torch.diag(mul)))
    return max_var
def Extract_coordinates_from_PDB(PDB_file,type):
    ''' Returns both the alpha carbon coordinates contained in the PDB file and the residues coordinates for the desired chains'''
    from Bio.PDB.PDBParser import PDBParser
    from Bio.PDB import MMCIFParser
    Name = ntpath.basename(PDB_file).split('.')[0]

    try:
        parser = PDB.PDBParser()
        structure = parser.get_structure('%s' % (Name), PDB_file)
    except:
        parser = MMCIFParser()
        structure = parser.get_structure('%s' % (Name), PDB_file)

    ############## Iterating over residues to extract all of them even if there is more than 1 chain
    if type=='models':
        CoordinatesPerModel = []
        for model in structure:
            model_coord =[]
            for chain in model:
                for residue in chain:
                    if is_aa(residue.get_resname(), standard=True):
                            model_coord.append(residue['CA'].get_coord())
            CoordinatesPerModel.append(model_coord)

        return CoordinatesPerModel
    elif type=='chains':
        CoordinatesPerChain=[]
        for model in structure:
            for chain in model:
                chain_coord = []
                for residue in chain:
                    if is_aa(residue.get_resname(), standard=True):
                        chain_coord.append(residue['CA'].get_coord())
                CoordinatesPerChain.append(chain_coord)
        return CoordinatesPerChain

    elif type =='all':
        alpha_carbon_coordinates = []
        for chain in structure.get_chains():
            for residue in chain:
                if is_aa(residue.get_resname(), standard=True):
                    alpha_carbon_coordinates.append(residue['CA'].get_coord())
        return alpha_carbon_coordinates
def Average_Structure(tuple_struct):
    average = sum(list(tuple_struct)) / len(
        tuple_struct)  # sum element-wise the list of tensors containing the coordinates #tf.add_n
    return average
def Center_numpy(Array):
    '''Centering to the origin the data'''
    mean = np.mean(Array, axis=0)
    centered_array = Array - mean
    return centered_array
def Center_torch(Array):
    '''Centering to the origin the data'''
    mean = torch.mean(Array, dim=0)
    centered_array = Array - mean
    return centered_array
def sample_R(ri_vec):
    """Inputs a sample of unit quaternion and transforms it into a rotation matrix"""
    # argument i guarantees that the symbolic variable name will be identical everytime this method is called
    # repeating a symbolic variable name in a model will throw an error
    # the first argument states that i will be the name of the rotation made
    theta1 = 2 * np.pi * ri_vec[1]
    theta2 = 2 * np.pi * ri_vec[2]

    r1 = torch.sqrt(1 - ri_vec[0])
    r2 = torch.sqrt(ri_vec[0])

    qw = r2 * torch.cos(theta2)
    qx = r1 * torch.sin(theta1)
    qy = r1 * torch.cos(theta1)
    qz = r2 * torch.sin(theta2)

    R = torch.eye(3, 3) # device =cuda
    # filling the rotation matrix
    # Evangelos A. Coutsias, et al "Using quaternions to calculate RMSD" In: Journal of Computational Chemistry 25.15 (2004)

    # Row one
    R[0, 0] = qw ** 2 + qx ** 2 - qy ** 2 - qz ** 2
    R[0, 1] = 2 * (qx * qy - qw * qz)
    R[0, 2] = 2 * (qx * qz + qw * qy)

    # Row two
    R[1, 0] = 2 * (qx * qy + qw * qz)
    R[1, 1] = qw ** 2 - qx ** 2 + qy ** 2 - qz ** 2
    R[1, 2] = 2 * (qy * qz - qw * qx)

    # Row three
    R[2, 0] = 2 * (qx * qz - qw * qy)
    R[2, 1] = 2 * (qy * qz + qw * qx)
    R[2, 2] = qw ** 2 - qx ** 2 - qy ** 2 + qz ** 2
    return R
def RMSD_numpy(X1, X2):
    import torch.nn.functional as F
    return F.pairwise_distance(torch.from_numpy(X1), torch.from_numpy(X2))
def RMSD(X1, X2):
    import torch.nn.functional as F
    return F.pairwise_distance(X1, X2)
def RMSD_biopython(x, y):
    sup = SVDSuperimposer()
    sup.set(x, y)
    sup.run()
    rot, tran = sup.get_rotran()
    return rot
def Read_Data(prot1,prot2,type='models',models =(0,1),RMSD=True):
    '''Reads different types of proteins and extracts the alpha carbons from the models, chains or all . The model,
    chain or aminoacid range numbers are indicated by the tuple models'''

    if type == 'models':
        X1_coordinates = Extract_coordinates_from_PDB('{}'.format(prot1),type)[models[0]]
        X2_coordinates = Extract_coordinates_from_PDB('{}'.format(prot2),type)[models[1]]
    elif type == 'chains':
        X1_coordinates = Extract_coordinates_from_PDB('{}'.format(prot1),type)[models[0]][1:141]
        X2_coordinates = Extract_coordinates_from_PDB('{}'.format(prot2),type)[models[1]][0:140]

    elif type == 'all':
        X1_coordinates = Extract_coordinates_from_PDB('{}'.format(prot1),type)[models[0]:models[1]]
        X2_coordinates = Extract_coordinates_from_PDB('{}'.format(prot2),type)[models[0]:models[1]]

    #Apply RMSD to the protein that needs to be superimposed
    X1_Obs_Stacked = Center_numpy(np.vstack(X1_coordinates))
    X2_Obs_Stacked = Center_numpy(np.vstack(X2_coordinates))
    if RMSD:
        X2_Obs_Stacked = torch.from_numpy(np.dot(X2_Obs_Stacked,RMSD_biopython(X1_Obs_Stacked,X2_Obs_Stacked)))
        X1_Obs_Stacked = torch.from_numpy(X1_Obs_Stacked)
    else:
        X1_Obs_Stacked = torch.from_numpy(X1_Obs_Stacked)
        X2_Obs_Stacked = torch.from_numpy(X2_Obs_Stacked)

    data_obs = (X1_Obs_Stacked,X2_Obs_Stacked)

    # ###PLOT INPUT DATA################
    x = Center_numpy(np.vstack(X1_coordinates))[:, 0]
    y=Center_numpy(np.vstack(X1_coordinates))[:, 1]
    z=Center_numpy(np.vstack(X1_coordinates))[:, 2]
    fig = plt.figure(figsize=(18, 16), dpi=80)
    ax = fig.add_subplot(111, projection='3d')
    plt.plot(x, y, z)
    ax.plot(x, y,z ,c='b', label='data1',linewidth=3.0)
    #orange graph
    x2 = Center_numpy(np.vstack(X2_coordinates))[:, 0]
    y2=Center_numpy(np.vstack(X2_coordinates))[:, 1]
    z2=Center_numpy(np.vstack(X2_coordinates))[:, 2]
    ax.plot(x2, y2,z2, c='r', label='data2',linewidth=3.0)
    ax.legend()
    plt.savefig(r"Initial.png")
    plt.clf() #Clear the plot, otherwise it will give an error when plotting the loss
    plt.close()

    # rmsd = RMSD_numpy(Center_numpy(np.vstack(X1_coordinates)),Center_numpy(np.vstack(X2_coordinates)))
    # plt.plot(rmsd.numpy())
    # plt.show()
    return data_obs
def model(data):
    max_var,data1, data2 = data
    ### 1. prior over mean M
    #M = pyro.sample("M", dist.Normal(0, 3).expand_by([data1.size(0),data1.size(1)]).to_event(2))
    M = pyro.sample("M", dist.StudentT(1,0, 3).expand_by([data1.size(0),data1.size(1)]).to_event(2))
    ### 2. Prior over variances for the normal distribution
    U = pyro.sample("U", dist.HalfNormal(1).expand_by([data1.size(0)]).to_event(1))
    U =  U.reshape(data1.size(0),1).repeat(1,3).view(-1)  #Triplicate the rows for the subsequent mean calculation
    ## 3. prior over translations T_i: Sample translations for each of the x,y,z coordinates
    T2 = pyro.sample("T2", dist.Normal(0, 1).expand_by([3]).to_event(1))

    ## 4. prior over rotations R_i
    ri_vec = pyro.sample("ri_vec",dist.Uniform(0, 1).expand_by([3]).to_event(1))  # Uniform distribution
    #ri_vec = pyro.sample("ri_vec",dist.Normal(0,0.5).expand_by([3]).to_event(1))  # Weak normal? Apply torch.sigmoid after to ri_vecs
    R = sample_R(ri_vec)
    M_T1 = M
    M_R2_T2 = M @ R + T2

    # 5. Sampling from several Univariate Distributions (approximating the posterior distribution ): The observations are conditionally independant given the U, which is sampled outside the loop
    #UNIVARIATE NORMALS
    with pyro.plate("plate_univariate", data1.size(0)*data1.size(1),dim=-1):
        pyro.sample("X1", dist.StudentT(1,M_T1.view(-1), U),obs=data1.view(-1))
        pyro.sample("X2", dist.StudentT(1,M_R2_T2.view(-1), U), obs=data2.view(-1))
def _get_initial_trace(data_obs, average):
    '''Initialize MCMC and NUTS. Pyro 0.41'''
    if use_cuda:
        data_obs = [data.cuda() for data in data_obs]
        average = average.cuda()
    else:
        pass
    #INITIALIZE PRIOR:
    def init_prior(site):
        if site["name"] == "ri_vec":
            return torch.tensor([0.9, 0.1, 0.9])
        elif site["name"] == "M":
            return average
        else:
            return init_to_median(site)
    # GUIDE
    global_guide = AutoDelta(model,init_loc_fn=init_prior)
    # OPTIMIZER
    optim = pyro.optim.AdagradRMSProp(dict()) #https://github.com/pyro-ppl/pyro/blob/58277184310ef76a62420a38c300e84cd12b88ad/pyro/optim/adagrad_rmsprop.py
    elbo = JitTrace_ELBO()
    # STOCHASTIC VARIATIONAL INFERENCE
    svi_engine = SVIEngine(model, global_guide, optim, loss=elbo)
    pbar = tqdm.tqdm()
    loss_list = []
    # INITIALIZING PRIOR : Changing in the first iteration the value for the prior in order to constrain the prior over the rotations
    #pyro.param("auto_ri_vec", torch.Tensor([0.9, 0.1, 0.9]),constraint=constraints.unit_interval)  # constraint = constraints.simplex doesn't work exactly
    # Initialize the Mean Structure (each coordinate separately): NO CONSTRAINTS!!!
    #pyro.param("auto_M",average)
    @svi_engine.on(Events.EPOCH_COMPLETED)
    def update_progress(svi_engine):
        pbar.update(1)
        loss_list.append(-svi_engine.state.output)
        pbar.set_description(
            "[epoch {}] avg train loss: {}".format(svi_engine.state.epoch, svi_engine.state.output))
    # HANDLER
    handler = EarlyStopping(patience=25, score_function=lambda eng: eng.state.output, trainer=svi_engine)
    # SVI
    svi_engine.add_event_handler(Events.EPOCH_COMPLETED, handler)
    svi_engine.run([data_obs], max_epochs=15000)
    #return svi_engine.svi.exec_traces
    return global_guide.median()
def Run(data_obs, average,name1):
    if use_cuda:
        data_obs = [data.cuda() for data in data_obs]
        average = average.cuda()
    else:
        pass
    ###POSTERIOR PROBABILITY calculations: MCMC and NUTS
    # I had to fix a problem at /home/lys/anaconda3/lib/python3.5/site-packages/pyro/util.py by initializing the seed to rng_seed = random.randint(0,2**32-1)
    map_points = _get_initial_trace(data_obs,average)
    # MCMC initialization
    chains=1
    warmup= 500
    samples = 1000
    initialize = False
    if initialize:
        print("Initializing with MAP estimation")
        #Initialize NUTS's trace with the MAP estimate of the model
        init_params, potential_fn, transforms, _ = initialize_model(model,model_args=(data_obs,), num_chains=chains,jit_compile=True,skip_jit_warnings=True)
        map_points = _get_initial_trace(data_obs, average)
        init_params = {name: transforms[name](value).detach() for name, value in map_points.items()}
        # Choose NUTS kernel given the initialized potential function and a maximum tree depth limit.
        nuts_kernel = NUTS(potential_fn=potential_fn, max_tree_depth=8, target_accept_prob=0.8,jit_compile=True)
        # Prepare MCMC with NUTS kernel with the given arguments. Run over the observed data X1 and X2
        mcmc = MCMC(nuts_kernel, num_samples=samples, warmup_steps=warmup, num_chains=chains,initial_params=init_params, transforms=transforms)
        mcmc.run(data_obs)
    else:
        # Running MCMC using NUTS as selected kernel
        nuts_kernel = NUTS(model, jit_compile=True, ignore_jit_warnings=True, max_tree_depth=8)
        nuts_kernel.initial_trace = _get_initial_trace(data_obs, average)
        mcmc = MCMC(nuts_kernel, num_samples=samples, warmup_steps=warmup, num_chains=chains)
        mcmc.run(data_obs)
    max_var, data1, data2 = data_obs
    #######SAMPLES FROM THE POSTERIOR
    mcmc_samples = mcmc.get_samples()
    #######PARAMETERS:
    #ROTATION
    ri_vec_post_samples = mcmc_samples["ri_vec"]
    ri_vec_mean = ri_vec_post_samples.mean(dim=0)
    ri_vec_variance = ri_vec_post_samples.var(dim=0)
    R = sample_R(ri_vec_mean)
    # MEAN STRUCTURE
    M_post_samples = mcmc_samples["M"]
    M = M_post_samples.mean(dim=0)
    M_variance = M_post_samples.var(dim=0)
    #TRANSLATION
    T2_post_samples = mcmc_samples["T2"]
    T2_mean = T2_post_samples.mean(dim=0)
    T2_variance = T2_post_samples.var(dim=0)

    def Old_sampling_extraction():
        '''Pyro 0.41 mcmc sampling'''
        #Rotation matrix stats
        ri_vec_marginal_1 = mcmc.marginal(sites=["ri_vec"])
        ri_vec_marginal_1 = torch.cat(list(ri_vec_marginal_1.support(flatten=True).values()), dim=-1).cpu().numpy() #Where the samples are stored
        params = ['ri_vec[0]','ri_vec[1]','ri_vec[2]']
        df = pd.DataFrame(ri_vec_marginal_1, columns= params).transpose()
        df_summary = df.apply(pd.Series.describe, axis=1)[["mean", "std", "25%", "50%", "75%"]]
        df_summary.to_csv("ri_vec_stats_{}.txt".format(name1),sep='\t')
        # Rotation matrix output
        ri_vec_marginal = mcmc.marginal(["ri_vec"]).empirical["ri_vec"]
        ri_vec_mean = ri_vec_marginal.mean
        ri_vec_variance = ri_vec_marginal.variance
        R = sample_R(ri_vec_mean)

        # Mean structure stats
        M_marginal_1 = mcmc.marginal(sites=["M"])
        M_marginal_1 = torch.cat(list(M_marginal_1.support(flatten=True).values()), dim=-1).cpu().numpy()
        params = ['M[{}]'.format(i) for i in range(0,len(data1))]
        label_one = np.array(params)
        label_two = np.array(['x', 'y', 'z'])
        cols = pd.MultiIndex.from_product([label_one, label_two])
        df = pd.DataFrame(M_marginal_1.T.reshape(samples, -1), columns=cols).transpose()
        df_summary = df.apply(pd.Series.describe, axis=1)[["mean", "std", "25%", "50%", "75%"]]
        df_summary.to_csv("M_stats_{}.txt".format(name1),sep='\t')

        # Mean structure M output
        M_vec_marginal = mcmc.marginal(["M"]).empirical["M"]
        M_vec_mean = M_vec_marginal.mean
        M_vec_variance= M_vec_marginal.variance
        M = M_vec_mean
        #M = Center_torch(M_vec_mean.detach())

        # Translation stats
        T_marginal_1 = mcmc.marginal(sites=["T2"])
        T_marginal_1 = torch.cat(list(T_marginal_1.support(flatten=True).values()), dim=-1).cpu().numpy()
        params = ['T[0]','T[1]','T[2]']
        df = pd.DataFrame(T_marginal_1, columns= params).transpose()
        df_summary = df.apply(pd.Series.describe, axis=1)[["mean", "std", "25%", "50%", "75%"]]
        df_summary.to_csv("T_stats_{}.txt".format(name1),sep='\t')
        # Translation T output
        T2_vec_marginal = mcmc.marginal(["T2"]).empirical["T2"]
        T2_vec_mean = T2_vec_marginal.mean
        T2_vec_variance = T2_vec_marginal.variance

    #Observed
    X1 = data1.detach().cpu().numpy() #- T1_vec_mean.cpu().numpy()  # X1 -T1
    X2 = np.dot(data2.detach().cpu().numpy() - T2_mean.cpu().numpy(), np.transpose(R.cpu()))  # (X2-T2)R-1

    import matplotlib
    matplotlib.rcParams['legend.fontsize'] = 10

    #################PLOTS################################################
    fig = plt.figure(figsize=(18, 16), dpi=80)
    ax = fig.add_subplot(111, projection='3d')

    # blue graph
    x = X1[:, 0]
    y = X1[:, 1]
    z = X1[:, 2]

    ax.plot(x, y, z, c='b', label='X1', linewidth=3.0)

    # red graph
    x2 = X2[:, 0]
    y2 = X2[:, 1]
    z2 = X2[:, 2]

    ax.plot(x2, y2, z2, c='r', label='X2', linewidth=3.0)

    ###green graph
    x3 = M.cpu().numpy()[:, 0]
    y3 = M.cpu().numpy()[:, 1]
    z3 = M.cpu().numpy()[:, 2]
    #ax.plot(x3, y3, z3, c='g', label='M', linewidth=3.0)
    ax.legend()

    plt.title("Initialized MCMC and NUTS model")
    plt.savefig("{}_PLOTS_and_FILES_PYRO/Bayesian_Result_Samples_{}_{}_chains_{}".format(name1,name1,samples + warmup,chains))

    plt.clf()
    plt.plot(RMSD(data1.cpu(),data2.cpu()).numpy(), linewidth = 8.0)
    plt.plot(RMSD(torch.from_numpy(X1),torch.from_numpy(X2)).numpy(), linewidth=8.0)
    plt.ylabel('Pairwise distances',fontsize='46')
    plt.xlabel('Amino acid position',fontsize='46')
    plt.title('{}'.format(name1.upper()),fontsize ='46')
    plt.gca().legend(('RMSD', 'Theseus-PP'),fontsize='40')
    plt.savefig("{}_PLOTS_and_FILES_PYRO/Distance_Differences_Average_Bayesian_{}".format(name1,name1))
    plt.close()

    return T2_mean.cpu().numpy(), R.cpu(), M.cpu().numpy(), X1, X2, ri_vec_post_samples.cpu().numpy(),M_post_samples.cpu().numpy(),T2_post_samples.cpu().numpy() #Values for the mean structure
def Write_PDB(initialPDB, Rotation, Translation):
    ''' Transform the atom coordinates from the original PDB file and overwrite it '''
    from Bio.PDB.PDBParser import PDBParser
    from Bio.PDB import MMCIFParser, PDBIO
    Name = ntpath.basename(initialPDB).split('.')[0]

    try:
        parser = PDB.PDBParser()
        structure = parser.get_structure('%s' % (Name), initialPDB)
    except:
        parser = MMCIFParser()
        structure = parser.get_structure('%s' % (Name), initialPDB)
    for atom in structure.get_atoms():
        atom.transform(Rotation, -Translation)
    io = PDBIO()
    io.set_structure(structure)
    io.save("Transformed_{}".format(ntpath.basename(initialPDB)))
def write_ATOM_line(structure, file_name):
    import os
    """Transform coordinates to PDB file: Add intermediate coordinates to be able to visualize Mean structure in PyMOL"""
    expanded_structure = np.ones(shape=(2 * len(structure) - 1, 3))  # The expanded structure contains extra rows between the alpha carbons
    averagearray = np.zeros(shape=(len(structure) - 1, 3))  # should be of size len(structure) -1
    for index, row in enumerate(structure):
        if index != len(structure) and index != len(structure) - 1:
            averagearray[int(index)] = (structure[int(index)] + structure[int(index) + 1]) / 2
        else:
            pass
    # split the expanded structure in sets , where each set will be structure[0] + number*average
    # The even rows of the 'expanded structure' are simply the rows of the original structure
    expanded_structure[0::2] = structure
    expanded_structure[1::2] = averagearray
    structure = expanded_structure
    aa_name = "ALA"
    aa_type = "CA"
    if os.path.isfile(file_name):
        os.remove(file_name)
        for i in range(len(structure)):
            with open(file_name, 'a') as f:
                f.write(
                    "ATOM{:7d} {}   {} A{:4d}{:12.3f}{:8.3f}{:8.3f}  0.00  0.00    X    \n".format(i, aa_type, aa_name,
                                                                                i, structure[i, 0],
                                                                                structure[i, 1],
                                                                                structure[i, 2]))
    else:
        for i in range(len(structure)):
            with open(file_name, 'a') as f:
                f.write(
                    "ATOM{:7d} {}   {} A{:4d}{:12.3f}{:8.3f}{:8.3f}  0.00  0.00    X    \n".format(i, aa_type, aa_name,
                                                                                i, structure[i, 0],
                                                                                structure[i, 1],
                                                                                structure[i, 2]))
def Pymol(*args):
    '''Visualization program'''
    #LAUNCH PYMOL
    launch=False
    if launch:
        pymol.pymol_argv = ['pymol'] + sys.argv[1:]
        pymol.finish_launching(['pymol'])
    def Colour_Backbone(selection,color,color_digit):
        #pymol.cmd.select("alphas", "name ca") #apparently nothing is ca
        #pymol.cmd.select("sidechains", "! alphas") #select the opposite from ca, which should be the side chains, not working :(
        pymol.cmd.show("sticks", selection)
        pymol.cmd.set_color(color,color_digit)
        pymol.cmd.color(color,selection)

    # Load Structures and apply the function
    #colornames=['red','green','blue','orange','purple','yellow','black','aquamarine']
    #Palette of colours
    pal1 = sns.color_palette("OrRd",1)
    pal2 = sns.color_palette("PuBuGn_d",100) #RGB numbers for the palette colours
    colornames1 = ["red_{}".format(i) for i in range(0, len(pal1))]#So that X1 is red
    colornames2 = ["blue_{}".format(i) for i in range(0,len(pal2))]
    colornames = colornames1 + colornames2
    pal = pal1 + pal2

    snames=[]
    for file,color,color_digit in zip(args,colornames,pal):
        sname = ntpath.basename(file)
        snames.append(sname)
        pymol.cmd.load(file, sname) #discrete 1 will create different sets of atoms for each model
        pymol.cmd.bg_color("white")
        pymol.cmd.extend("Colour_Backbone", Colour_Backbone)
        Colour_Backbone(sname,color,color_digit)
    print(snames[0].split('_')[1])
    exit()
    pymol.cmd.png("{}_PLOTS_and_FILES_PYRO/Superposition_Bayesian_Pymol_{}".format(snames[0].split('_')[1],snames[0].split('_')[1]))
def Pymol_Samples(data1,data2,name1,R_samples,T_samples,samples):
    '''Create the PDB files to be sent to plot to PyMOL'''
    #Process the dataframes
    R_samples = torch.from_numpy(R_samples)
    indexes = random.sample(range(0, samples), samples) #not warm up samples
    X1 = data1.detach().cpu().numpy()
    plt.clf()
    plt.plot(RMSD(data1.cpu(), data2.cpu()).numpy(), linewidth=2.0)
    for i in range(0,samples):
        Rotation = sample_R(R_samples[i,:]) #torch
        Translation = T_samples[i,:] #numpy
        X2 = np.dot(data2.numpy() - Translation, np.transpose(Rotation.cpu().numpy()))
        write_ATOM_line(X2, os.path.join("{}_PLOTS_and_FILES_PYRO".format(name1),'Result_MCMC_{}_X2_{}.pdb'.format(name1,i)))
        plt.plot(RMSD(torch.from_numpy(X1), torch.from_numpy(X2)).numpy(), linewidth=0.5,color = plt.cm.autumn(i))
    plt.ylabel('Pairwise distances', fontsize='10')
    plt.xlabel('Amino acid position', fontsize='10')
    plt.title('{}'.format(name1.upper()), fontsize='10')
    plt.gca().legend(('RMSD', 'Theseus-PP'), fontsize='10')
    plt.savefig("{}_PLOTS_and_FILES_PYRO/Distance_Differences_Bayesian_{}.png".format(name1,name1),dpi=600)
    plt.close()
    names = [os.path.join("{}_PLOTS_and_FILES_PYRO".format(name1),'Result_MCMC_{}_X2_{}.pdb'.format(name1,i)) for i in indexes] #exchange indexes with range(0,samples)
    Pymol(*names)
def Folders(folder_name):
    """ Folder for all the generated images It will updated everytime!!! Save the previous folder before running again. Creates folder in current directory"""
    import os
    import shutil
    basepath = os.getcwd()
    if not basepath:
        newpath = folder_name
    else:
        newpath = basepath + "/%s" % folder_name

    if not os.path.exists(newpath):
        try:
            original_umask = os.umask(0)
            os.makedirs(newpath, 0o777)
        finally:
            os.umask(original_umask)
    else:
        shutil.rmtree(newpath)  # removes all the subdirectories!
        os.makedirs(newpath,0o777)

#if __name__ == "__main__":
name1 = '2lmp' #2nl7:139 #2do0=114
name2 ='2lmp'
models = (0,3)
samples =100
print(name1 + "\t" + str(models))
Folders("{}_PLOTS_and_FILES_PYRO".format(name1))
data_obs = Read_Data('../PDB_files/{}.pdb'.format(name1), '../PDB_files/{}.pdb'.format(name2),type='models',models =models,RMSD=True)
max_var = Max_variance(data_obs[0])
average = Average_Structure(data_obs)
data1, data2 = data_obs

write_ATOM_line(data1, os.path.join("{}_PLOTS_and_FILES_PYRO".format(name1),'RMSD_{}_data1.pdb'.format(name1)))
write_ATOM_line(data2, os.path.join("{}_PLOTS_and_FILES_PYRO".format(name1),'RMSD_{}_data2.pdb'.format(name1)))
#Pymol('{}_PDB_files/RMSD_{}_data1.pdb'.format(name1,name1), '{}_PDB_files/RMSD_{}_data2.pdb'.format(name2,name2))
data_obs = max_var, data1, data2
start  = time.time()
T2, R, M, X1, X2, ri_vec_samples,M_samples,T_samples = Run(data_obs, average,name1)
stop = time.time()

print("Time:", stop-start)
print('Memory Usage:')
print('Allocated:', round(torch.cuda.memory_allocated(0) / 1024 ** 3, 1), 'GB')
print('Cached:   ', round(torch.cuda.memory_cached(0) / 1024 ** 3, 1), 'GB')
write_ATOM_line(M, 'M.pdb')
write_ATOM_line(X1, os.path.join("{}_PLOTS_and_FILES_PYRO".format(name1),'Result_MCMC_{}_X1.pdb'.format(name1)))
write_ATOM_line(X2, os.path.join("{}_PLOTS_and_FILES_PYRO".format(name1),'Result_MCMC_{}_X2.pdb'.format(name2)))
#Write_PDB(r"../PDB_files/{}.pdb".format(name1), np.transpose(R), T1)
#Write_PDB(r"../PDB_files/{}.pdb".format(name2), np.transpose(R), T2)
#Pymol("Result_MCMC_{}_X1.pdb".format(name1), "Result_MCMC_{}_X2.pdb".format(name2))
Pymol_Samples(data1,data2,name1,ri_vec_samples,T_samples,samples)


import os
import collections

# scandir_out is the output of os.scandir
# scandir output gives a list of string with 'DirEntry <file_name>' this func removes Dir entry from the list
def format_scandir_output(scandir_out):
    dir_list = []
    file_list = []
    for entry in scandir_out :
        if entry.is_dir():
            dir_list.append(entry.name)
        if entry.is_file():
            file_list.append(entry.name)

    scandir_out.close()
    return dir_list, file_list

def get_latest_dir(scandir_out):
    dir_list = []
    file_list = []
    for entry in scandir_out :
        if entry.is_dir():
            dir_list.append(entry.name)
        if entry.is_file():
            file_list.append(entry.name)

    scandir_out.close()
    return dir_list, file_list

def sortandlist_os_scandir(out):
    dirs_dct = {}
    files_dct = {}
    for entry in out:
        if entry.is_dir():
            dirs_dct[os.path.getctime(entry)] = entry.name # sorting by key is easier (?) than sorting by values
            print (os.path.getctime(entry), entry.name)
        if entry.is_file():
            files_dct[os.path.getctime(entry)] = entry.name # sorting by key is easier (?) than sorting by values
            print (os.path.getctime(entry), entry.name)

    sorted_dirs_dct = collections.OrderedDict(sorted(dirs_dct.items())) # sort them using ordered dict
    sorted_dirs_list = list(sorted_dirs_dct.values())
    sorted_file_dct = collections.OrderedDict(sorted(files_dct.items())) # sort them using ordered dict
    sorted_file_list = list(sorted_file_dct.values())

    return sorted_dirs_list, sorted_file_list


with open("liczby.txt", "r") as file:
    lines = file.readlines()

numbers = []

for line in lines:
    numbers.append(line[:-1])

lista = ['1011', '100', '11000']

def zadanie2(numbers):
    print(numbers)
    for i in range(0, len(numbers)):
        if numbers[i][len(numbers)-1] ==0:
            print("tak")

print(zadanie1(numbers))
print(zadanie2(lista))


from datetime import *
from dateutil.rrule import *
from dateutil.parser import *
import time

def get_datetime():
    '''This function retrun current and
    change time for TDT and TOT tables'''

    current_time = time.gmtime()
    current_local_time = time.localtime()

    # italian rules
    current_offset_polarity = 0x0 # positive
    if current_local_time.tm_isdst == 1:
        start = 'DTSTART:%(year)04d%(month)02d%(day)02dT%(hour)02d%(minute)02d%(second)02d\n' % {
            "year": current_time[0],
            "month": current_time[1],
            "day": current_time[1],
            "hour" : 3,
            "minute": 0,
            "second": 0
        }
        start += 'RRULE:FREQ=YEARLY;BYDAY=-1SU;BYMONTH=10;COUNT=1' # last sunday of october
        current_offset = 0x02
        new_offset = 0x01
    else:
        start = 'DTSTART:%(year)04d%(month)02d%(day)02dT%(hour)02d%(minute)02d%(second)02d\n' % {
            "year": current_time[0],
            "month": current_time[1],
            "day": current_time[1],
            "hour" : 2,
            "minute": 0,
            "second": 0
        }
        start += 'RRULE:FREQ=YEARLY;BYDAY=-1SU;BYMONTH=3;COUNT=1' # last sunday of march
        current_offset = 0x01
        new_offset = 0x02

    change_time = list(rrulestr(start))[0]

    return (current_time, change_time, new_offset)


from  PyDSS.pyPlots.pyPlotAbstract import PlotAbstract
from bokeh.plotting import figure, curdoc
from bokeh.io import output_file
from bokeh.client import push_session
from bokeh.layouts import column
from bokeh.palettes import Set1
from ast import literal_eval as LE

class TimeSeries(PlotAbstract):
    def __init__(self,PlotProperties, dssBuses, dssObjects, dssCircuit, dssSolver):
        super(TimeSeries).__init__()
        self.__dssBuses = dssBuses
        self.__dssObjs = dssObjects
        self.__PlotProperties = PlotProperties
        self.__Properties =  LE(PlotProperties['Properties'])
        self.__index = LE(PlotProperties['index'])
        self.__multiplier = LE(PlotProperties['yScaler'])

        if self.__PlotProperties['yObjectType'] == 'Element':
            self.yObj = self.__dssObjs[self.__PlotProperties['yObjName']]
        elif self.__PlotProperties['yObjectType'] == 'Bus':
            self.yObj = self.__dssBuses[self.__PlotProperties['yObjName']]
        elif self.__PlotProperties['yObjectType'] == 'Circuit':
            self.yObj = dssCircuit

        self.Figures = []
        self.Y = [[] for x in self.__Properties]
        self.X = [[] for x in self.__Properties]
        self.Lines = [None for x in self.__Properties]
        output_file(PlotProperties['FileName'])
        for i, yProperties in enumerate(self.__Properties):
            Multiplier = self.__multiplier[i]
            pptyType ,Property = yProperties.split('.')
            if pptyType == 'p':
                newValue = self.yObj.GetParameter(Property)
            elif pptyType == 'v' and self.__PlotProperties['yObjectType'] != 'Circuit':
                newValue = self.yObj.GetVariable(Property)
            elif pptyType == 'v' and self.__PlotProperties['yObjectType'] == 'Circuit':
                newValue = getattr(self.yObj, Property)()

            if newValue is not None:
                if isinstance(newValue, list):
                    if self.__index[i] == 'SumEven':
                        self.Y[i].append(Multiplier * sum(newValue[::2]))
                    elif self.__index[i] == 'SumOdd':
                        self.Y[i].append(Multiplier * sum(newValue[1::2]))
                    elif self.__index[i] == 'Even' :
                        self.Y[i]= [[Multiplier * x] for x in newValue[::2]]
                    elif self.__index[i] == 'Odd':
                        self.Y[i] = [[Multiplier * x] for x in newValue[1::2]]
                    elif 'Index=' in self.__index[i]:
                        c = int(self.__index[i].replace('Index=', ''))
                        self.Y[i].append(Multiplier * newValue[c])
                else:
                    self.Y[i].append(Multiplier * newValue)
                self.X[i] = range(len(self.Y[i])+1)

            Figure = figure(plot_width=self.__PlotProperties['Width'],
                            plot_height=self.__PlotProperties['Height'],
                            title=self.__PlotProperties['yObjName'])
            if isinstance(self.Y[i][0], list):
                self.Lines[i] = []
                for j, Y in enumerate(self.Y[i]):
                    a = Figure.line(self.X[i], Y, line_width=3, line_alpha=1, color=Set1[9][j])
                    self.Lines[i].append(a)
            else:
                self.Lines[i] = Figure.line(self.X[i], self.Y[i], line_width=3, line_alpha=1)
            if self.__index[i] is not None:
                Figure.yaxis.axis_label = Property + '-' + self.__index[i]
            else:
                Figure.yaxis.axis_label = Property
            Figure.xaxis.axis_label = 'Timesteps'
            self.Figures.append(Figure)

        self.Layout = column(self.Figures)
        self.doc = curdoc()
        self.doc.add_root(self.Layout)
        self.doc.title = "PyDSS"
        self.session = push_session(self.doc)
        #self.session.show(self.Layout)  # open the document in a browser

        return

    def GetSessionID(self):
        return self.session.id

    def UpdatePlot(self):
        for i, yProperties in enumerate(self.__Properties):
            Multiplier = self.__multiplier[i]
            pptyType ,Property = yProperties.split('.')
            if pptyType == 'p':
                newValue = self.yObj.GetParameter(Property)
            elif pptyType == 'v' and self.__PlotProperties['yObjectType'] != 'Circuit':
                newValue = self.yObj.GetVariable(Property)
            elif pptyType == 'v' and self.__PlotProperties['yObjectType'] == 'Circuit':
                newValue = getattr(self.yObj, Property)()

            if newValue is not None:
                if isinstance(newValue, list):
                    if self.__index[i] == 'SumEven':
                        self.Y[i].append(Multiplier * sum(newValue[::2]))
                    elif self.__index[i] == 'SumOdd':
                        self.Y[i].append(Multiplier * sum(newValue[1::2]))
                    elif self.__index[i] == 'Even':
                        for j, Value in enumerate(newValue[::2]):
                            self.Y[i][j].append(Multiplier * Value)
                    elif self.__index[i] == 'Odd':
                        for j, Value in enumerate(newValue[1::2]):
                            self.Y[i][j].append(Multiplier * Value)
                    elif 'Index=' in self.__index[i]:
                        c = int(self.__index[i].replace('Index=',''))
                        self.Y[i].append(Multiplier * newValue[c])
                else:
                    self.Y[i].append(Multiplier * newValue)
                self.X[i] = range(len(self.Y[i])+1)

            if isinstance(self.Y[i][0], list):
                for j, Y in enumerate(self.Y[i]):
                    self.Lines[i][j].data_source.data['x'] = range(len(Y)+1)
                    self.Lines[i][j].data_source.data['y'] = Y
            else:
                if self.__PlotProperties['Dynamic'] and len(self.X[i]) > self.__PlotProperties['nSamples']:
                    self.X[i] = self.X[i][-self.__PlotProperties['nSamples']:]
                    self.Y[i] = self.Y[i][-self.__PlotProperties['nSamples']:]
                self.Lines[i].data_source.data['x'] = self.X[i]
                self.Lines[i].data_source.data['y'] = self.Y[i]
        return


import math
import datetime
import string

class Employee:
    def __init__(self, data):
        """

        Args:
            data layout: {'active':True,
                      'id': idNum, 'name': nameCapitalized, 'salonName':salon, 'pay':pay,
                      'fees':fees, 'rent':rent, 'printchecks': printchecks,
                      'paygrade':{'regType':regType, 'cashType': cashtype,
                                  'janitorType': janitortype, 'checkdealType': checkdealtype, 'owner': owner,
                                  'regular':{'commission': commission, 'check': check},
                                  'special':{'commissionspecial': comspec, 'checkdeal': checkdeal,
                                             'checkoriginal': checkoriginal, 'cashrate': cashrate}
                      }}}
        """
        self.id = data['id']
        self.name = data['name']
        self.salonName = data['salonName']
        self.basePay = data['pay']
        self.rent = data['rent']
        self.fees = data['fees']
        self.active = data['active']
        self.cashtype = data['paygrade']['cashType']
        self.janitortype = data['paygrade']['janitorType']
        self.checkdealtype = data['paygrade']['checkdealType']
        self.owner = data['paygrade']['owner']
        self.regtype = data['paygrade']['regType']
        self.commission = data['paygrade']['regular']['commission']
        self.check = data['paygrade']['regular']['check']
        self.commissionspecial = data['paygrade']['special']['commissionspecial']
        self.checkdeal = data['paygrade']['special']['checkdeal']
        self.checkoriginal = data['paygrade']['special']['checkoriginal']
        self.printchecks = data['printchecks']
        self.cashrate = float(data['paygrade']['special']['cashrate'])

        self.sales = {}
        self.sDate = ''         # m.d.y for saving text purpose
        self.payrollSummary = dict()
        self.xlreport= {'check': 0, 'checkdeal': 0, 'cash': 0}

        self.payrollPrint = ''

    def calculatePayroll(self, sales):
        """

        Args:
            sales: dictionary of daily sales, keys are datetime

        Returns:

        """
        if not self.janitortype:
            self.sales = sales.copy()
            self.genericCalculate()

    def genericCalculate(self):
        tips,commissionSales,totalSales,cleaningFees,daysWorked = [0 for i in range(1,6)]
        for days,amt in self.sales.items():
            tips += amt[2]
            commissionSales += amt[1]
            totalSales += amt[0]
            if amt[0] > 0:
                daysWorked += 1

        # self.commission is an integer 1-10 convert it to decimal percent
        check = (commissionSales * (self.check / 10))
        cash = commissionSales * ((10 - self.check) / 10)
        # check if made enough
        basePayPerDay = self.basePay / 6
        basePayPerRange = basePayPerDay * daysWorked
        basePayCheck = (basePayPerRange * (self.check / 10))
        basePayCash = basePayPerRange * ((10 - self.check) / 10)

        metgoal = False
        if daysWorked >= 6 and self.salonName != 'posh':
            if (commissionSales / daysWorked) >= basePayPerDay:
                metgoal = True
        else:
            if commissionSales > basePayPerRange:
                metgoal = True

        if self.rent >= 0:
            fees = (daysWorked * self.fees) + self.rent
        else:
            fees = daysWorked * self.fees

        self.payrollSummary = {
            'totalsale':totalSales,
            'commission':commissionSales,
            'check':check,
            'cash':cash,
            'basepaycheck':basePayCheck,
            'basepaycash':basePayCash,
            'tips':tips,
            'fees':fees,
            'daysworked': daysWorked,
            'metGoal': metgoal,
        }

        output = f'{"  "+string.capwords(self.salonName)+"  ":=^40}\n'
        output += f'{"Name":<10}{" ":10}{self.name:>20}\n'
        output += f'{" Summary ":-^40}\n'
        output += f'{"Total Sale":<10}{" ":20}{totalSales:>10.2f}\n'
        output += f'{"Commission":<10}{" ":20}{commissionSales:>10.2f}\n'
        output += f'{"Tips":<10}{" ":20}{tips:>10.2f}\n'
        output += f'{" ":30}{"-":->10}\n'
        output += f'{"Total Pay":<10}{" ":20}{commissionSales + tips:>10.2f}\n'
        output += f'{" Daily ":-^40}\n'
        output += f'{"Day":<10}{"Total":>10}{"Comm":>10}{"Tips":>10}\n'
        for day,amt in self.sales.items():
            d = datetime.datetime.strftime(day,'%m/%d:%a')
            output += f'{d:<10}{amt[0]:>10.2f}{amt[1]:>10.2f}{amt[2]:>10.2f}\n'
        output += f'{" ":12}{"-":->8}{" ":2}{"-":->8}{" ":2}{"-":->8}\n'
        output += f'{" ":10}{totalSales:>10.2f}{commissionSales:>10.2f}{tips:>10.2f}\n'
        output += f'{"=":=^40}\n'
        self.payrollPrint = output
        if self.regtype:
            self.regTypeSummaryAddOn()

    def regTypeSummaryAddOn(self):
        """
            This is the summary at the end of the payroll print out. When rent is positive, then that means
            owner is paying rent and will deduct from cash payroll. When rent is negative, then that means
            owner is helping employee pay rent and will be added to check payroll. Note in calculations
            if rent is negative, we subtract rent from check amount. Minus a negative is a positive
        Returns:
            None
        """
        check = self.payrollSummary['check'] if self.payrollSummary['metGoal'] else self.payrollSummary['basepaycheck']
        cash = self.payrollSummary['cash'] if self.payrollSummary['metGoal'] else self.payrollSummary['basepaycash']
        if self.rent < 0:
            self.xlreport['check'] = math.ceil(check + self.payrollSummary["tips"] - self.rent)
        else:
            self.xlreport['check'] = math.ceil(check + self.payrollSummary["tips"])
        self.xlreport['cash'] = math.ceil(cash - self.payrollSummary["fees"])
        output = ''
        if self.rent < 0:
            output = f'{"Check":<10} + {"Tip"} + {"Rent"}\n'
            output += f'{check:<10.2f} + {self.payrollSummary["tips"]:<8.2f} + {abs(self.rent):<5.0f} = ' \
                      f'${math.ceil(check + self.payrollSummary["tips"] - self.rent):<10}\n'
        else:
            output = f'{"Check":<10} + {"Tip"}\n'
            output += f'{check:<10.2f} + {self.payrollSummary["tips"]:<8.2f} = ' \
                      f'${math.ceil(check + self.payrollSummary["tips"]):<10}\n'
        output += f'{"Tien Mat":<10} - {"Le Phi":<8}\n'
        output += f'{cash:<10.2f} - {self.payrollSummary["fees"]:<8.2f} = ' \
                  f'${math.ceil(cash - self.payrollSummary["fees"]):<10}\n\n'
        self.payrollPrint += output

    def getPayrollSummary(self):
        return self.payrollSummary

    def getPrintOut(self):
        return self.payrollPrint

    def getStatus(self):
        if self.commission == 0 and self.regtype:
            return False
        try:
            result = {
                'total': math.ceil(self.payrollSummary['totalsale']),
                'comm': math.ceil(self.payrollSummary['commission']),
                'tips': math.ceil(self.payrollSummary['tips']),
                'daysWorked': math.ceil(self.payrollSummary['daysworked']),
                'metGoal': self.payrollSummary['metGoal']
            }
            return result
        except Exception:
            print('cannot get from emp {}'.format(self.name))

    def exportPayroll(self):
        sdate = ''
        counter = 0
        while counter < 1:
            for keys in self.sales:
                sdate = datetime.datetime.strftime(keys, '%m.%d.%Y')
            counter += 1

        fname = f'{self.salonName[0]}.{sdate}.{self.name}.txt'
        path = '../tmp/'
        with open(path + fname, 'w+') as file:
            file.writelines(self.payrollPrint)
        return path + fname

    def getId(self):
        return self.id

    def getName(self):
        return self.name

    def getInfo(self):
        return {'paygrade': {'regType':self.regtype, 'cashType': self.cashtype,
                            'janitorType': self.janitortype, 'checkdealType': self.checkdealtype, 'owner': self.owner,
                            'regular': {'commission': self.commission, 'check': self.check},
                            'special': {'commissionspecial': self.commissionspecial, 'checkdeal': self.checkdeal,
                                        'checkoriginal': self.checkoriginal, 'cashrate': self.cashrate}},
                'id': self.id, 'active': self.active,'salonName': self.salonName, 'name': self.name,
                'rent': self.rent, 'fees': self.fees, 'pay': self.basePay, 'printchecks': self.printchecks}

    def getXlReport(self):
        return self.xlreport

class EmployeeSpecial(Employee):
    """
        Nguoi nay can phai khai income thap cho nen ky check it ma khai so thiet
    """
    def __init__(self, data):
        super().__init__(data)
        self.payrollSummaryEtc = dict()

    def calculatePayroll(self,sales):

        self.sales = sales.copy()
        self.genericCalculate()

        '''
            so tien check ky ra va so tien check deal se co khac biet.
            check deal la so ky ra va check binh thuong la ho phai khai cuoi nam
        '''

        # self.commission is an integer 1-10 convert it to decimal percent
        check = self.payrollSummary['commission'] * (self.checkoriginal / 10)
        cash = self.payrollSummary['commission'] * ((10 - self.checkoriginal) / 10)

        checkdeal = self.payrollSummary['commission'] * (self.checkdeal / 10)
        cashdeal = self.payrollSummary['commission'] * ((10 - self.checkdeal) / 10)
        # check if made enough
        basePayPerDay = self.basePay / 6
        basePayPerRange = basePayPerDay * self.payrollSummary['daysworked']

        basePayCheck = basePayPerRange * (self.checkoriginal / 10)
        basePayCash = basePayPerRange * ((10 - self.checkoriginal) / 10)

        basePayCheckDeal = basePayPerRange * (self.checkdeal / 10)
        basePayCashDeal = basePayPerRange * ((10 - self.checkdeal) / 10)

        self.payrollSummaryEtc = {
            'check':check,
            'cash':cash,
            'basepaycheck': basePayCheck,
            'basepaycash': basePayCash,
            'metGoal':True if self.payrollSummary['commission'] > basePayPerRange else False,
            'checkdeal':checkdeal,
            'cashdeal':cashdeal,
            'basepaycheckdeal':basePayCashDeal,
            'basepaycashdeal':basePayCashDeal
        }
        check = self.payrollSummaryEtc['check'] if self.payrollSummaryEtc['metGoal'] else self.payrollSummaryEtc['basepaycheck']
        cash = self.payrollSummaryEtc['cash'] if self.payrollSummaryEtc['metGoal'] else self.payrollSummaryEtc['basepaycash']
        output = f'{"Check":<10} + {"Tip"}\n'
        output += f'{check:<10.2f} + {self.payrollSummary["tips"]:<8.2f} = ' \
                  f'${math.ceil(check + self.payrollSummary["tips"]):<10}\n'
        output += f'{"Tien Mat":<10} - {"Le Phi":<8}\n'
        output += f'{cash:<10.2f} - {self.payrollSummary["fees"]:<8.2f} = ' \
                  f'${math.ceil(cash - self.payrollSummary["fees"]):<10}\n\n'

        # add additional info to payroll for special case
        checkdeal = self.payrollSummaryEtc['checkdeal'] if self.payrollSummaryEtc['metGoal'] else self.payrollSummaryEtc['basepaycheckdeal']
        cashdeal = self.payrollSummaryEtc['cashdeal'] if self.payrollSummaryEtc['metGoal'] else self.payrollSummaryEtc['basepaycashdeal']
        self.xlreport['check'] = math.ceil(checkdeal + self.payrollSummary["tips"])
        self.xlreport['checkdeal'] = math.ceil(check + self.payrollSummary["tips"])
        self.xlreport['cash'] = math.ceil(cashdeal - self.payrollSummary["fees"])
        output += f'{"*":*^40}\n'
        output += f'{"Check":<10} + {"Tip"}\n'
        output += f'{checkdeal:<10.2f} + {self.payrollSummary["tips"]:<10.2f} = ' \
                             f'{math.ceil(checkdeal + self.payrollSummary["tips"]):<10}\n'
        output += f'{"Tien Mat":<10} - {"Le Phi":<8}\n'
        output += f'{cashdeal:<10.2f} - {self.payrollSummary["fees"]:<8} = ' \
                             f'{math.ceil(cashdeal - self.payrollSummary["fees"])}\n\n'
        self.payrollPrint += output

class EmployeeCash(Employee):
    def __init__(self, data):
        super().__init__(data)
        self.payrollSummaryEtc = {}

    def calculatePayroll(self,sales):
        self.sales = sales.copy()
        self.genericCalculate()

        # self.commission is an integer 1-10 convert it to decimal percent
        '''
            the next three lines are what makes the difference because check is
            deducted 17% and converted to cash
        '''
        check = self.payrollSummary['commission'] * (self.checkoriginal / 10)
        cash = self.payrollSummary['commission'] * ((10 - self.checkoriginal) / 10)
        # check if made enough
        basePayPerDay = self.basePay / 6
        basePayPerRange = basePayPerDay * self.payrollSummary['daysworked']
        basePayCheck = basePayPerRange * (self.checkoriginal / 10)
        basePayCash = basePayPerRange * ((10 - self.checkoriginal) / 10)
        self.payrollSummaryEtc = {
            'cash': cash,
            'check': check,
            'basepaycash': basePayCash,
            'basepaycheck': basePayCheck,
            'cashrate': self.cashrate,
            'metGoal': True if cash > basePayCash else False
        }

        # add additional info to payroll for special case
        cash = self.payrollSummaryEtc['cash'] if self.payrollSummaryEtc['metGoal'] else self.payrollSummaryEtc['basepaycash']
        check = self.payrollSummaryEtc['check'] if self.payrollSummaryEtc['metGoal'] else self.payrollSummaryEtc['basepaycheck']
        cashdeal = float(check + self.payrollSummary['tips']) *  self.payrollSummaryEtc['cashrate']
        self.xlreport['cash'] = math.ceil(cash + cashdeal - self.payrollSummary["fees"])
        output = f'{"Check":<10} + {"Tip"}\n'
        output += f'{check:<10.2f} + {self.payrollSummary["tips"]:<8.2f} = ' \
                  f'${math.ceil(check + self.payrollSummary["tips"]):<10}\n'
        output += f'{"Tien Mat":<10} - {"Le Phi":<8}\n'
        output += f'{cash:<10.2f} - {self.payrollSummary["fees"]:<8.2f} = ' \
                  f'${math.ceil(cash - self.payrollSummary["fees"]):<10}\n\n'

        output += f'{"*":*^40}\n'
        output += f'{"Check Qua Tien Mat":<30}:{cashdeal:>10.2f}\n'
        output += f'{"Tien Mat":<10} - {"Le Phi":<8}\n'
        output += f'{cash:<8.2f} - {self.payrollSummary["fees"]:<6} = ' \
                             f'{cash - self.payrollSummary["fees"]:>8.2f}\n'
        subTotal = cash + cashdeal - self.payrollSummary["fees"]
        output += f'{"Ca hai cong loi:":<30}{math.ceil(subTotal):>10}\n\n'
        self.payrollPrint += output

class EmployeeJanitor(Employee):
    def __init__(self, data):
        super().__init__(data)

    def calculatePayroll(self, sales=None):
        self.xlreport['check'] = self.basePay * (self.check * 0.1)
        self.xlreport['cash'] = self.basePay * ((10 - self.check) * 0.1)
        if self.xlreport['cash'] == 0:
            self.xlreport['check'] = self.xlreport['check'] - self.rent
        self.payrollPrint = f'Check: {self.xlreport["check"]}    Cash: {self.xlreport["cash"]}'


#!/usr/bin/python
#-*- coding: utf-8 -*-

'''

'''
__author__ = 'wpxiao'

from django.shortcuts import render,HttpResponse,redirect
from web import models
from django.urls import reverse
from utils import pagination_new
def index(request,*args,**kwargs):
    article_type_list = models.Article.type_choice
    if kwargs:
        article_type_id = int(kwargs["article_type_id"])
        base_url = reverse("index",kwargs=kwargs)
        print(base_url)
    else:
        article_type_id = None
        base_url = "/"
    # 获取文章分类的数量
    data_count = article_list = models.Article.objects.filter(**kwargs).count()

    page_obj = pagination_new.Pagination(current_page=request.GET.get("p"),data_count=data_count)
    article_list = models.Article.objects.filter(**kwargs).order_by('-nid')[page_obj.start:page_obj.end]
    page_str = page_obj.page_str(base_url)
    print(request.session["user_info"])
    return render(
        request,
        'index.html',
        {
            'article_list': article_list,
            'article_type_id': article_type_id,
            'article_type_list': article_type_list,
            'page_str': page_str,
        }
    )

from web.forms import article
from django.db import transaction
from utils.xss import XSSFilter
def article_add(request):
    if request.method == "GET":
        articleForm = article.ArticleForm(request)
        return render(request,"article_add.html",{"articleForm":articleForm})
    elif request.method == "POST":
        articleForm = article.ArticleForm(request=request,data=request.POST)
        if articleForm.is_valid():
            # 开启事务操作
            with transaction.atomic():
                tags = articleForm.cleaned_data.pop("tags")
                content = articleForm.cleaned_data.pop("content")
                content = XSSFilter().process(content)
                articleForm.cleaned_data["blog_id"] = request.session["user_info"]["blog__nid"]
                #创建文章
                article_obj = models.Article.objects.create(**articleForm.cleaned_data)
                # models.ArticleDetail.objects.create(content=content,article=article_obj)
                #功能同上，创建文章详细内容
                models.ArticleDetail.objects.create(content=content, article_id=article_obj.nid)
                tag_list = []
                for tag_id in tags:
                    tag_id = int(tag_id)
                    tag_list.append(models.Article2Tag(article_id = article_obj.nid,tag_id=tag_id))
                # 批量创建文章与标签关系，如果不调用批量创建接口的话，那就需要手动save
                models.Article2Tag.objects.bulk_create(tag_list)
            return redirect("/")
        else:
            # 表单验证错误的情况
            return render(request, "article_add.html", {"articleForm": articleForm})

def article_show(request,**kwargs):
    blog_site = kwargs.pop("blog_site")
    article_nid = int(kwargs.pop("article_nid"))
    article_obj = models.Article.objects.filter(nid=article_nid,blog__site=blog_site).first()
    # 一对一关系的反向查询
    article_content = article_obj.articledetail.content
    print(article_content)
    return render(request,"article_show.html",{"article_obj":article_obj,"article_content":article_content})

def upload_img(request):
    print("..........")
    result = {"error":0,"url":None}
    if request.method == "POST":
        dir = request.GET.get("dir",None)
        if dir == "image":
            img_obj = request.FILES.get("imgFile", None)
            print("img_obj:", img_obj.name)
            import os, json
            img_path = os.path.join("static/imgs/upload_img", img_obj.name)
            with open(img_path, mode="wb") as f:
                for i in img_obj.chunks():
                    f.write(i)
                result["url"] = "/static/imgs/upload_img/" + img_obj.name
            return HttpResponse(json.dumps(result))


from unittest import TestCase

from scrapy.downloadermiddlewares.useragent import UserAgentMiddleware
from scrapy.http import Request
from scrapy.spiders import Spider
from scrapy.utils.test import get_crawler

class UserAgentMiddlewareTest(TestCase):
    def get_spider_and_mw(self, default_useragent):
        crawler = get_crawler(Spider, {"USER_AGENT": default_useragent})
        spider = crawler._create_spider("foo")
        return spider, UserAgentMiddleware.from_crawler(crawler)

    def test_default_agent(self):
        spider, mw = self.get_spider_and_mw("default_useragent")
        req = Request("http://scrapytest.org/")
        assert mw.process_request(req, spider) is None
        self.assertEqual(req.headers["User-Agent"], b"default_useragent")

    def test_remove_agent(self):
        # settings USER_AGENT to None should remove the user agent
        spider, mw = self.get_spider_and_mw("default_useragent")
        spider.user_agent = None
        mw.spider_opened(spider)
        req = Request("http://scrapytest.org/")
        assert mw.process_request(req, spider) is None
        assert req.headers.get("User-Agent") is None

    def test_spider_agent(self):
        spider, mw = self.get_spider_and_mw("default_useragent")
        spider.user_agent = "spider_useragent"
        mw.spider_opened(spider)
        req = Request("http://scrapytest.org/")
        assert mw.process_request(req, spider) is None
        self.assertEqual(req.headers["User-Agent"], b"spider_useragent")

    def test_header_agent(self):
        spider, mw = self.get_spider_and_mw("default_useragent")
        spider.user_agent = "spider_useragent"
        mw.spider_opened(spider)
        req = Request(
            "http://scrapytest.org/", headers={"User-Agent": "header_useragent"}
        )
        assert mw.process_request(req, spider) is None
        self.assertEqual(req.headers["User-Agent"], b"header_useragent")

    def test_no_agent(self):
        spider, mw = self.get_spider_and_mw(None)
        spider.user_agent = None
        mw.spider_opened(spider)
        req = Request("http://scrapytest.org/")
        assert mw.process_request(req, spider) is None
        assert "User-Agent" not in req.headers


from numpy import *
from spivet import flolib, pivlib

import unittest, StringIO, sys

class test_floftle(unittest.TestCase):
    def setUp(self):
        self.eps = 1.e-6

        # Redirect stdout for chatty functions.
        self.sobfr = StringIO.StringIO()
        sys.stdout = self.sobfr

        self.velkey   = "VEL"
        self.tssdiv   = 12
        self.eisz     = 15
        self.ntrpc    = [1,4,4]
        self.ncalls   = 1
        self.irsbpath = "test-output"

        pdifname = "data/TGYRE.ex2"
        self.pd  = pivlib.loadpivdata(pdifname)

    def tearDown(self):
        self.sobfr.close()
        sys.stdout = sys.__stdout__

    def testForwardAdvection(self):
        irspath = "%s/%s" % (self.irsbpath,"FTLE-FWD")
        epslc   = slice(0,1,1)
        kpd     = pivlib.loadpivdata("data/TFTLE-FWD.ex2")

        ftledict = flolib.ftleinit(self.pd,
                                   epslc,
                                   self.eisz,
                                   self.ntrpc,
                                   self.ncalls,
                                   self.tssdiv)
        flolib.ftletrace(self.pd,self.velkey,irspath,ftledict)
        fpd = flolib.ftlecomp(irspath,ftledict)

        kvar = kpd[0]['FTLE']
        fvar = fpd[0]['FTLE']

        fdlta = abs(fvar -kvar)
        self.assertEqual( (fdlta < self.eps).all(), True )

    def testReverseAdvection(self):
        irspath = "%s/%s" % (self.irsbpath,"FTLE-BWD")
        epslc   = slice(15,14,-1)
        kpd     = pivlib.loadpivdata("data/TFTLE-BWD.ex2")

        ftledict = flolib.ftleinit(self.pd,
                                   epslc,
                                   self.eisz,
                                   self.ntrpc,
                                   self.ncalls,
                                   self.tssdiv)
        flolib.ftletrace(self.pd,self.velkey,irspath,ftledict)
        fpd = flolib.ftlecomp(irspath,ftledict)

        kvar = kpd[0]['FTLE']
        fvar = fpd[0]['FTLE']

        rdlta = abs(fvar -kvar)
        self.assertEqual( (rdlta < self.eps).all(), True)

def suite():
    suite = unittest.TestSuite()
    suite.addTest( unittest.makeSuite( test_floftle ) )

    return suite

if __name__ == '__main__':
    unittest.TextTestRunner(verbosity=2).run(suite())


s = str(input())
upper = 0
lower = 0

for x in s:
    if 65 <=ord(x)<=90:
        upper = upper+1
    elif 97 <= ord(x) <= 122:
        lower = lower+1
print( "UPPER: ", upper)
print("LOWER: ", lower)


# Databricks notebook source
# DBTITLE 1,Web Scraping from Hindenburg Research
from bs4 import BeautifulSoup
import requests

url = "https://hindenburgresearch.com/page/{}/"
page_num = 1
page = requests.get(url.format(page_num))

# COMMAND ----------

soup = BeautifulSoup(page.text, 'html')
print(soup.prettify())

# COMMAND ----------

headline = soup.find('h1').text
print(headline)
print(type(headline))

# COMMAND ----------

titles_arr = []

next_text = soup.find("a", class_="next page-numbers")

while next_text is not None:
    #titles = soup.find_all(['h2','span'], class_ = ['post-title','entry-title-primary'])
    titles = soup.find_all(['h2'], class_ = ['post-title'])
    for title in titles:
        print(title.text)
        titles_arr.append(title.text)
    page_num+=1
    page = requests.get(url.format(page_num))
    soup = BeautifulSoup(page.text, 'html')
    next_text = soup.find("a", class_="next page-numbers")

# COMMAND ----------

#Last page that does not have next button
if next_text is None:
    titles = soup.find_all(['h2'], class_ = ['post-title'])
    for title in titles:
        print(title.text)
        titles_arr.append(title.text)

# COMMAND ----------

print(page_num)
print(url.format(page_num))

# COMMAND ----------

print(len(titles_arr))

# COMMAND ----------

for title_text in titles_arr:
    print(title_text)


import Tkinter

class Meter(Tkinter.Frame):
    def __init__(self, master, width=300, height=20, bg='white', fillcolor='orchid1',\
                 value=0.0, text=None, font=None, textcolor='black', *args, **kw):
        Tkinter.Frame.__init__(self, master, bg=bg, width=width, height=height, *args, **kw)
        self._value = value

        self._canv = Tkinter.Canvas(self, bg=self['bg'], width=self['width'], height=self['height'],\
                                    highlightthickness=0, relief='flat', bd=0)
        self._canv.pack(fill='both', expand=1)
        self._rect = self._canv.create_rectangle(0, 0, 0, self._canv.winfo_reqheight(), fill=fillcolor,\
                                                 width=0)
        self._text = self._canv.create_text(self._canv.winfo_reqwidth()/2, self._canv.winfo_reqheight()/2,\
                                            text='', fill=textcolor)
        if font:
            self._canv.itemconfigure(self._text, font=font)

        self.set(value, text)
        self.bind('<Configure>', self._update_coords)

    def _update_coords(self, event):
        '''Updates the position of the text and rectangle inside the canvas when the size of
        the widget gets changed.'''
        # looks like we have to call update_idletasks() twice to make sure
        # to get the results we expect
        self._canv.update_idletasks()
        self._canv.coords(self._text, self._canv.winfo_width()/2, self._canv.winfo_height()/2)
        self._canv.coords(self._rect, 0, 0, self._canv.winfo_width()*self._value, self._canv.winfo_height())
        self._canv.update_idletasks()

    def get(self):
        return self._value, self._canv.itemcget(self._text, 'text')

    def set(self, value=0.0, text=None):
        #make the value failsafe:
        if value < 0.0:
            value = 0.0
        elif value > 1.0:
            value = 1.0
        self._value = value
        if text  is None:
            #if no text is specified use the default percentage string:
            text = str(int(round(100 * value))) + ' %'
        self._canv.coords(self._rect, 0, 0, self._canv.winfo_width()*value, self._canv.winfo_height())
        self._canv.itemconfigure(self._text, text=text)
        self._canv.update_idletasks()


#Day 1: Calorie Counting

f = open("day_1_input.txt")

contents = f.read().split("\n")

print(contents)

stacks = []
stack_list = []

for i in contents:
    if i == "":
        stack_list.append(stacks)
        stacks = []
    if i != "":
        stacks.append(int(i))

print(stack_list)

sum_list = []

for i in stack_list:
    sum_list.append(sum(i))

print(sum_list)

max_value = max(sum_list)
print(max_value)

first = 0
second = 0
third = 0
for i in sum_list:
    if i > first:
        third = second
        second = first
        first = i
    elif i > second:
        third = second
        second = i
    elif i > third:
        third = i

print(first + second + third)


from django.db import models
from phonenumber_field.modelfields import PhoneNumberField
from stdimage.models import StdImageField
from django.utils.text import slugify
from django.contrib.auth.models import User
from django.contrib.postgres.fields import ArrayField
from knox.models import AuthToken

STATES = ['AC', 'AL', 'AP', 'AM', 'BA', 'CE', 'DF', 'ES', 'GO', 'MA', 'MT', 'MS', 'MG', 'PA',
    'PB', 'PR', 'PE', 'PI', 'RJ', 'RN', 'RS', 'RO', 'RR', 'SC', 'SP', 'SE', 'TO']

STATES_CHOICES = [(state, state) for state in STATES]

class Base(models.Model):
    created_at = models.DateField("Created", auto_now_add=True)
    updated_at = models.DateField("Updated", auto_now=True)
    is_active = models.BooleanField("Active", default=True)

    class Meta:
        abstract = True

class Market(Base):
    name = models.CharField("Name", max_length=100)
    logo = StdImageField('Logo', upload_to='logos', variations={'thumb': (124, 124)})
    website = models.URLField("Website", max_length=200)
    ecommerce = models.URLField("E-commerce", max_length=200)
    phone = PhoneNumberField('Phone Number', unique=True)
    address = models.CharField("Address", max_length=300)
    s3_scrapy_path = models.CharField("Scrapy data path", max_length=128, null=True, blank=True)

    class Meta:
        verbose_name = 'Market'
        verbose_name_plural = 'Markets'

    def __str__(self):
        return self.name

class Category(models.Model):
    name = models.CharField("Name", max_length=128)
    icon = models.CharField("Icon name", null=True, max_length=32)

    def __str__(self):
        return f"{self.name}"

class Product(Base):
    name = models.CharField("Name", max_length=128)
    slug = models.SlugField("Slug", max_length=128, allow_unicode=True, null=True, blank=True)
    image_url = models.URLField("Image", blank=True, null=True)
    market = models.ForeignKey(Market, on_delete=models.PROTECT, related_name="products")
    category = models.ForeignKey(Category, on_delete=models.PROTECT, related_name="products", null=True)
    department = models.CharField("Department", max_length=128, null=True, blank=True)
    sub_department = models.CharField("Sub-department", max_length=128, null=True, blank=True)
    tags = ArrayField(models.CharField('Tag', max_length=128), null=True, blank=True)
    price = models.DecimalField("Price", decimal_places=2, max_digits=10)
    is_available = models.BooleanField("Is available?", default=True)
    link = models.URLField("Link")
    description = models.TextField("Description", null=True, blank=True)
    code = models.IntegerField("Market Code", null=True, blank=True)
    is_offer = models.BooleanField("Is it on offer?", default=False)
    similar = models.ManyToManyField("Product", null=True, blank=True, related_name="similar_products")
    comparison_name = models.CharField("Comparison Name", max_length=128, null=True, blank=True)
    max_price = models.DecimalField("Max Price", decimal_places=2, max_digits=10, null=True, blank=True)
    min_price = models.DecimalField("Min Price", decimal_places=2, max_digits=10, null=True, blank=True)
    similar_to = models.ManyToManyField("Product", null=True, blank=True, related_name="similar_products_set", through="Similar")

    def __str__(self):
        return f"{self.name} - {self.market.name}"

    def save(self, *args, **kwargs):
        self.slug = slugify(self.name)
        super(Product, self).save(*args, **kwargs)

class Similar(models.Model):
    product1 = models.ForeignKey("Product", on_delete=models.CASCADE, related_name="product1")
    product2 = models.ForeignKey("Product", on_delete=models.CASCADE, related_name="product2")
    similarity = models.DecimalField("Max Price", decimal_places=2, max_digits=5, null=True, blank=True)

class ProductHistory(Base):
    product = models.ForeignKey(Product, on_delete=models.PROTECT, related_name="history")
    price = models.DecimalField("Price", decimal_places=2, max_digits=10)

class PriceAlert(Base):
    user = models.ForeignKey(User, on_delete=models.PROTECT, related_name="alerts")
    product = models.ForeignKey(Product, on_delete=models.PROTECT, related_name="alerts")
    expected_price = models.DecimalField("Expected Price", decimal_places=2, max_digits=10)
    was_reached = models.BooleanField("Price reached?", default=False)
    reached_date = models.DateField("Reached date", blank=True, null=True)
    reached_price = models.DecimalField("Reached Price", decimal_places=2, max_digits=10, blank=True, null=True)

class Favorite(Base):
    user = models.ForeignKey(User, on_delete=models.PROTECT, related_name="favorites_products")
    product = models.ForeignKey(Product, on_delete=models.PROTECT, related_name="favorites")

class Notification(Base):
    title = models.CharField("Notification Title", max_length=128)
    body = models.TextField("Notification Body")
    payload = models.JSONField("Notification Payload", null=True)
    user = models.ForeignKey(User, on_delete=models.PROTECT, related_name="notifications")
    was_sent = models.BooleanField("Notification sent?", default=False)

class Account(Base):
    user = models.OneToOneField(User, related_name="account", on_delete=models.PROTECT)
    phone = PhoneNumberField('Phone Number', null=True, blank=True)
    street = models.CharField("Street", max_length=128, null=True, blank=True)
    district = models.CharField("District", max_length=128, null=True, blank=True)
    address_number = models.IntegerField("Number", null=True, blank=True)
    zip_code = models.CharField("Zip Code", max_length=8, blank=True)
    city = models.CharField("City", max_length=128, null=True, blank=True)
    state = models.CharField("State", max_length=2, null=True, choices=STATES_CHOICES, blank=True)
    fcm_token = models.CharField("Firebase Token", max_length=255, blank=True, null=True)
    google_id = models.CharField("Google Account Id", max_length=128, null=True, blank=True)


import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patheffects as path_effects

def plotConf(df_confusion,titulo,savefile):
  plt.ioff()
  fig = plt.figure()
  plt.clf()
  ax = fig.add_subplot(111)
  ax.set_aspect(1)
  res = ax.imshow(np.array(df_confusion), cmap=plt.cm.jet,
                  interpolation='nearest')

  width, height = np.array(df_confusion).shape

  for x in xrange(width):
      for y in xrange(height):
          ax.annotate(str(df_confusion[x][y]), xy=(y, x),
                      horizontalalignment='center',
                      verticalalignment='center', color='white',
                      fontweight='bold',
                      path_effects=[path_effects.withSimplePatchShadow(alpha=0.5)])

  plt.yticks( np.arange(height) )
  plt.xticks( np.arange(width) )
  plt.ylabel('Valor real')
  plt.xlabel('Valor predicho')
  cb = fig.colorbar(res)
  plt.title(titulo,fontsize=16)

  plt.savefig(savefile)
  return


# Given an array that contains at least three integers, we are asked to write a function that is going to find the three largest numbers in the input array and return them in an array sorted in ascending order. We are not allowed to sort the input array.

def findThreeLargestNumbers(array):
    top_three = [None] * 3
    for elem in array:
       if top_three[2] is None or elem > top_three[2]:
           shiftAndUpdate(top_three, elem, 2)
       elif top_three[1] is None or elem > top_three[1]:
           shiftAndUpdate(top_three, elem, 1)
       elif top_three[0] is None or elem > top_three[0]:
           shiftAndUpdate(top_three, elem, 0)
    return top_three

def shiftAndUpdate(array, num, idx):
    for i in range(idx + 1):
        if i == idx:
            array[i] = num # place highest value here
        else:
            array[i] = array[i + 1] # shift back 1


from django.shortcuts import redirect, render
from django.urls import reverse

from .models import Brand, Feature, Person, Rating, Project, TelegramUser
from .forms import RequestForm
from rest_framework.generics import CreateAPIView
from .serializers import TelegramBotSerializer
from telegram.ext import (
    Updater,
    CommandHandler,
    MessageHandler,
    Filters,
    ConversationHandler,
)

class TelegramUserCreateAPIView(CreateAPIView):
    serializer_class = TelegramBotSerializer
    queryset = TelegramUser.objects.all()

def home_view(request):
    brands = Brand.objects.all()
    features = Feature.objects.all()
    ratings = Rating.objects.all()
    projects = Project.objects.all()

    data = {
        'brands': brands,
        'features': features,
        'ratings': ratings,
        'projects': projects
    }

    return render(request, 'main/home.html', context=data)

def contact_us_view(request):
    return render(request, 'main/contact.html')

def success_view(request):
    return render(request, 'main/success.html')

def about_us_view(request):
    persons = Person.objects.all()
    data = {
        'persons': persons
    }
    return render(request, 'main/about-us.html', context=data)

def projects_view(request):
    projects = Project.objects.all()
    data = {
        'projects': projects
    }
    return render(request, 'main/projects.html', context=data)

def submit_request(request):
    updater = Updater("5478835430:AAEqT3Kw8QKDlTPRzIfSRVYAlBOW2uIxfPQ")
    telegram_users = TelegramUser.objects.all()
    if request.method == "POST" or None:
        form = RequestForm(request.POST or None)
        if form.is_valid():
            form.save(commit=False)
            obj = form.save()
            for user in telegram_users:
                updater.bot.send_message(chat_id=user.chat_id, text=f"""
                Ух ты у нас новая заявка № {obj.id} !!!  \n Номер телефона:  {obj.phone} \n Имя:  {obj.name}  \n Время создание: {obj.created_at} \n
                """, parse_mode="HTML")
            return redirect(reverse('success_view'))
        else:
            return render(request, 'main/contact.html', {'form': form})
    else:
        return render(request, 'main/404.html')