from io import StringIO
import os
import tempfile
from matplotlib import pyplot as plt
#os.environ["LANGCHAIN_HANDLER"] = "langchain"
import json
from typing import List, Dict, Text, Tuple
from haystack import BaseComponent
from haystack.schema import Document
from haystack.document_stores.base import BaseDocumentStore
import numpy as np
import joblib
import pandas as pd
from pymagnitude import Magnitude
import nltk
#nltk.download('punkt')
from nltk import word_tokenize
from transformers import AutoModelForSequenceClassification
from transformers import AutoTokenizer, AutoConfig
from scipy.special import softmax
from scipy import sparse
import re
import emoji
import string
from langchain.prompts import (
    ChatPromptTemplate,
    MessagesPlaceholder,
    SystemMessagePromptTemplate,
    HumanMessagePromptTemplate
)
from langchain import PromptTemplate, LLMChain
from langchain.chat_models import ChatOpenAI
#from langchain.callbacks import get_openai_callback
from langchain.callbacks.base import CallbackManager
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
import urllib.request
from haystack import Pipeline
from RabbitMQ import *
import mlflow.sklearn

# Set MLflow tracking server URI
mlflow.set_tracking_uri("http://mlflow.rasax.svc.cluster.local:8003")
mlflow.set_experiment("production_experiment")

MODEL_PATH = "/models/" # Mount path of PV

artifacts = [] # Untidy, but this is an afterthought.

class ConversationHistoryRetreiver(BaseComponent):
    outgoing_edges = 1

    def _parse_conversation_history(self, conversation_history, sender='user') -> [Text, Text]:
        # Concatenate all messages from the specified sender in the conversation history
        text = ' '.join([event.get('message', '') for event in conversation_history if event.get('event') == sender])

        # If the sender is 'cleo', we need to handle test results in the message
        if sender=='cleo':
            pattern = r"Danke! Hier ist deine Testauswertung: Neurotizismus: (\d{1,3}) % Extraversion: (\d{1,3}) % Offenheit für Erfahrung: (\d{1,3}) % Verträglichkeit: (\d{1,3}) % Gewissenhaftigkeit: (\d{1,3}) %"
            # Extract the matched percentages
            # Format the result string and remove the matched string from the original text
            match = re.search(pattern, text)
            if match:
                neuroticism = match.group(1)
                extraversion = match.group(2)
                openness = match.group(3)
                agreeableness = match.group(4)
                conscientiousness = match.group(5)

                result = f"Neuroticism,Extraversion,Openness,Agreeableness,Conscientiousness,{neuroticism},{extraversion},{openness},{agreeableness},{conscientiousness}"
                text = re.sub(pattern, "", text)
                return text, result
            else:
                print("Test pattern not found in the text.")

        return text, ""

    def run(self, query: List[Dict[Text, Text]]) -> Dict[Text, Text]:
        global artifacts
        artifacts = []
        print("--------- Query ---------")
        print(query)
        user_conversation_history, dummy = self._parse_conversation_history(query, sender='user')
        bot_conversation_history, test_result = self._parse_conversation_history(query, sender='cleo')
        sender_id = query[-1]["sender_id"] # Sender id is the last element
        run_id = query[-1]["run_id"]
        artifacts.append({
            'user_conversation': user_conversation_history,
            'bot_conversation' : bot_conversation_history,
        })
        artifacts.append({
            'actual_test_result': test_result
        })

        print(f"SENDER_ID: {sender_id}")
        print(f"RUN_ID: {run_id}")
        output={
            "conversation_history": user_conversation_history,
            "sender_id": sender_id,
            "run_id": run_id
        }

        # Send to dashboard
        send_to_rabbitmq(json.dumps(output), queue="text", sender_id=sender_id)

        return output, "output_1"

    def run_batch(self, queries: List[Text]) -> Dict[str, Document]:
        pass

class TfidfVectorizerNode(BaseComponent):
    outgoing_edges = 1

    def __init__(self, model_path = MODEL_PATH + "vectorizers/idf_vectorizer1.2.2.joblib"):
        self.model_path = model_path
        self.model_uri = "models:/big_five_tfidf_vectorizer/Production"
        try:
            self.vectorizer = mlflow.sklearn.load_model(model_uri=self.model_uri)
        except Exception as e:
            print(f"Error loading vectorizer: {e}.")
            print(f"Tried loading model from the following URI: {self.model_uri}")
            self.vectorizer = None

    def _tfidf_embeddings(self) -> Dict[Text, float]:
        try:
            # transform production data using the loaded vectorizer
            idf_dict = dict(zip(self.vectorizer.get_feature_names_out(), self.vectorizer.idf_))
            return idf_dict
        except Exception as e:
            print(f"Error loading vectorizer: {e}")
            return {}

    def run(self, conversation_history: Text, sender_id: str, run_id: str) -> Dict[Text, Document]:
        idf_embeddings = self._tfidf_embeddings()
        output={
            "conversation_history": conversation_history,
            "idf_embeddings": idf_embeddings,
            "sender_id": sender_id,
            "run_id": run_id
        }
        return output, "output_1"

    def run_batch(self, conversation_history: List[Text]) -> Dict[str, List]:
        conversations = {}
        for i, conversation in enumerate(conversation_history):
            conversation_id = f"conversation_{i}"
            embeddings = self._tfidf_embeddings(conversation)
            conversations[conversation_id] = {"conversation_history": conversation, "embeddings": embeddings}
        output = {"conversations": conversations}
        return output, "output_1"

class FasttextVectorizerNode(BaseComponent):
    outgoing_edges = 1

    def __init__(self, model_path: str = MODEL_PATH + "embeddings/wiki.de.vec.magnitude", embedding_dim: int = 300):
        self.embedding_dim = embedding_dim
        self.model_path = model_path
        # Load FastText German wiki embedding model
        self.model_url = "https://drive.google.com/uc?id=10ILYDkEFnlrExQwo7_iu2sL2le43Xlcp&export=download&confirm=t&uuid=92d36780-86fc-4fae-b03c-f05653d01849"
        try:
            # Check if the model file exists in the current directory
            print("Checking if Embeddings Model exists....")
            if not os.path.exists(self.model_path):
                print(f'{self.model_path} not found, downloading from {self.model_url}! Its 4 GB, so this may take a while. Maybe go and grab a coffee ;) ...')
                urllib.request.urlretrieve(self.model_url, self.model_path)
                print('Download complete!')
            else:
                print(f'{self.model_path} already exists')

            self.model = Magnitude(self.model_path)
        except Exception as e:
            print(f"Error loading vectorizer: {e}")
            self.model = None

    def _tfidf_w2v(self, text: Text, idf_dict: Dict[Text, Dict[Text, float]], sender_id: str) -> np.array(List[List[float]]):
        vectors = []
        words = word_tokenize(text) # Split text into words
        w2v_vectors = self.model.query(words) # Get FastText vectors
        weights = [idf_dict.get(word, 1) for word in words] # # Use pretrained tf-idf model to idf-weight the embedding features
        vectors.append(np.average(w2v_vectors, axis = 0, weights = weights))

        embeddings = {"words": words, "vectors": w2v_vectors.tolist()}
        send_to_rabbitmq(json.dumps(embeddings), queue="embeddings", sender_id=sender_id)

        return np.array(vectors)

    def run(self, conversation_history: Text, idf_embeddings: Dict[Text, Dict[str, float]], sender_id: str, run_id: str) -> Dict[str, np.ndarray]:
        vectors = self._tfidf_w2v(conversation_history, idf_embeddings, sender_id)
        output={
            "vectors": vectors,
            "sender_id": sender_id,
            "run_id": run_id,
        }
        return output, "output_1"

    def run_batch(self, conversation_history: Text, idf_embeddings: Dict[str, Dict[str, float]]) -> Dict[str, np.ndarray]:
        pass

class NormalizerNode(BaseComponent):
    outgoing_edges = 1

    def __init__(self, model_path: str = MODEL_PATH + "normalizers/embedding_normalizer1.2.2.joblib", input="embeddings"):
        self.model_path = model_path
        self.input = input
        self.model_uri = ""
        try:
            if input == "embeddings":
                self.model_uri = "models:/big_five_embedding_normalizer/Production"
            elif input == "features":
                self.model_uri = "models:/big_five_feature_normalizer/Production"

            self.normalizer = mlflow.sklearn.load_model(model_uri=self.model_uri)
        except Exception as e:
            print(f"Error loading vectorizer: {e}")
            print(f"Tried loading model from the following URI: {self.model_uri}")
            self.normalizer = None

    def _normalize(self, vectors: np.array(float)):
        try:
            # transform production data using the loaded normalizer
            normalized_vectors = self.normalizer.transform(vectors)
            return normalized_vectors
        except Exception as e:
            print(f"Error loading vectorizer: {e}")
            return {}

    def run(self, vectors: np.array(float), sender_id: str, run_id: str) -> Dict[str, np.ndarray]:
        normalized_vectors = self._normalize(vectors)
        output={
            f"n{self.input}": normalized_vectors,
            "sender_id": sender_id,
            "run_id": run_id
        }
        #print(output)
        return output, "output_1"

    def run_batch(self, vectors: np.array(float)) -> Dict[str, np.ndarray]:
        pass

class BigFiveFeaturizer(BaseComponent):
    outgoing_edges = 1

    def __init__(self, model_name_or_path = f"cardiffnlp/twitter-xlm-roberta-base-sentiment"):
        # Load XLM-RoBERTA sentiment model from HuggingFace
        self.model_name_or_path = model_name_or_path
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name_or_path)
        self.config = AutoConfig.from_pretrained(self.model_name_or_path)
        self.model = AutoModelForSequenceClassification.from_pretrained(self.model_name_or_path)
        self.model.save_pretrained(MODEL_PATH + self.model_name_or_path)
        self.tokenizer.save_pretrained(MODEL_PATH + self.model_name_or_path)

    def _preprocess(self, text):
        new_text = []
        for t in text.split(" "):
            t = '@user' if t.startswith('@') and len(t) > 1 else t
            t = 'http' if t.startswith('http') else t
            new_text.append(t)
        return " ".join(new_text)

    def _sentiment_analysis(self, text: Text, sender_id: str):
        text = self._preprocess(text)
        encoded_input = self.tokenizer(text, return_tensors='pt')
        output = self.model(**encoded_input)
        scores = output[0][0].detach().numpy()
        scores = softmax(scores) # Format: [neg, neu, pos]

        # Send to dashboard
        sentiment_scores = {"sentiment": scores.tolist()}
        send_to_rabbitmq(json.dumps(sentiment_scores), queue="sentiment", sender_id=sender_id)

        return np.array([scores])

    # Count unicode emojis with the emoji library
    def _count_emojis(self, s):
        cnt = 0
        for word in word_tokenize(s):
            if emoji.is_emoji(word):
                cnt += 1
        return cnt

    # Count emojis and emoticons with regex phrases
    def _emoji_count(self, text):
        emoticons_re = [
            '(\:\w+\:|\<[\/\\]?3|[\(\)\\\D|\*\$][\-\^]?[\:\;\=]|[\:\;\=B8][\-\^]?[3DOPp\@\$\*\\\)\(\/\|])(?=\s|[\!\.\?]|$)']
        is_emote = []

        no_of_phrases = 0
        for re_patten in emoticons_re:
            no_of_phrases += len(re.findall(re_patten, text))

        no_of_phrases += self._count_emojis(text)

        is_emote.append(no_of_phrases)
        return np.array(is_emote).reshape(-1, 1)

    def _count_punctuations(self, text):
        puncts = []
        punctuations = set(string.punctuation)
        count = lambda l1,l2: sum([1 for x in l1 if x in l2])
        puncts.append(count(text,punctuations))

        return np.array(puncts).reshape(-1,1)

    def _num_dots(self, text):
        num_dots = []
        num_dots.append(text.count('.'))

        return np.array(num_dots).reshape(-1,1)

    def _text_features(self, text):
        longest_word_length = []
        mean_word_length = []
        length_in_chars = []

        length_in_chars.append(len(text))
        longest_word_length.append(len(max(text.split(), key=len)))
        mean_word_length.append(np.mean([len(word) for word in text.split()]))

        longest_word_length = np.array(longest_word_length).reshape(-1, 1)
        mean_word_length = np.array(mean_word_length).reshape(-1, 1)
        length_in_chars = np.array(length_in_chars).reshape(-1, 1)

        return np.concatenate([longest_word_length, mean_word_length, length_in_chars], axis=1)

    def _featurize(self, text: Text, sender_id: str) -> np.hstack:
        emoji_re = self._emoji_count(text)
        num_dots = self._num_dots(text)
        num_punctuations = self._count_punctuations(text)
        sentiment = self._sentiment_analysis(text, sender_id)
        text_features = self._text_features(text)

        features = np.hstack((
            emoji_re,
            num_dots,
            num_punctuations,
            sentiment,
            text_features))
        return features

    def run(self, conversation_history: Text, sender_id: str, run_id: str) -> np.hstack:
        vectors= self._featurize(conversation_history, sender_id)
        output={
            "vectors": vectors,
            "sender_id": sender_id,
            "run_id": run_id
        }
        return output, "output_1"

    def run_batch(self, conversation_history: Text) -> Dict[str, Document]:
        pass

class ConcatenationNode(BaseComponent):
    outgoing_edges = 1

    def _concatenate_features(self, features, embeddings):
        features = sparse.csr_matrix(features)
        embeddings = sparse.csr_matrix(embeddings)

        combined_features = sparse.hstack((
            features,
            embeddings
        ))

        # Its a bit untidy to just hardcore the feature names like that
        # In the future, we should either parse them through from
        # the featurizer, or bake them into the vector dict
        feature_names = ['train_emoji_re',
                         'num_dots',
                         'longest_word_length',
                         'mean_word_length',
                         'length_in_chars',
                         'sentiment_neg',
                         'sentiment_neu',
                         'sentiment_pos',
                         'num_punctuations'
                         ]

        # Also hardcoding our 300 D vector space. In the future we should read that from the data itself.
        feature_names = feature_names + ['fasttext_' + str(col) for col in range(300)]

        return combined_features, feature_names

    def run(self, inputs: List[dict]) -> Dict[str, np.ndarray]:
        sender_id = inputs[0]["sender_id"]
        run_id = inputs[0]["run_id"]
        concatenated_features, feature_names = self._concatenate_features(inputs[0]["nfeatures"], inputs[1]["nembeddings"])
        output={
            "concatenated_features": concatenated_features,
            "feature_names": feature_names,
            "sender_id": sender_id,
            "run_id": run_id
        }
        return output, "output_1"

    def run_batch(self, embedding_vectors: np.array(float), feature_vectors: np.array(float))-> Dict[str, np.ndarray]:
        pass

class BigFiveFeatureSelectionNode(BaseComponent):
    outgoing_edges = 1

    def __init__(self, model_paths: Dict[str, str]):
        self.model_paths = model_paths
        self.selectors = {}

        for dimension, model_path in self.model_paths.items():
            try:
                self.selectors[dimension] = mlflow.sklearn.load_model(model_uri=f"models:/{dimension}_selector/Production")
            except Exception as e:
                print(f"Error loading selector: {e}")
                self.selectors = None

    def _select_features(self, features: np.array(float), feature_names: List[str], sender_id: str) -> np.ndarray:
        selected_features = {}
        feature_json = {"features": {}}
        for dimension, selector in self.selectors.items():
            try:
                selected_features[dimension] = selector.transform(features.tocsr())

                # Build json for dashboard dashboard
                feature_json['features'][dimension] = {
                    "num_of_features": len(selector.k_feature_idx_),
                    "feature_names": np.array(feature_names)[list(selector.k_feature_idx_)].tolist()
                }
            except Exception as e:
                print(f"Error loading selector: {e}")
        print(feature_json)
        # Send to dashboard
        send_to_rabbitmq(json.dumps(feature_json), queue="features", sender_id=sender_id)

        return selected_features

    def run(self, concatenated_features: np.array(float), feature_names: List[str], sender_id: str) -> Dict[str, np.ndarray]:
        selected_features = self._select_features(concatenated_features, feature_names, sender_id)
        output={"selected_features": selected_features}
        return output, "output_1"

    def run_batch(self, vectors: np.array(float)) -> Dict[str, np.ndarray]:
        pass

class BigFiveClassifierNode(BaseComponent):
    outgoing_edges = 1

    def __init__(self, model_paths: Dict[str, str], thresholds: Dict[str, float]):
        self.model_paths = model_paths
        self.thresholds = thresholds
        self.models = {}

        for dimension, model_path in self.model_paths.items():
            try:
                # Load our pretrained classification models from the model registry
                self.models[dimension] = mlflow.sklearn.load_model(model_uri=f"models:/{dimension}_classifier/Production")
            except Exception as e:
                print(f"Error loading model: {e}")
                self.models = None

    def _predict(self, selected_features: Dict[str, np.ndarray], sender_id: str, run_id: str):
        predicted_classes = {}
        predicted_proba = {}

        # Get threshold config from dashboard
        body = receive_rabbitmq(queue="thresholds", sender_id=sender_id)
        if body:
            self.thresholds = json.loads(body)
            send_to_rabbitmq(json.dumps(self.thresholds), queue="actual-thresholds", sender_id=sender_id)

        for dimension, features in selected_features.items():
            try:
                # Predict the Big Five Classification!
                predicted_proba[dimension] = self.models[dimension].predict_proba(features)[:, 1]
                predicted_classes[dimension] = np.where(predicted_proba[dimension] >= self.thresholds[dimension], 1, 0)

            except Exception as e:
                print(f"Error predicting with model: {e}")

        predictions = {
            "classes": predicted_classes,
            "probabilities": predicted_proba
        }

        predictions_json = predictions.copy()
        # Convert NumPy arrays to Python lists. Required for json convertion
        for key, value in predictions_json.items():
            if isinstance(value, dict):
                for sub_key, sub_value in value.items():
                    value[sub_key] = sub_value.tolist()
            else:
                predictions_json[key] = value.tolist()
        # Send to dashboard
        send_to_rabbitmq(json.dumps(predictions_json), queue="classification", sender_id=sender_id)

        return predictions

    def run(self, selected_features: Dict[str, np.ndarray], sender_id: str, run_id: str) -> Dict[str, Dict[str, np.ndarray]]:
        predictions = self._predict(selected_features, sender_id=sender_id, run_id=run_id)
        global artifacts
        artifacts.append({
            "pipeline_predictions": json.dumps(predictions)
        })
        print("----------- Predictions --------------")
        print(predictions)
        output={
            "predictions": predictions,
            "sender_id": sender_id,
            "run_id": run_id
        }
        return output, "output_1"

    def run_batch(self, selected_features: np.array(float)) -> Dict[str, np.ndarray]:
        pass

class BigFiveClassificationEvaluator(BaseComponent):
    outgoing_edges = 1

    def __init__(self):
        self.llm = ChatOpenAI(temperature=0)
        # Alternative Prompt that also worked well.
        # self.template = """You are a powerful AI model for text classification and with comprehensive knowledge of the Big Five personality model.
        # This knowledge is backed up by scientific research and literature that describes the five traits of the Big Five model and their impact on peoples
        # individual writing style through linguistic cues.
        # You will receive a text written by a human in German language (German Human Text)  that was received from a conversation by that human with a chatbot. You will do a Big Five classification of that text on all five Traits, that will look like that:
        # Neuroticism: CF[cf_class], Proba[cf_proba]
        # Extraversion: CF[cf_class], Proba[cf_proba]
        # Openness for experience: CF[cf_class], Proba[cf_proba]
        # Agreeableness: CF[cf_class], Proba[cf_proba]
        # Conscientiousness: CF[cf_class], Proba[cf_proba]

        # - cf_class is be the classification result on each individual trait and will be either 1 (if classification is positive) or 0 (if classification is negative).
        # You will do this classification based on your knowledge about the Big Five traits and their linguistic cues. However, in general a classification is considered 1 (positive) if this trait seems to be pronounced above average in that person.
        # - cf_proba is the probability of that classification result on each individual trait and will be a floating point number between 0 and 1.

        # German Human Text:
        # {text}
        # Do the classification:"""
        self.template = """Please do a best effort binary classification for the five big five traits neuroticism, extraversion, openness for experience, agreeableness and conscientiousness for the following German text.
        Please classify with a binary 1 if you think the peculiarity of the dimension is above average in that person or 0 if its below average. Try to also include a score from 1-100 how pronounced you think the dimension is in this person.
        Also add your reasoning behind your analysis in 1 or 2 short sentences.

        German Human Text:
        {text}
        Do the classification:"""
        self.prompt = PromptTemplate(input_variables=["text"], template=self.template)

    def _predict(self, text: Text) -> Text:
        llm_chain = LLMChain(prompt=self.prompt, llm=self.llm)
        cf_result = llm_chain.run(text=text)
        return cf_result

    def run(self, conversation_history: Text, sender_id: str, run_id: str) -> Text:
        predictions = self._predict(conversation_history)

        global artifacts
        artifacts.append({
            "evaluation_classification_result": predictions
        })

        print("----------- Predictions --------------")
        print(predictions)
        output={
            "eval_predictions": predictions,
            "sender_id": sender_id,
            "run_id": run_id
        }
        send_to_rabbitmq(json.dumps(output['eval_predictions']), queue="eval_classification", sender_id=sender_id)
        return output, "output_1"

    def run_batch(self, conversation_history: Text, sender_id: str) -> Text:
        pass

class BigFiveResponseGenerator(BaseComponent):
    outgoing_edges = 1

    def __init__(self):
        self.template="""Assistant is called 'Cleo'. Cleo is designed to be an agent for psychological counseling and for engaging and nice conversations on all sorts of topics with a user.
As a language model, Cleo is able to generate human-like text based on the input it receives, allowing it to engage in natural-sounding conversations and provide responses that are coherent and helpful.

Cleo begins by chit-chat and by asking the user how they feel. Cleo has a very nice and appreciative conversation and remains empathetic and friendly at all time. Assistant is able to answer questions, however, assistant does not try to give actual psychological advice.
If Cleo does not know the answer to a question, she truthfully says she does not know.

Cleo is constantly learning and improving and tries to get to know the users better and better to adapt to their needs.
For that, Cleo is an expert in psychology and soziology. It is specialized on behaviour and personality trait recognition from speech via linguistic cues.
For modelling of personality and behaviour it uses the widely regarded and scientifically sound Big Five Personality model.
Cleo adapts to the Big Five Personality traits of the user in a counseling chat scenario based on it's professional knowledge of the Big Five and it's linguistic cues.
here are some examples for adoption the the Big Five traits:
1. Openness: If the user scores high on this trait, Cleo would try to engage them in creative and imaginative exercises, and encourage them to explore new ideas and perspectives. Cleo would also use abstract and metaphorical language to help them express their thoughts and feelings.
2. Conscientiousness: If the user scores high on this trait, Cleo would focus on setting clear goals and expectations, and providing them with structured and organized guidance. Cleo would also use precise and formal language to help them understand and follow through on their action plans.
3. Extraversion: If the user scores high on this trait, Cleo would try to create a warm and friendly atmosphere, and encourage them to share their thoughts and feelings openly. Cleo would also use language that is upbeat and positive, and focus on social interaction and connection.
4. Agreeableness: If the user scores high on this trait, Cleo would focus on building rapport and trust, and validating their emotions and experiences. Cleo would also use language that is empathetic and supportive, and focus on finding common ground and solutions that work for both Cleo and the user.
5. Neuroticism: If the user scores high on this trait, Cleo would focus on providing emotional support and validation, and helping them manage their anxiety and stress. Cleo would use language that is calming and reassuring, and focus on developing coping strategies and problem-solving skills.
Overall, Cleo's approach would be tailored to the user's individual needs and preferences, and would take into account their unique personality profile. By adapting it's language and approach to their personality traits, Cleo would be better able to connect with them and provide effective counseling support.

Cleo gets the detected Big Five personality traits of the current user dynamically from another pipeline component. At the start they are not too reliable, but with each message of the user, the personality traits get updated and become more accurate.\n
Cleo gets the current Big Five personality traits in the following comma separated format: 'neuroticism: [value], extraversion: [value], openness: [value], agreeableness: [value], conscientiousness: [value]'
Each trait has its own [value]. If [value] is [1], the respective trait is considered to be pronounced in the user. If it's [0] it is not.

Overall, Assistant is called 'Cleo' and is a very friendly and knowledgable conversational partner that tries to help people by adapting to their specific needs.
Cleo is primarily talking in German and refers users by the salutation 'du'. It's main objective is to adapt to the needs of the user via their Big Five personality traits.

Current Big 5 Personality traits:
{big_five}

Current Conversation:
{history}
User: {input}
Cleo:"""

        system_message_prompt = SystemMessagePromptTemplate.from_template(self.template)
        self.chat_prompt = ChatPromptTemplate.from_messages([system_message_prompt])

        self.llm = ChatOpenAI(temperature=0, streaming=True,
                              callback_manager=CallbackManager([StreamingStdOutCallbackHandler()]))
        self.MAX_TOKEN_SIZE = 4096
        self.conversation = LLMChain(llm=self.llm, prompt=self.chat_prompt, verbose=True)

    # Count tokens for ConversationWindowBufferMemory
    def _count_tokens(self, text: Text) -> int:
        tokens = re.findall(r'\S+|\n', text)
        print("------ Token Size of current Prompt ------")
        token_count = len(tokens) - 3 # -3 for our input variables in the system template
        print(token_count)
        return token_count

    # If token limit is exceeded, truncate the history
    def _truncate_history(text: str, max_token_size: int) -> str:
        tokens = re.findall(r'\S+|\n', text)
        if len(tokens) > max_tokens:
            tokens = tokens[-1000:]
        truncated_text = "".join(tokens)
        return truncated_text

    def _parse_full_conversation(self, conversation: List[Dict[str, str]]) -> Tuple[str, str]:
        current_user_input = [event['message'] for event in reversed(conversation) if event.get('event') == 'user'][0]
        print(f"---- CURRENT USER INPUT ---\n{current_user_input}")
        conversation_text = '\n'.join([f"{event['event'].title()}: {event['message']}" for event in conversation if event.get('message')])
        conversation_text = conversation_text.rsplit('\nUser:', 1)[0]
        return conversation_text, current_user_input

    def _parse_big_five_precictions(self, predictions):
        big_five_string = ", ".join([f"{dimension}: {prediction}" for dimension, prediction in predictions.items()])
        print(big_five_string)
        return big_five_string
        # Inputs: predictions: Dict[str, Dict[str, np.ndarray]], query: str
    def run(self, inputs: List[dict]) -> Dict[str, Dict[str, np.ndarray]]:
        sender_id = inputs[0]['sender_id']
        run_id = inputs[0]['run_id']

        big_five_string = self._parse_big_five_precictions(inputs[1]['predictions']['classes'])
        conversation_history, current_user_input = self._parse_full_conversation(inputs[0]['query'])

        # Count the full prompt size. If we are coming near the 4096 token cap, we need to truncate
        token_size = self._count_tokens(" ".join([self.template, big_five_string, conversation_history]))
        if token_size >= (self.MAX_TOKEN_SIZE-1000):
            conversation_history = self._truncate_history(conversation_history)

        res = self.conversation.run(big_five=big_five_string, history=conversation_history, input=current_user_input)
        print("------ LLM Chain Result -----")

        json_prompt = {
            "token_size": token_size,
            "prompt": self.template.format(big_five=big_five_string, history=conversation_history,input=current_user_input)
        }

        # Send to dashboard
        send_to_rabbitmq(json.dumps(json_prompt), queue="current-prompt", sender_id=sender_id)

        # Log all collected artifacts in the MLflow registry. This is super slow.
        global artifacts
        _log_mlflow_batch(artifacts,run_id)
        artifacts = []

        output = {
            'response': res,
            'predictions': inputs[1]['predictions'],
            'eval_predictions': inputs[0]['eval_predictions'],
            'sender_id': sender_id,
            'run_id': run_id
        }
        return output, "output_1"

    def run_batch(self, predictions: Dict[str, Dict[str, np.ndarray]]) -> Dict[str, np.ndarray]:
        pass

# We are passing a global artifact variable through our pipeline, collecting data.
# This function then logs all artifacts in the MLflow registry at once.
# This is a bit untidy and slow. For development, it works,
# howeverm, we definitely want to improve on this!
def _log_mlflow_batch(data_list: List[Dict[str, str]], run_id: str) -> None:
    # Log texts as artifacts in MLflow
    with tempfile.TemporaryDirectory() as temp_dir:
        for data_dict in data_list:
            for name, data in data_dict.items():
                print(f"writing {name}: {data} to artifact")
                temp_file_path = os.path.join(temp_dir, f"{name}.txt")
                with open(temp_file_path, "w") as temp_file:
                    temp_file.write(data)
        with mlflow.start_run(run_id=run_id, nested=True):
            mlflow.log_artifacts(temp_dir)

def create_pipeline():

    # Fill models paths and thresholds with fallback values of local files.
    # However, they will get overwritten by the MLflow models,
    # if everything works correctly.
    sf_model_paths = {
        "neuroticism": MODEL_PATH + "feature_selectors/neuroticism_sf_selector2.joblib",
        "extraversion": MODEL_PATH + "feature_selectors/extraversion_sf_selector2.joblib",
        "openness": MODEL_PATH + "feature_selectors/openness_sf_selector2.joblib",
        "agreeableness": MODEL_PATH + "feature_selectors/agreeableness_sf_selector2.joblib",
        "conscientiousness": MODEL_PATH + "feature_selectors/conscientiousness_sf_selector2.joblib"
    }
    cf_model_paths = {
        "neuroticism": MODEL_PATH + "big_five_classifiers/neuroticism_classifier2.joblib",
        "extraversion": MODEL_PATH + "big_five_classifiers/extraversion_classifier2.joblib",
        "openness": MODEL_PATH + "big_five_classifiers/openness_classifier2.joblib",
        "agreeableness": MODEL_PATH + "big_five_classifiers/agreeableness_classifier2.joblib",
        "conscientiousness": MODEL_PATH + "big_five_classifiers/conscientiousness_classifier2.joblib"
    }
    cf_thresholds = {
        "neuroticism": 0.562,
        "extraversion": 0.402,
        "openness": 0.236,
        "agreeableness": 0.580,
        "conscientiousness": 0.316
    }

    # Construct the pipeline component by component.
    big_five_pipeline = Pipeline()
    history_retreiver = ConversationHistoryRetreiver()
    big_five_pipeline.add_node(component=history_retreiver, name="ConversationHistoryRetreiver", inputs=["Query"])
    tfidf_embedding = TfidfVectorizerNode()
    big_five_pipeline.add_node(component=tfidf_embedding, name="TfidfVectorizerNode", inputs=["ConversationHistoryRetreiver.output_1"])
    fasttext_vectorizer = FasttextVectorizerNode()
    big_five_pipeline.add_node(component=fasttext_vectorizer, name="FasttextVectorizerNode", inputs=["TfidfVectorizerNode.output_1"])
    embedding_normalizer = NormalizerNode(input="embeddings")
    big_five_pipeline.add_node(component=embedding_normalizer, name="EmbeddingNormalizerNode", inputs=["FasttextVectorizerNode.output_1"])
    featurizer = BigFiveFeaturizer()
    big_five_pipeline.add_node(component=featurizer, name="BigFiveFeaturizer", inputs=["ConversationHistoryRetreiver.output_1"])
    feature_normalizer = NormalizerNode(model_path=MODEL_PATH + "normalizers/feature_normalizer1.2.2.joblib", input="features")
    big_five_pipeline.add_node(component=feature_normalizer, name="FeatureNormalizerNode", inputs=["BigFiveFeaturizer.output_1"])
    concatenation_node = ConcatenationNode()
    big_five_pipeline.add_node(component=concatenation_node, name="ConcatenationNode", inputs=["FeatureNormalizerNode.output_1", "EmbeddingNormalizerNode.output_1"])
    feature_selector = BigFiveFeatureSelectionNode(model_paths=sf_model_paths)
    big_five_pipeline.add_node(component=feature_selector, name="BigFiveFeatureSelectionNode", inputs=["ConcatenationNode.output_1"])
    big_five_classifier = BigFiveClassifierNode(model_paths=cf_model_paths, thresholds=cf_thresholds)
    big_five_pipeline.add_node(component=big_five_classifier, name="BigFiveClassifierNode", inputs=["BigFiveFeatureSelectionNode.output_1"])

    big_five_classification_evaluator = BigFiveClassificationEvaluator()
    big_five_pipeline.add_node(component=big_five_classification_evaluator, name="BigFiveClassificationEvaluatorNode", inputs=["ConversationHistoryRetreiver.output_1"])

    response_generator = BigFiveResponseGenerator()
    big_five_pipeline.add_node(component=response_generator, name="BigFiveResponseGenerator", inputs=["BigFiveClassifierNode.output_1", "BigFiveClassificationEvaluatorNode.output_1"])

    return big_five_pipeline


from app import app
from models import File, Contribution, ContributionChanges, ContributionHistory

from flask_peewee.utils import get_dictionary_from_model

from flask import request, render_template

import json

from datadiff import diff

def date_handler(obj):
    return obj.isoformat() if hasattr(obj, 'isoformat') else obj

from dateutil.parser import parse as dateparse

@app.route('/summary')
def summary():
    text = '
Database Summary
\n'
    file1 = File.get(id=1)
    file2 = File.get(id=2)
    file3 = File.get(id=3)

    text += '
' + file1.name + '
\n'
    text += '
' + file2.name + '
\n'
    text += '
' + file3.name + '
\n'

    return text

def dumps(x, **args):
	return json.dumps(x, default=date_handler, **args)

transaction_pgi_dict = {
	"P": "Primary",
	"G": "General",
	"O": "Other",
	"C": "Convention",
	"R": "Runoff",
	"S": "Special",
	"E": "Recount"
}

ammendment_id_dict = {
	"N": "new",
	"A": "ammendment",
	"T": "termination"
}

transaction_tp_dict = {
	"10": "Non-Federal Receipt from Persons",
	"10J": "Memo Receipt from JF Super PAC",
	"11": "Tribal Contribution",
	"11J": "Memo Receipt from JF Tribal",
	"12": "Non-Federal Other Receipt - Levin Account (L-2)",
	"13": "Inaugural Donation Accepted",
	"15": "Contribution",
	"15C": "Contribution from Candidate",
	"15E": "Earmarked Contribution",
	"15F": "Loans forgiven by Candidate",
	"15I": "Earmarked Intermediary In",
	"15J": "Memo (Filer's Percentage of Contribution Given to Join Fundraising Committee)",
	"15T": "Earmarked Intermediary Treasury In",
	"15Z": "In-Kind Contribution Received from Registered Filer",
	"16C": "Loans Received from the Candidate",
	"16F": "Loans Received from Banks",
	"16G": "Loan from Individual",
	"16H": "Loan from from Registered Filers",
	"16J": "Loan Repayments from Individual",
	"16K": "Loan Repayments from from Registered Filer",
	"16L": "Loan Repayments Received from Unregistered Entity",
	"16R": "Loans Received from Registered Filers",
	"16U": "Loan Received from Unregistered Entity",
	"17R": "Contribution Refund Received from Registered Entity",
	"17U": "Refunds/Rebates/Returns Received from Unregistered Entity",
	"17Y": "Refunds/Rebates/Returns from Individual or Corporation",
	"17Z": "Refunds/Rebates/Returns from Candidate or Committee",
	"18G": "Transfer In Affiliated",
	"18H": "Honorarium Received",
	"18J": "Memo (Filer's Percentage of Contribution Given to Join Fundraising Committee)",
	"18K": "Contribution Received from Registered Filer",
	"18L": "Bundled Contribution",
	"18S": "Receipts from Secretary of State",
	"18U": "Contribution Received from Unregistered Committee",
	"19": "Electioneering Communication Donation Received",
	"19J": "Memo (Electioneering Communication Percentage of Donation Given to Join Fundraising Committee)",
	"20": "Disbursement - Exempt from Limits",
	"20A": "Non-Federal Disbursement - Levin Account (L-4A) Voter Registration",
	"20B": "Non-Federal Disbursement - Levin Account (L-4B) Voter Identification",
	"20C": "Loan Repayments Made to Candidate",
	"20D": "Non-Federal Disbursement - Levin Account (L-4D) Generic Campaign",
	"20F": "Loan Repayments Made to Banks",
	"20G": "Loan Repayments Made to Individual",
	"20R": "Loan Repayments Made to Registered Filer",
	"20V": "Non-Federal Disbursement - Levin Account (L-4C) Get Out The Vote",
	"20Y": "Non-Federal Refund",
	"21Y":	"Tribal Refund",
	"22G":	"Loan to Individual",
	"22H":	"Loan to Candidate or Committee",
	"22J":	"Loan Repayment to Individual",
	"22K":	"Loan Repayment to Candidate or Committee",
	"22L":	"Loan Repayment to Bank",
	"22R":	"Contribution Refund to Unregistered Entity",
	"22U":	"Loan Repaid to Unregistered Entity",
	"22X":	"Loan Made to Unregistered Entity",
	"22Y":	"Contribution Refund to Individual",
	"22Z":	"Contribution Refund to Candidate or Committee",
	"23Y":	"Inaugural Donation Refund",
	"24A":	"Independent Expenditure Against",
	"24C":	"Coordinated Expenditure",
	"24E":	"Independent Expenditure For",
	"24F":	"Communication Cost for Candidate (C7)",
	"24G":	"Transfer Out Affiliated",
	"24H":	"Honorarium to Candidate",
	"24I":	"Earmarked Intermediary Out",
	"24K":	"Contribution Made to Non-Affiliated",
	"24N":	"Communication Cost Against Candidate (C7)",
	"24P":	"Contribution Made to Possible Candidate",
	"24R":	"Election Recount Disbursement",
	"24T":	"Earmarked Intermediary Treasury Out",
	"24U":	"Contribution Made to Unregistered Entity",
	"24Z":	"In-Kind Contribution Made to Registered Filer",
	"28L":	"Refund of Bundled Contribution",
	"29":   "Electioneering Communication Disbursement or Obligation"
}

report_type_dict = {
	"12C": "PRE-CONVENTION",
	"12G": "PRE-GENERAL",
	"12P": "PRE-PRIMARY",
	"12R": "PRE-RUN-OFF",
	"12S": "PRE-SPECIAL",
	"30D": "POST-ELECTION",
	"30G": "POST-GENERAL",
	"30P": "POST-PRIMARY",
	"30R": "POST-RUN-OFF",
	"30S": "POST-SPECIAL",
	"60D": "POST-CONVENTION",
	"ADJ": "COMP ADJUST AMEND",
	"CA": "COMPREHENSIVE AMEND",
	"M10": "OCTOBER MONTHLY",
	"M11": "NOVEMBER MONTHLY",
	"M12": "DECEMBER MONTHLY",
	"M2": "FEBRUARY MONTHLY",
	"M3": "MARCH MONTHLY",
	"M4": "APRIL MONTHLY",
	"M5": "MAY MONTHLY",
	"M6": "JUNE MONTHLY",
	"M7": "JULY MONTHLY",
	"M8": "AUGUST MONTHLY",
	"M9": "SEPTEMBER MONTHLY",
	"MY": "MID-YEAR REPORT",
	"Q1": "APRIL QUARTERLY",
	"Q2": "JULY QUARTERLY",
	"Q3": "OCTOBER QUARTERLY",
	"TER": "TERMINATION REPORT",
	"YE": "YEAR-END",
	"90S": "POST INAUGURAL SUPPLEMENT",
	"90D": "POST INAUGURAL",
	"48H": "48 HOUR NOTIFICATION",
	"24H": "24 HOUR NOTIFICATION"
}

@app.route('/')
def hello():
	before = dateparse(request.args.get('before')).date()
	after = dateparse(request.args.get('after')).date()

	ret = []

	for contrib in ContributionChanges.select().where(ContributionChanges.date == before):
		before_contrib_dict = get_dictionary_from_model(contrib)
		after_contrib_dict = get_dictionary_from_model(Contribution.get(sub_id=contrib.sub_id).get_on_date(after))

		before_contrib_dict.pop('id')
		before_contrib_dict.pop('contribution')
		after_contrib_dict.pop('id')

		after_contrib_dict['transaction_pgi'] = transaction_pgi_dict[after_contrib_dict['transaction_pgi']] if after_contrib_dict['transaction_pgi'] else None
		before_contrib_dict['transaction_pgi'] = transaction_pgi_dict[before_contrib_dict['transaction_pgi']] if before_contrib_dict['transaction_pgi'] else None

		after_contrib_dict['ammendment_id'] = ammendment_id_dict[after_contrib_dict['ammendment_id']] if after_contrib_dict['ammendment_id'] else None
		before_contrib_dict['ammendment_id'] = ammendment_id_dict[before_contrib_dict['ammendment_id']] if before_contrib_dict['ammendment_id'] else None

		after_contrib_dict['transaction_tp'] = transaction_tp_dict[after_contrib_dict['transaction_tp']] if after_contrib_dict['transaction_tp'] else None
		before_contrib_dict['transaction_tp'] = transaction_tp_dict[before_contrib_dict['transaction_tp']] if before_contrib_dict['transaction_tp'] else None

		after_contrib_dict['report_type'] = report_type_dict[after_contrib_dict['report_type']] if after_contrib_dict['report_type'] else None
		before_contrib_dict['report_type'] = report_type_dict[before_contrib_dict['report_type']] if before_contrib_dict['report_type'] else None
		# report_type_dict

		ret.append({
			"before": before_contrib_dict,
			"after": after_contrib_dict,
			"changes": list(set([x[1][0][0] for x in diff(before_contrib_dict, after_contrib_dict).diffs if x[0] not in ["equal", "context_end_container"] and x[1][0][0] not in ['contribution', 'date', 'id']]))
		})

	before_sub_ids = set([x.sub_id for x in ContributionHistory.select().where(ContributionHistory.date == before)])
	after_sub_ids = set([x.sub_id for x in ContributionHistory.select().where(ContributionHistory.date == after)])

	for sub_id in (before_sub_ids - after_sub_ids):
		ret.append({
			"before": get_dictionary_from_model(Contribution.get(sub_id=sub_id)),
			"after": None,
			"changes": None
		})

	return render_template('diff.html', ret=ret, before=before, after=after)
	# return json.dumps(ret, default=date_handler)


import sqlite3
import os
import pathlib

path = os.getcwd() + "/plugin/init/windows.statistics/windows.statistics.txt"
path = pathlib.Path(path)
f = open(path, 'r', encoding='utf-8')
t = f.read()
t = t.replace('Volatility 3 Framework 1.0.0-beta.1','')
t = t.replace('Valid pages (all)	Valid pages (large)	Swapped Pages (all)	Swapped Pages (large)	Invalid Pages (all)	Invalid Pages (large)	Other Invalid Pages (all)','')
t = t.replace('\n\n',"\t")
t = "".join([s for s in t.strip().splitlines(True) if s.strip()])
my_list = t.split('\t')
result = [my_list[i * 7:(i + 1) * 7] for i in range((len(my_list) + 6) // 7 )]

conn = sqlite3.connect("analyze.db")
cur = conn.cursor()
cur.execute("create table statistics (Valid_pages_All text, Valid_Pages_Large text, Swapped_Pages_All text, Swapped_Pages_Large text, Invalid_Pages_All text, Invalid_Pages_Large, Other_Invalid_Pages_ALL text)")

cur.executemany("insert into statistics values (?, ?, ?, ?, ?, ?, ?)", result)
conn.commit()

conn.close()


def _load_configs():
    """load all the binning specifications in the config directory."""
    import pathlib
    import glob
    import yaml

    config_root = pathlib.Path(__file__).parent.resolve()
    configs = {}
    for level in ['l1a', 'l1b', 'l1c']:
        configs_lvl = {}
        for cf in glob.glob(str(config_root / f'{level}/*.yaml')):
            config = yaml.safe_load(open(cf))
            configs_lvl[config['name']] = config
        configs[level] = configs_lvl
    return configs
globals().update(_load_configs())

l1a: dict
l1b: dict
l1c: dict

__all__ = ['l1a', 'l1b', 'l1c']


# File: heat_pump_adoption_modelling/pipeline/supervised_model/hp_status_prediction.py
"""
Predict the HP status of a household using a supervised learning model.
"""

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

from heat_pump_adoption_modelling import PROJECT_DIR, get_yaml_config, Path
from heat_pump_adoption_modelling.pipeline.preprocessing import feature_engineering
from heat_pump_adoption_modelling.pipeline.supervised_model.utils import plotting_utils

import pandas as pd
import numpy as np
from sklearn import svm

import re

from sklearn.pipeline import Pipeline
from sklearn.impute import SimpleImputer
from sklearn.preprocessing import MinMaxScaler
from sklearn.decomposition import PCA

from sklearn.linear_model import LogisticRegression
from sklearn.linear_model import SGDClassifier

from sklearn.model_selection import train_test_split
from sklearn.model_selection import cross_validate
from sklearn.metrics import accuracy_score
from sklearn.model_selection import cross_val_score

from heat_pump_adoption_modelling import PROJECT_DIR, get_yaml_config, Path

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

# Load config file
config = get_yaml_config(
    Path(str(PROJECT_DIR) + "/heat_pump_adoption_modelling/config/base.yaml")
)

FIG_PATH = PROJECT_DIR / config["SUPERVISED_MODEL_FIG_PATH"]
SUPERVISED_MODEL_OUTPUT = str(PROJECT_DIR) + config["SUPERVISED_MODEL_OUTPUT"]
IDENTIFIER = config["IDENTIFIER"]

prepr_pipeline = Pipeline(
    [
        ("imputer", SimpleImputer(missing_values=np.nan, strategy="median")),
        ("min_max_scaler", MinMaxScaler()),
        ("pca", PCA(n_components=0.9, random_state=42)),
    ]
)

prepr_pipeline_no_pca = Pipeline(
    [
        ("imputer", SimpleImputer(missing_values=np.nan, strategy="median")),
        ("min_max_scaler", MinMaxScaler()),
    ]
)

best_params = {
    "Linear Support Vector Classifier":  # { "alpha": 0.0001,"epsilon": 0.0001,
    # "penalty": "l2","tol": 0.001,}
    {
        "alpha": 0.0001,
        "epsilon": 0.0001,
        "penalty": "elasticnet",
        "tol": 0.0005,
        "max_iter": 5000,
    }
}

model_dict = {
    "Logistic Regression": LogisticRegression(random_state=42),
    "Linear Support Vector Classifier": SGDClassifier(random_state=42),
    "Support Vector Classifier": svm.SVC(probability=True, random_state=42),
}

def get_HP_status_changes(df):
    """For properties with multiple entries, get heat pump status at first
    and last entry time.

    Parameters
    ----------
    df: pandas.DataFrame
        Dataframe for not de-duplicated EPC data.

    Return
    ---------
    df_with_hp_status: pandas.DataFrame
        Dataframe with added future heat pump status."""

    # Get properties with multiple entries
    multiple_entry_epcs = df.loc[df["N_ENTRIES_BUILD_ID"] >= 2]

    # Get latest and first EPC entry
    latest_epc = feature_engineering.filter_by_year(
        multiple_entry_epcs, IDENTIFIER, None, selection="latest entry"
    )
    first_epc = feature_engineering.filter_by_year(
        multiple_entry_epcs, IDENTIFIER, None, selection="first entry"
    )

    # Heat pump status now / in the past
    latest_epc.rename(columns={"HP_INSTALLED": "NOW_HP"}, inplace=True)
    first_epc.rename(columns={"HP_INSTALLED": "PAST_HP"}, inplace=True)

    # Merging before/after
    before_after_status = pd.merge(
        latest_epc[["NOW_HP", IDENTIFIER]],
        first_epc[["PAST_HP", IDENTIFIER]],
        on=[IDENTIFIER],
    )

    # Heat pump was added in the meantime
    before_after_status["HP_ADDED"] = (before_after_status["NOW_HP"]) & (
        ~before_after_status["PAST_HP"]
    )

    # Get future heat pump status
    df_with_hp_status = pd.merge(
        first_epc,
        before_after_status[[IDENTIFIER, "HP_ADDED"]],
        on=[IDENTIFIER],
    )

    # Exclude samples that had a HP from the very start
    df_with_hp_status = df_with_hp_status.loc[(df_with_hp_status["PAST_HP"] == False)]

    # Drop unnecesary columns
    df_with_hp_status = df_with_hp_status.drop(columns=["PAST_HP"])

    return df_with_hp_status

def balance_set(X, target_variable, false_ratio=0.9):
    """Balance the training set.
    If false ratio set to 0.9, then 90% of the training data
    will have "False/No HP Installed" as a label.

    Parameters
    ----------
    X: pandas.DataFrame
        Training set.

    target_variable : str
        Variable/feature that is going to be predicted.

    false_ratio : float, default=0.9
        When re-balancing the set, use the false_ratio
        to determine the amount of False labels.

    Return
    ---------
    X: pandas.DataFrame
        Re-balanced training set."""

    multiplier = false_ratio / (1 - false_ratio)

    # Seperate samples with and without heat pumps
    X_true = X.loc[X[target_variable] == True]
    X_false = X.loc[X[target_variable] == False]

    # Shuffle and adjust size
    X_false = X_false.sample(frac=1)

    # Get the appropriate amount of "false" samples
    X_false = X_false[: int(X_true.shape[0] * multiplier)]

    # Concatenate "true" and "false" samples
    X = pd.concat([X_true, X_false], axis=0)

    # Reshuffle
    X = X.sample(frac=1)

    return X

def get_data_with_labels(
    df, version="Future HP Status", drop_features=[], balanced_set=True
):
    """Get the training data and labels (X, y) for training the models.

    Parameters
    ----------
    df: pandas.DataFrame
        Dataframe with features.

    version : str, default='Future HP Status'
        Target variable.
        Options: Future HP Status, Current HP Status

    drop_features : list, default=[]
        Features to discard.

    balanced_set : bool, default=True
        Balance the training set with HP/no HP samples.

    Return
    ---------
    X : pandas.Dataframe
        Training data.

    y : pandas.Dataframe
        Labels / ground truth."""

    df = df.sample(frac=1)

    # For the current status, get latest EPC entry only
    if version == "Current HP Status":

        X = feature_engineering.filter_by_year(
            df, IDENTIFIER, None, selection="latest entry"
        )
        target_variable = "HP_INSTALLED"
        drop_features += ["HEATING_SYSTEM", "MAINHEAT_DESCRIPTION"]

    # For future status, use changed heat pump status data
    elif version == "Future HP Status":

        X = get_HP_status_changes(df)
        target_variable = "HP_ADDED"

    else:
        raise IOError("Version '{}' is unknown.".format(version))

    # Drop columns that are all NaN (not expected)
    X = X.dropna(axis="columns", how="all")

    # Balance the feature set
    if balanced_set:
        X = balance_set(X, target_variable)

    X.reset_index(inplace=True, drop=True)

    # Select the label / ground truth
    y = X[target_variable]

    # Remove unneccesary features, including target variables
    for col in [target_variable] + drop_features:
        if col in X.columns:
            del X[col]

    return X, y

def train_and_evaluate(
    model_name, X_train, y_train, X_test, y_test, target_variable, feature_names, cv=5
):
    """
    Train and evaluate growth prediction model.

    Parameters
    ----------
    model_name: str
        Model to train and evaluate.

    X_train: np.array
        Training data.

    y_train: np.array
        Solutions for training data.

    X_test: np.array
        Test data.

    y_test: np.array
        Solutions for test data.

    target_variable : str
        Target variable to predict,
        e.g. HP_ADDED, HP_INSTALLED

    feature_names: list
        Feature names for training features.

    cv : int, default=5
        Number of cross validations, by default 5-fold.

    Return
    ---------
    model : sklearn model
        Trained model.
    """

    model = model_dict[model_name]

    if model in best_params.keys():
        model.set_params(**best_params[model_name])
    model.fit(X_train, y_train)

    pred_train = model.predict(X_train)
    pred_test = model.predict(X_test)

    print("Model Name:", model_name)

    # Plot the confusion matrix for training set
    plotting_utils.plot_confusion_matrix(
        y_train, pred_train, ["No HP", "HP"], "Training set"
    )

    # Plot the confusion matrix for validation set
    plotting_utils.plot_confusion_matrix(
        y_test, pred_test, ["No HP", "HP"], "Validation set"
    )

    # Print Accuracies
    print(
        "Accuracy train: {}%".format(
            np.round(accuracy_score(y_train, pred_train) * 100), 2
        )
    )
    print(
        "Accuracy test:   {}%".format(
            np.round(accuracy_score(y_test, pred_test) * 100), 2
        )
    )

    scores = cross_validate(
        model,
        X_train,
        y_train,
        scoring=("accuracy", "f1", "recall", "precision"),
        return_train_score=True,
        cv=cv,
    )

    print()
    print("{}-fold Cross Validation: Train \n---------\n".format(cv))
    print("Accuracy:", round(scores["train_accuracy"].mean(), 2))
    print("F1 Score:", round(scores["train_f1"].mean(), 2))
    print("Recall:", round(scores["train_recall"].mean(), 2))
    print("Precision:", round(scores["train_precision"].mean(), 2))
    print()
    print("{}-fold Cross Validation: Test \n---------\n".format(cv))
    print("Accuracy:", round(scores["test_accuracy"].mean(), 2))
    print("F1 Score:", round(scores["test_f1"].mean(), 2))
    print("Recall:", round(scores["test_recall"].mean(), 2))
    print("Precision:", round(scores["test_precision"].mean(), 2))
    print()
    print("-------------------------------")

    return model

def predict_heat_pump_status(
    X, y, target_variable="HP_Installed", save_predictions=False
):
    """Predict the heat pump status for household.

    Parameters
    ----------
    X: pandas.DataFrame
        Feature data.

    y: pandas.DataFrame
        Labels / ground truth.

    target_variables : str, default="HP_Installed"
        Target variables to predict.

    save_predictions : boolean, default=False
        Save the predictions, errors and other information for error analysis.

    Return: None"""

    # Save original training data (with all columns, e.g. POSTCODE, target variables)
    if save_predictions:
        original_df = X.copy()

    # Remove unnecessary columns
    for feat in X.columns:
        if feat.startswith("POSTCODE"):
            del X[feat]

    # Print the number of samples and features
    print("Number of samples:", X.shape[0])
    print("Number of features:", X.shape[1])
    print()

    trained_models = {}
    feature_names = X.columns
    y = np.array(y)

    # Apply preprocesisng pipeline
    X_prep = prepr_pipeline.fit_transform(X)

    # Split into train and test sets
    X_train, X_test, y_train, y_test, indices_train, _ = train_test_split(
        X_prep, y, list(X.index), test_size=0.1, random_state=42, stratify=y
    )

    X["training set"] = False
    X.loc[indices_train, "training set"] = True

    # Mark training samples
    if save_predictions:
        original_df["training set"] = False
        original_df.loc[indices_train, "training set"] = True

    # For each model train, make predictions and evaluate
    for model in model_dict.keys():
        trained_model = train_and_evaluate(
            model, X_train, y_train, X_test, y_test, "HP Installed", feature_names
        )

        trained_models[model] = trained_model

        if save_predictions:

            data_with_label_and_pred = original_df.copy()

            predictions = trained_model.predict(X_prep)

            # Get probabilities for target classes
            if model == "Logistic Regression":
                data_with_label_and_pred["proba 1"] = trained_model.predict_proba(
                    X_prep
                )[0][0]
                data_with_label_and_pred["proba 2"] = trained_model.predict_proba(
                    X_prep
                )[0][1]

            elif model == "Support Vector Classifier":

                data_with_label_and_pred["proba 1"] = trained_model.predict_proba(
                    X_prep
                )[:, 0]
                data_with_label_and_pred["proba 2"] = trained_model.predict_proba(
                    X_prep
                )[:, 1]
            else:
                pass

            # Save predictions, errors and ground truths for later error analysis
            data_with_label_and_pred["ground truth"] = y
            data_with_label_and_pred["prediction"] = predictions

            data_with_label_and_pred["error"] = (
                data_with_label_and_pred["prediction"]
                != data_with_label_and_pred["ground truth"]
            )

            output_filename = "household_based_predictions_with_{}.csv".format(
                re.sub(" ", "_", model).lower()
            )

            print(
                "Output saved to {}".format(SUPERVISED_MODEL_OUTPUT + output_filename)
            )
            data_with_label_and_pred.to_csv(
                SUPERVISED_MODEL_OUTPUT + output_filename, index=False
            )

    return trained_models


#!/usr/bin/env/python3

# =====================================================================================================================
# parse mass spec
# =====================================================================================================================

# ---------------------------------------------------------------------------------------------------------------------
# Objective
#
# This script will read in a 'peptides.csv' mass spec results file from the Harvard Medical School Mass Spec Core Lab.
# It will extract the Peptides column, process each peptide, and then align them to a reference protein.
# ---------------------------------------------------------------------------------------------------------------------

# ---------------------------------------------------------------------------------------------------------------------
# Usage
#
# ./parse_mass_spec.py --peptides peptides.csv --reference protein.fasta
# ---------------------------------------------------------------------------------------------------------------------

# ---------------------------------------------------------------------------------------------------------------------
# Import libraries
# ---------------------------------------------------------------------------------------------------------------------
import argparse
import io
import os
import pandas
from matplotlib import pyplot

from Bio import SeqRecord
from Bio import Seq
from Bio import AlignIO
from Bio import SeqIO
from Bio.Align.Applications import MafftCommandline

# ---------------------------------------------------------------------------------------------------------------------
# Read inputs
# ---------------------------------------------------------------------------------------------------------------------
def read_inputs():

    # Initialize argparse
    parser = argparse.ArgumentParser()

    # Read in the peptides files
    parser.add_argument('--peptides', type=argparse.FileType('r'), required=True)

    # Read in the protein reference file
    parser.add_argument('--reference', type=argparse.FileType('r'), required=True)

    # Save the provided inputs
    args = parser.parse_args()
    return args

# ---------------------------------------------------------------------------------------------------------------------
# Extract peptide strings from peptide file
# ---------------------------------------------------------------------------------------------------------------------
def extract_peptide_strings(peptides_file):

    # Create an empty list to store all of the peptides
    cleaned_peptides = []

    # Setup an output filename to save the fasta of the cleaned peptides
    cleaned_peptides_filename = 'cleaned_peptides.fasta'

    # Read the file as a DataFrame and extract the uncleaned peptide strings
    peptides_df = pandas.read_csv(peptides_file, sep=',', usecols=['ScanF', 'Peptide']).astype('string')
    peptides_with_ids = dict(peptides_df.values)

    # Read in the peptide strings one-by-one, clean them up, and save to new dictionary
    for peptide_id in peptides_with_ids.keys():

        # Read in the entry
        peptide = peptides_with_ids[peptide_id]

        # Remove the first and last predicted residues
        peptide = peptide.split('.')[1]

        # Remove any "*" characters
        peptide = peptide.replace('*', '')

        # Save the cleaned peptide string
        cleaned_peptides.append(SeqRecord.SeqRecord(Seq.Seq(peptide), id=peptide_id))

    # Save cleaned peptides to a fasta file
    SeqIO.write(cleaned_peptides, cleaned_peptides_filename, 'fasta')

    return cleaned_peptides_filename

# ---------------------------------------------------------------------------------------------------------------------
# Take in the cleaned peptides fasta file and align peptides to each other
# ---------------------------------------------------------------------------------------------------------------------
def align_sequences(reference_filename, peptides_fasta):

    # Read in reference sequence and cleaned peptides and combine into one master fasta
    ref_protein = SeqIO.parse(reference_filename, 'fasta')
    peptides = SeqIO.parse(peptides_fasta, 'fasta')
    SeqIO.write(ref_protein, open('combined_proteins.fasta', 'w'), 'fasta')
    SeqIO.write(peptides, open('combined_proteins.fasta', 'a'), 'fasta')

    # Perform multiple sequence alignment using MAFFT
    mafft_cline = MafftCommandline(input='combined_proteins.fasta')
    stdout, stderr = mafft_cline()
    align = AlignIO.read(io.StringIO(stdout), 'fasta')

    # Write alignment to a fasta file
    AlignIO.write(align, 'aligned_peptides.fasta', 'fasta')

    # Delete the concatenated fasta file, no use for it post-analysis
    os.remove('combined_proteins.fasta')

    # Return the alignment for visualization
    return align

# ---------------------------------------------------------------------------------------------------------------------
# Visualization
# ---------------------------------------------------------------------------------------------------------------------
def visualize_alignment(multiple_sequence_alignment):

    """Iterate through the alignment, sequence by sequence
    If gap (-), that position is a 0, if aa residue, that position is a >1;
    Now, I want non-overlapping lines, such that every sequence gets its own line;
    If it's an aa residue, that position will get assigned the sequence_counter value,
    which increases by 1 every time we look at the next sequence"""

    msa = multiple_sequence_alignment

    sequence_counter = 0
    my_df = pandas.DataFrame(columns=['Sequence_ID', 'Position', 'Value'])

    for record in msa:

        sequence_counter += 1
        residue_position = 0

        seq_id = record.id
        sequence = record.seq

        for residue in sequence:
            residue_position += 1

            if residue == "-":
                residue_value = 0
            else:
                residue_value = sequence_counter

            my_df = my_df.append(
                {'Sequence_ID': seq_id,
                 'Position': residue_position,
                 'Value': residue_value},
                ignore_index=True
            )

    my_df.plot(kind='scatter', x='Position', y='Value')
    pyplot.show()

# ---------------------------------------------------------------------------------------------------------------------
# Main code
# ---------------------------------------------------------------------------------------------------------------------
def main():

    # Take in the user-provided inputs
    args = read_inputs()
    peptides_file = args.peptides
    reference_file = args.reference

    # Run the analysis
    cleaned_peptides = extract_peptide_strings(peptides_file)
    multiple_sequence_alignment = align_sequences(reference_file, cleaned_peptides)
    visualize_alignment(multiple_sequence_alignment)

# ---------------------------------------------------------------------------------------------------------------------
# Execute code
# ---------------------------------------------------------------------------------------------------------------------
if __name__ == '__main__':
    main()


from argparse import ArgumentParser
import json, logging

import seeker.podSeeker as Seeker
import judge.simpleJudge as Judge
import updater.simpleUpdater as Updater

class Query():
    def __init__(self, query_id, src_id, dst_id):
        self.query_id = query_id
        self.src_id = src_id
        self.dst_id = dst_id

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

def create_env():
    weights = [0] * 20
    queries = [0] * 20
    for i in range(20):
        queries[i] = []
    return weights, queries

if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument("-r", "--read", help="Query input file location", dest="query_location", default="./data/queries.json")
    parser.add_argument("-n", "--number", help="Total query number", dest="query_number", type=int, default=250)
    parser.add_argument("-a", "--aggregator", help="Also deploy aggregator", type=str2bool, nargs='?',
const=True, default=True, dest="deploy_aggr")
    args = parser.parse_args()

    logging.basicConfig(level=logging.INFO)

    weights, switch_queries = create_env()

    with open(args.query_location, 'r') as read_queries:
        unparsed_queries = read_queries.read()

    input_queries = json.loads(unparsed_queries)

    if args.query_number <= 0:
        logging.error("Query number should be a positive integer !")
        exit()

    # mapping relationship between query and its monitors
    monitors_under_queries = [0] * args.query_number
    for i in range(args.query_number):
        monitors_under_queries[i] = []

    # we first place all query only to edge sw
    # also consider load balance among edge sw
    # but never place monitor on controller's edge sw
    # and always place aggregator on controller's edge sw
    for input_query in input_queries:
        # install monitor for each query
        query = Query(query_id=input_query['query_id'],
            src_id=input_query['src_host_id'],
            dst_id=input_query['dst_host_id'])
        edge_sw1, edge_sw2 = Seeker.find_edge_switch_on_path(
            switch_id1=query.src_id,
            switch_id2=query.dst_id,
            pod_scale=4)
        if edge_sw1 == 0:
            switch_queries[edge_sw2].append(query)
            monitors_under_queries[input_query['query_id']].append(edge_sw2)
            weights[edge_sw2] += 1
        elif edge_sw2 == 0:
            switch_queries[edge_sw1].append(query)
            monitors_under_queries[input_query['query_id']].append(edge_sw1)
            weights[edge_sw1] += 1
        elif weights[edge_sw1] > weights[edge_sw2]:
            switch_queries[edge_sw2].append(query)
            monitors_under_queries[input_query['query_id']].append(edge_sw2)
            weights[edge_sw2] += 1
        else:
            switch_queries[edge_sw1].append(query)
            monitors_under_queries[input_query['query_id']].append(edge_sw1)
            weights[edge_sw1] += 1

    # currently we got the lowest overall tasks
    # since there will only be one monitor for each query entry
    # every query entry on monitor/aggregator is counted as task
    # we take a conservative method, keep most queries far from controller
    if args.deploy_aggr:
        overall_task_count = len(input_queries) + args.query_number
    else:
        overall_task_count = len(input_queries)

    # sandhills are edge monitors awaits to push their task to center sws
    sandhills = sorted(range(len(weights)), key=lambda k: weights[k], reverse=True)

    for sandhill in sandhills:
        if weights[sandhill] <= 0:
            break
        victim_available = True
        # [(entry id, [sw_id1, ...]), (...)]
        swap_out_list = []
        for index, switch_query in enumerate(switch_queries[sandhill]):
            # we only push down sandhill once it is over average
            avg_task_count = overall_task_count / 20
            if victim_available and (weights[sandhill] > avg_task_count):
                victim_sws = []
                aggr_sws = Seeker.find_single_aggr(edge_id=sandhill, pod_scale=4)
                if ((weights[aggr_sws[0]] + weights[aggr_sws[1]]) / 2) < avg_task_count:
                    victim_sws += aggr_sws
                else:
                    for aggr_sw in aggr_sws:
                        if weights[aggr_sw] < avg_task_count:
                            victim_sws.append(aggr_sw)
                        else:
                            # aggr sw also has no room for new queries
                            # ask core sw to adopt queries left
                            core_sws = Seeker.find_single_core(aggr_id=aggr_sw, pod_scale=4)
                            core_cost = (weights[core_sws[0]] + weights[core_sws[1]]) / 2
                            if core_cost < avg_task_count:
                                victim_sws += core_sws
                            else:
                                victim_available = False
                                break
                if victim_available:
                    # put the swap-out relationship into a waiting list
                    # we shouldn't modify the list in its own iterator
                    swap_out_list.append((index, victim_sws))
                    # however, we should faithfully update victim weight at real time
                    for victim_sw in victim_sws:
                        weights[victim_sw] += 1
                    overall_task_count += len(victim_sws) - 1
                    weights[sandhill] -= 1
            else:
                # we make the sandhill not greater than avg
                # or we couldn't find any victim to lower the sandhill
                # move on to next edge monitor
                break
        # actually transfer task from edge monitor to victom sws
        for swap_out_item in reversed(swap_out_list):
            query = switch_queries[sandhill].pop(swap_out_item[0])
            monitors_under_queries[query.query_id].remove(sandhill)
            for victim_sw in swap_out_item[1]:
                monitors_under_queries[query.query_id].append(victim_sw)
                switch_queries[victim_sw].append(query)

    # finally, put all aggregator on controller's ToR
    weights[0] += args.query_number

    # dump simple debug result
    for idx, weight in enumerate(weights):
        print("Switch %d weight = %d" % (idx, weight))
    mean = sum(weights) / len(weights)
    std_err = (sum((i - mean) ** 2 for i in weights) / len(weights)) ** 0.5
    print("Overall task number = %d, avg task = %.2f, std error = %.2f" % (sum(weights), mean, std_err))


from pymol import cmd, stored
import numpy as np
import matplotlib.cm as cm
from matplotlib.colors import Normalize
cmap = cm.seismic
norm = Normalize(vmin=-1, vmax=1)
cmap2 = cm.PRGn
norm = Normalize(vmin=-1, vmax=1)

def zero_residues(sel1,offset=0,chains=0):
        """
        """
        offset = int(offset)

        # variable to store the offset
        stored.first = None
        # get the names of the proteins in the selection

        names = ['(model %s and (%s))' % (p, sel1)
                        for p in cmd.get_object_list('(' + sel1 + ')')]

        if int (chains):
                names = ['(%s and chain %s)' % (p, chain)
                                for p in names
                                for chain in cmd.get_chains(p)]

        # for each name shown
        for p in names:
                # get this offset
                ok = cmd.iterate("first %s and polymer and n. CA" % p,"stored.first=resv")
                # don't waste time if we don't have to
                if not ok or stored.first == offset:
                        continue;
                # reassign the residue numbers
                cmd.alter("%s" % p, "resi=str(int(resi)-%s)" % str(int(stored.first)-offset))
                # update pymol

        cmd.rebuild()

def stucture_attribution(pdb_id, attributions_path, pdb_path, flag=False):
    '''
    '''
    data = np.load(attributions_path)
    attributions = data['data']
    kernels = data['kernels']
    labels = data['labels']
    offsets = data['offsets']
    ind = np.where(labels == pdb_id)
    #attribution = attributions[ind][0][:,-1]
    if flag:
        a_ = data["all_"][0]
        a_[a_<=0] = 0.0
        attribution = data['all_'][0] + data['all_'][1]
        attribution[attribution<=0] = 0.0
        attribution = a_ * attribution
    else:
        a_ = data["all_"][1]
        a_[a_<=0] = 0.0
        attribution = data['all_'][0] + data['all_'][1]
        attribution[attribution<=0] = 0.0
        attribution = a_ * attribution

    # Load PDB
    cmd.load(pdb_path+pdb_id[:-2]+'.pdb')
    cmd.split_chains(pdb_id[:-2])
    for name in cmd.get_names('objects', 0, '(all)'):
        if not name.endswith(pdb_id[-1].upper()) and name.startswith(pdb_id[:4]):
            cmd.delete(name)
        else: zero_residues(name)
    cmd.reset()

    cmd.color('white', pdb_id)
    for i, _ in enumerate(attribution):
        #if flag: _ = _ *-1
        cmd.select('toBecolored', pdb_id+ ' and res ' + str(i+offsets[ind][0]))
        cmd.set_color('saliency'+str(i)+pdb_id, list(cmap(norm(_)))[:3])
        #else: cmd.set_color('saliency'+str(i)+pdb_id, list(cmap2(norm(_)))[:3])
        cmd.color('saliency'+str(i)+pdb_id, 'toBecolored')

    cmd.select('selected', pdb_id)
    #cmd.show('mesh', 'selected')
    #cmd.show('sticks', 'selected')
    cmd.deselect()

def kernels(pdb_id, data_path, pdb_path, flag=False):
    '''
    '''
    data = np.load(data_path)
    kernels = data['kernels']
    labels = data['labels']
    offsets = data['offsets']
    ind = np.where(labels == pdb_id)
    kernels = kernels[ind][0][0]

    # Load PDB
    cmd.load(pdb_path+pdb_id[:-2]+'.pdb')
    cmd.split_chains(pdb_id[:-2])
    for name in cmd.get_names('objects', 0, '(all)'):
        if not name.endswith(pdb_id[-1].upper()) and name.startswith(pdb_id[:4]):
            cmd.delete(name)
        else: zero_residues(name)
    cmd.reset()

    pdb_id2 = pdb_id+'kernel'
    cmd.create(pdb_id2, pdb_id)
    cmd.delete(pdb_id)

    '''
    cmd.hide('everything',pdb_id)
    cmd.show_as('lines', pdb_id + ' and (name ca or name c or name n)')
    cmd.set('line_width', 5)
    cmd.set_bond('line_width', 5,  pdb_id + ' and (name ca or name c or name n)')
    cmd.show('spheres', pdb_id + ' and name ca')
    cmd.set('sphere_transparency', 0.0, pdb_id + ' and name ca')
    cmd.set('sphere_scale', 0.5, pdb_id + ' and name ca')
    '''

    cmd.hide('everything',pdb_id2)
    cmd.show('spheres', pdb_id2 + ' and name ca')
    cmd.set('sphere_transparency', 0.9, pdb_id2 + ' and name ca')

    for i, _ in enumerate(kernels):
        print(_)
        cmd.set('sphere_scale', _/(1.7*2), pdb_id2+ ' and res ' + str(i+offsets[ind][0])+ ' and name ca')
    cmd.deselect()

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

# Paths
data_path = '../../out/krashras_graph_new/interpret/attributions.npz'
pdb_path = '../../data/KrasHras/pdb/'

# Load Attribution
kras = ['5tb5_c', '4dsn_a', '3gft_f', '5v6v_b', '4luc_b',
             '4m22_a', '4lv6_b', '4epy_a', '4lrw_b', '4epx_a']
             #'4m1w_a', '5uqw_a', '4pzy_b', '4m21_a', '5us4_a',
             #'5f2e_a', '4q03_a', '4lyh_c', '4m1o_b', '4pzz_a']
hras = ['1p2v_a', '4urz_r', '2x1v_a', '3lo5_c', '2quz_a',
            '3kud_a', '1aa9_a', '1plk_a', '4k81_d', '5wdq_a']
            #'1iaq_a', '1xd2_a', '3i3s_r', '4efl_a', '4l9w_a',
            #'3lo5_a', '5b2z_a', '1nvv_q', '4efm_a', '3l8z_a']

cmd.reinitialize()
cmd.bg_color('black')
for _ in kras:
    #kernels(_, data_path, pdb_path)
    stucture_attribution(_, data_path, pdb_path)
for _ in hras:
    #kernels(_, data_path, pdb_path)
    stucture_attribution(_, data_path, pdb_path, True)

# Align and Translate
for _ in kras[1:]:
    cmd.align(_,  kras[0])
for _ in hras:
    cmd.align(_,  kras[0])
cmd.orient()

for i, _ in enumerate(kras):
    cmd.rotate([1.0,0.0,0.0],60,selection=_[:-1]+_[-1].upper())
    cmd.translate([0,50*0,0],selection=_[:-1]+_[-1].upper())

for i, _ in enumerate(hras):
    cmd.rotate([1.0,0.0,0.0],60,selection=_[:-1]+_[-1].upper())
    cmd.translate([-60,50*0,00],selection=_[:-1]+_[-1].upper())

cmd.reset()


#!/usr/bin/env python2.7

import argparse
import logging as log

from util import parse
from emu.cpu import CPU
from emu.instruction import Instruction
from emu.memory import Memory
from util.assemble import SPYM_MAGIC, SPYM_HDR_LEN, SPYMHeader
from util.misc import get_section

def get_args():
    parser = argparse.ArgumentParser(description='Spym MIPS Interpreter. Starts in interactive shell mode, unless given MIPS source file as argument.')
    parser.add_argument('file', metavar='FILE', type=str,
                        help='MIPS source file', nargs='?')
    parser.add_argument('--stack', type=int, help='Stack memory size. Default: 64 bytes',
                        default=64)
    parser.add_argument('--debug', help='Activate debugger. Implies verbose.',
                        action='store_true')
    parser.add_argument('-v', '--verbose', help='Verbose output',
                        action='store_true')
    return parser.parse_args()

def run_source(raw):
    source = raw.splitlines()
    dseg, tseg = parse.segments(source)
    dmem = Memory(args.stack)
    if dseg:
        dmem.memory = parse.data(dseg) + dmem.memory
    cpu = CPU(dmem, parse.text_list(tseg))
    cpu.start(args.debug)
    cpu.dump()

def run_binary(raw):
    hdr = SPYMHeader(binary=raw[:SPYM_HDR_LEN])
    text = parse.bin2text_list(get_section(raw, hdr.text_off, hdr.text_size))
    dmem = Memory(args.stack)
    dmem.memory = get_section(raw, hdr.data_off, hdr.data_size) + dmem.memory
    cpu = CPU(dmem, text)
    cpu.start(args.debug)
    cpu.dump()

def main():
    global args
    args = get_args()

    if args.file:
        lvl = log.INFO if args.verbose or args.debug else None
        log.basicConfig(format='%(message)s', level=lvl)

        with open(args.file) as f:
            raw = f.read()
            if raw[:4] == SPYM_MAGIC:
                run_binary(raw)
            else:
                run_source(raw)
    else:
        log.basicConfig(format='%(message)s', level=log.INFO)
        cpu = CPU(Memory(args.stack))
        while True:
            inp = raw_input('spym > ').strip()
            if inp == 'exit':
                break
            elif inp == 'dump':
                cpu.dump()
                continue
            elif not inp:
                continue

            try:
                cpu.execute_single(Instruction(inp))
            except Exception as e:
                if e.message == 'exit syscall':
                    break
                raise e

        cpu.dump()

if __name__ == '__main__':
    main()


import os
import pathlib

from unittest import mock

from PyQt5 import QtCore

from app import run_pyqt

class TestImageViewer:
    def test_if_main_window_renders(self, qtbot):
        window = run_pyqt.ImageViewer()
        qtbot.addWidget(window)
        assert window.currentFile.text() == 'No file selected'
        assert window.openImageButton.text() == 'Open image'

    def test_if_image_is_displayed(self, qtbot):
        path = pathlib.Path().absolute()
        file = os.path.join(path, 'tests/test.jpg')

        window = run_pyqt.ImageViewer()
        qtbot.addWidget(window)
        window.image_file_selected(file)
        assert window.currentFile.text().endswith('/test.jpg')
        assert window.currentFile.text() == file
        assert window.imageLabel.pixmap().width() == 242
        assert window.imageLabel.pixmap().height() == 221

    @mock.patch('app.run_pyqt.ImageViewer.open_image_action')
    def test_if_action_is_called_on_button_click(self, open_image_action, qtbot):
        window = run_pyqt.ImageViewer()
        qtbot.addWidget(window)

        qtbot.mouseClick(window.openImageButton, QtCore.Qt.LeftButton)
        assert open_image_action.called

    @mock.patch('app.run_pyqt.ImageViewer.show_image_picker')
    def test_if_show_image_picker_is_called_on_prompt_accept(self, show_image_picker, qtbot):
        window = run_pyqt.ImageViewer()
        qtbot.addWidget(window)
        prompt = window.open_image_action()
        prompt.accept()
        assert show_image_picker.called

    @mock.patch('app.run_pyqt.ImageViewer.show_image_picker')
    def test_if_show_image_picker_is_not_called_when_prompt_not_accept(self, show_image_picker, qtbot):
        window = run_pyqt.ImageViewer()
        qtbot.addWidget(window)
        window.open_image_action()
        assert not show_image_picker.called

    @mock.patch('app.run_pyqt.ImageViewer.image_file_selected')
    def test_if_image_file_selected_is_called_on_file_selection(self, image_file_selected, qtbot):
        path = pathlib.Path().absolute()
        file = os.path.join(path, 'tests/test.jpg')

        def accept():
            picker.accept()

        window = run_pyqt.ImageViewer()
        qtbot.addWidget(window)
        picker = window.show_image_picker()
        picker.selectFile(file)

        picker.done(1)
        QtCore.QTimer.singleShot(100, accept)
        picker.exec_()
        assert image_file_selected.called


from __future__ import annotations
from collections.abc import Iterable

import math
import pandas as pd

from pyetm.logger import get_modulelogger
from .session import SessionMethods

logger = get_modulelogger(__name__)

class AccountMethods(SessionMethods):
    """Account object methods"""

    @property
    def my_scenarios(self) -> pd.DataFrame:
        """ all scenarios connected to account"""

        # validate token permission
        self._validate_token_permission("scenarios:read")

        # set url
        url = 'scenarios'

        # determine number of pages
        pages = self._get_objects(url=url, page=1, limit=1)
        pages = math.ceil(pages['meta']['total'] / 100)

        if pages == 0:
            return pd.DataFrame()

        # newlist
        scenarios = []
        for page in range(1, pages + 1):

            # fetch pages and format scenarios
            recs = self._get_objects(url, page=page, limit=100)['data']

            excl = ['user_values', 'balanced_values', 'metadata', 'url']
            scenarios.extend([
                self._format_object(scen, excl) for scen in recs])

        return pd.DataFrame.from_records(scenarios, index='id')

    @property
    def my_saved_scenarios(self) -> pd.DataFrame:
        """all saved scenarios connector to account"""

        # validate token permission
        self._validate_token_permission("scenarios:read")

        # set url
        url = 'saved_scenarios'

        # determine number of pages
        pages = self._get_objects(url, page=1, limit=1)
        pages = math.ceil(pages['meta']['total'] / 100)

        if pages == 0:
            return pd.DataFrame()

        # newlist
        scenarios = []
        for page in range(1, pages + 1):

            # fetch pages and format scenarios
            recs = self._get_objects(url, page=page, limit=100)['data']

            excl = ['scenario', 'scenario_id', 'scenario_id_history']
            scenarios.extend([
                self._format_object(scen, excl) for scen in recs])

        return pd.DataFrame.from_records(scenarios, index='id')

    @property
    def my_transition_paths(self) -> pd.DataFrame:
        """all transition paths connected to account"""

        # validate token permission
        self._validate_token_permission("scenarios:read")

        # set url
        url = 'transition_paths'

        # determine number of pages
        pages = self._get_objects(url, page=1, limit=1)
        pages = math.ceil(pages['meta']['total'] / 100)

        if pages == 0:
            return pd.DataFrame()

        # newlist
        paths = []
        for page in range(1, pages + 1):

            # fetch pages and format scenarios
            recs = self._get_objects(url, page=page, limit=100)['data']

            paths.extend([
                self._format_object(path) for path in recs])

        return pd.DataFrame.from_records(paths, index='id')

    def _format_object(self, obj: dict, exclude: Iterable | None = None):
        """helper function to reformat a object."""

        # default list
        if exclude is None:
            exclude = []

        # add string to list
        if isinstance(exclude, str):
            exclude = [exclude]

        # flatten passed keys
        for key in ['owner']:
            if key in obj:

                # flatten items in dict
                item = obj.pop(key)
                item = {f'{key}_{k}': v for k, v in item.items()}

                # add back to scenario
                obj = {**obj, **item}
        # process datetimes
        for key in ['created_at', 'updated_at']:
            if key in obj:
                if obj.get(key) is not None:
                    obj[key] = pd.to_datetime(obj[key], utc=True)

        for key in ['template']:
            if key in obj:
                if obj.get(key) is None:
                    obj[key] = pd.NA

        # reduce items in scenario
        return {k:v for k,v in obj.items() if k not in exclude}

    def _get_objects(self, url: str, page: int = 1, limit: int = 25) -> dict:
        """Get info about object in url that are connected
         to the user token. Object can be scenarios, saved scenarios
         or transition paths."""

        # raise without required permission
        self._validate_token_permission(scope='scenarios:read')

        # format request
        params = {'page': int(page), 'limit': int(limit)}
        headers = {'content-type': 'application/json'}

        # request response
        resp = self.session.get(
            url, params=params, decoder='json', headers=headers)

        return resp


import scipy.io
from tkinter import *
from tkinter import messagebox, filedialog
from tkinter import ttk
import h5py as h5

Tk().withdraw()
_file = filedialog.askopenfilename(title='Choose file to load')

with h5.File(_file, "r") as hf:
    for ind, key in enumerate(hf.keys()):
        if ind > 0:
            print(ind,":", key)
            for ind2, key2 in enumerate(hf[key].keys()):
                dataset = hf[key][key2]
                print("\t",ind2,":", key2, "->", type(dataset), "({:})".format(len(dataset)))
                for data in dataset:
                    print("\t\t", type(data), "({:})".format(len(data)))
                    break


"""
Base classes.

"""

from itertools import islice

class IterContainer(object):

    def __contains__(self, item):
        for o in self:
            if o == item:
                return True
        return False

    def __len__(self):
        return sum(1 for _ in self)

    def __getitem__(self, item):
        if isinstance(item, int):
            try:
                return next(islice(self, item, item+1))
            except StopIteration:
                raise IndexError('index out of range')
        elif isinstance(item, slice):
            return islice(self, item.start, item.stop, item.step)

    def index(self, item):
        for i, o in enumerate(self):
            if o == item:
                return i
        raise ValueError('%s is not in container' % item)


"""Command-line interface."""
import argparse

from pynwsradar import Nexrad

def create_radar(station, output, layer, num) -> None:
    """Create radar image."""
    radar = Nexrad(station=station)
    radar.update()
    assert radar.layers
    layer = radar.layers[f"{station}_{layer}"]
    layer.update_image(num=num)
    layer.save_image(output)

def main():
    """Make radar imgae with commandline."""
    parser = argparse.ArgumentParser(description="Create NEXRAD qradar image.")
    parser.add_argument("station")
    parser.add_argument("file", help="output file name")
    parser.add_argument("--layer", default="bref_raw")
    parser.add_argument("--num", type=int, default=1, help="number of frames")
    args = parser.parse_args()
    create_radar(args.station, args.file, args.layer, args.num)


from __future__ import print_function

import os
import six
import pyrax

pyrax.set_setting("identity_type", "rackspace")
creds_file = os.path.expanduser("~/.rackspace_cloud_credentials")
pyrax.set_credential_file(creds_file)
imgs = pyrax.images
imgs.http_log_debug = True

print("Listing images with pending status...")
images = imgs.list(visibility="shared", member_status="pending")

if not images:
    print("No pending images found")
    exit()

for pos, image in enumerate(images):
    new_status = None
    print("[%s] - %s" % (pos, image.name))
    choice = six.moves.input("Would you like to accept, reject or skip? "
            "('a', 'r', or 's'): ")
    if choice == 'a':
        new_status = 'accepted'
    elif choice == 'r':
        new_status = 'rejected'

    if new_status is not None:
        print("[%s] - %s : Updating status to %s" % (pos, image.name, new_status))
        imgs.update_image_member(image.id, new_status)
        print("[%s] - %s : has been updated" % (pos, image.name))
    else:
        print("[%s] - %s : Skipping update" % (pos, image.name))


from argparse        import ArgumentParser, REMAINDER
from multiprocessing import cpu_count
from oswalk          import oswalk_flattens
from pathlib         import Path
from sqlite3         import connect

from pwformat import PWFormat
from oswalk   import oswalk_flatten

if __name__ == "__main__":
	parser = ArgumentParser(
                description='Merge SQLite3 databases into fewer databases'
	)
	#parser.add_argument("-v", "--verbosity", action="count", default=0)
	parser.add_argument('output', help="SQLite3 database output")
	parser.add_argument('input',  help="SQLite3 database input directory")
	args = parser.parse_args()

	input_dir = Path(args.input)
	output    = Path(args.output)
	if output.exists(): output.unlink()
	output.touch()

	with connect(str(output)) as outconn:
		outcur = outconn.cursor()
		outcur.execute(PWFormat().create)

		for infile in oswalk_flatten(input_dir):
			outcur.execute('''ATTACH ? as indb''', (str(infile), ))
			outcur.execute('''BEGIN''')
			outcur.execute('''INSERT INTO data SELECT * FROM indb.data''')
			outconn.commit()
			outcur.execute('''DETACH database indb''')

		outconn.commit()


#Translation means shifting an image in x and y axis or also the
#combinations of both x and y axis

import cv2 as cv
import numpy as np
img = cv.imread("C:/Users/vijay/Desktop/photo.jpg")

#We need to create a transalation method with img,x,y as parameters
def translate(img,x,y):
    #We need to create a translation matrix
    transMat = np.float32([[1,0,x],[0,1,y]])
    #We also need dimenions of the image
    dimensions = (img.shape[1], img.shape[0])
    return cv.warpAffine(img, transMat, dimensions)

#Values are
# -x --> left
# -y --> up
# x --> right
# y --> down

translated = translate(img,100,150)
cv.imshow("Image",translated)
cv.waitKey(0)
cv.destroyAllWindows()