import sys

from mainwindow import Ui_MainWindow
from dataloader import DataLoader
from xbee import XBeeCommunicator

from functools import partial

from PyQt5 import QtCore, QtGui, QtWidgets
import matplotlib.pyplot as plt
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.backends.backend_qt5agg import NavigationToolbar2QT as NavigationToolbar
from matplotlib.figure import Figure

from mpl_toolkits.mplot3d import Axes3D

from random import randint
from datetime import datetime

import warnings

HEADERS = ["TeamID", "Time", "Packet", "Altitude", "Pressure", "Airspeed", "Temperature", "Voltage", "Latitude", "Longitude",
"GPSAlt", "Satellites", "GPSSpeed", "Heading", "ImageCount", "State"]

# Values below graph
v1 = "Time"
v2 = "Altitude"
v3 = "GPSSpeed"
v4 = "Airspeed"
v5 = "Temperature"

# Values below graph - units
v1_u = "s"
v2_u = "m"
v3_u = "mph"
v4_u = "mph"
v5_u = "°F"

class Wrapper():

    def __init__(self, ui, xbee_communicator, dataloader=None):
        self.ui = ui
        self.ax = None
        self.xbee_communicator = xbee_communicator

        # reference to 3d plot
        self.plot3d = None

        # hold the smallest min and largest max limits for the y axis
        # this lets us know the highest or lowest we can scale the graph
        self.yLimits = [0, 0]
        self.xLimits = [0, 0]

        # hold the current min and max limits for the y axis
        # this keeps track of the values the sliders are at
        self.minY = 0
        self.maxY = 0
        self.minX = 0
        self.maxX = 0

        # populate plot select menu
        self.ui.comboBox.clear()
        # self.ui.comboBox.addItem("All")
        # should create separate array of items we want to plot
        self.ui.comboBox.addItems(HEADERS[2:])

        self.ui.horizontalSlider_2.setValue(100)
        self.ui.horizontalSlider_4.setValue(100)

        # store the current option select in plot select
        self.currentPlot = self.ui.comboBox.currentText()

        if dataloader:
            # Create default dataloader? Dataloader should be initialized with directory/file name
            self.dataloader = dataloader
            pass

    def codeUpdatesToUI(self):
        self.ui.figure = Figure()
        self.ui.canvas = FigureCanvas(self.ui.figure)
        self.ui.canvas.setObjectName("gridCanvas")
        # add the plot to the main layout
        self.ui.gridLayout.addWidget(self.ui.canvas, 0, 0, 1, 1)

        # add a plot to the sub layout
        self.ui.figure2 = Figure()
        self.ui.canvas2 = FigureCanvas(self.ui.figure2)
        self.ui.canvas2.setObjectName("gridCanvas2")
        self.ui.gridLayout_2.addWidget(self.ui.canvas2, 0, 0, 1, 1)

    def initNewUI(self):

        # Change session name
        if self.dataloader is not None:

            # Update session name
            self.update_session_name(self.dataloader.file_name)

            if self.ax is not None:
                # Clear old plot to plot new data
                self.plotClear()

            # Plot initial data
            data = self.dataloader.fetch(["Time", self.currentPlot])

            # Compute new limits
            self.compute_plot_limits(data)

            # Update limits
            self.update_plot_limits()

            # Update plot controls
            self.update_plot_controls(True)

            # Plot
            self.plot(data, "Time", self.currentPlot)

            # plot 3d
            self.plotOrientation(0, 0, 0, 0.2, -0.4, 1.8)

            # Update text vals
            self.update_text_vals()

        else:
            # Update session name
            self.update_session_name("No File Loaded")

    def compute_plot_limits(self, data):

        data = data.astype(float)

        if not data.empty:
            self.maxX = int(data["Time"].max())
            self.maxY = int(data[self.currentPlot].max()) + 1
            self.minX = int(data["Time"].min())
            self.minY = int(data[self.currentPlot].min()) + 1
        else:
            self.maxX = 10
            self.maxY = 10
            self.minX = 0
            self.minY = 0

    def update_plot_limits(self):
        axes = self.ui.figure.gca()
        axes.set_xlim([self.minX, self.maxX])
        axes.set_ylim([self.minY, self.maxY])
        self.ui.canvas.draw()

    def update_plot_controls(self, set_slider_max):

        minVals = [self.minX, self.minY]
        maxVals = [self.maxX, self.maxY]
        mintxtedits = [self.ui.textEdit, self.ui.textEdit_3]
        maxtxtedits = [self.ui.textEdit_2, self.ui.textEdit_4]
        minsliders = [self.ui.horizontalSlider_3, self.ui.horizontalSlider]
        maxsliders = [self.ui.horizontalSlider_4, self.ui.horizontalSlider_2]

        for minval, minte, minslider, maxval, maxte, maxslider in zip(minVals, mintxtedits, minsliders, maxVals, maxtxtedits, maxsliders):
            minte.setText(str(minval))
            maxte.setText(str(maxval))

            if set_slider_max:
                minslider.setMinimum(minval)
                minslider.setMaximum(maxval)
                maxslider.setMinimum(minval)
                maxslider.setMaximum(maxval)

            minslider.setValue(minval)
            maxslider.setValue(maxval)

    # plot the data
    """
    need to use data.astype(float) so it can process the data
    need ax = self.ui.figure.add_subplot(111) to show the graph
    set x axis to Time and y axis to altitude
    """
    def plot(self, data, x, y):
        self.ax = self.ui.figure.add_subplot(111)
        self.ax.plot(data[x].astype(float), data[y].astype(float),
            color = "xkcd:teal")
        self.ax.set_title("{} vs {}".format(x, y))
        self.ax.set_xlabel(x)
        self.ax.set_ylabel(y)
        self.ax.grid()
        self.ui.canvas.draw()

    # plots new points
    def plot_points(self, data, x, y):
        self.ax.plot(data[x].tail(2).astype(float),
            data[y].tail(2).astype(float), color = "xkcd:teal")
        self.ui.canvas.draw()

    # Clear the plot
    def plotClear(self):
        self.ax.clear()
        self.ui.canvas.draw()

    # plots 3d
    # use u v w for orientation/direction of arrow
    # x y z is for initial position of tail
    def plotOrientation(self, x, y, z, u, v, w):
        ax = self.ui.figure2.add_subplot(111, projection = "3d")
        ax.set_xlabel("x")
        ax.set_ylabel("y")
        ax.set_zlabel("z")
        ax.quiver(x, y, z, u, v, w)
        ax.set_xlim([-1, 1])
        ax.set_ylim([-1, 1.5])
        ax.set_zlim([-1, 2])
        self.ui.canvas2.draw()

    # Connects buttons to given functions
    def setUpHandlers(self):
        self.ui.actionConnect.triggered.connect(self.select_port)
        self.ui.actionStart.triggered.connect(self.xbee_start)
        self.ui.actionPause.triggered.connect(self.xbee_pause)
        self.ui.actionStop.triggered.connect(self.xbee_stop)

        # detect when a new plot has been selected
        self.ui.comboBox.activated.connect(self.setComboBox)

        # detect when max sliders have changed
        self.ui.horizontalSlider_2.valueChanged.connect(partial(self.maxSliderChange, self.ui.horizontalSlider_2)) # Y Max
        self.ui.horizontalSlider_4.valueChanged.connect(partial(self.maxSliderChange, self.ui.horizontalSlider_4)) # X Max

        # detect when min sliders have changed
        self.ui.horizontalSlider.valueChanged.connect(partial(self.minSliderChange, self.ui.horizontalSlider)) # Y Min
        self.ui.horizontalSlider_3.valueChanged.connect(partial(self.minSliderChange, self.ui.horizontalSlider_3)) # X Min

        self.ui.actionNew_Session.triggered.connect(self.new_session)
        self.ui.actionLoad_Session.triggered.connect(self.load_session)

        # Snapshot detection
        self.ui.pushButton.clicked.connect(self.snapshot)

    # Starts a new session (new data file)
    def new_session(self):
        if self.yesno_prompt("New Session", "Are you sure you want to start a new session?"):

            # Prompt user to choose filename
            fn = self.saveFileNameDialog()
            if fn == '':
                return

            # Save current data
            if self.dataloader is not None:
                self.dataloader.save_as_csv()

            # Use current time as filename
            self.dataloader = DataLoader(fn, HEADERS)
            self.dataloader.save_as_csv()

            self.initNewUI()

    # Loads a session from an old csv file
    def load_session(self):
        if self.yesno_prompt("Load Session", "Are you sure you want to load a session?"):

            # Prompt user to choose filename
            fn = self.openFileNameDialog()
            if fn == '':
                return

            # Save current data loader
            if self.dataloader is not None:
                self.dataloader.save_as_csv()

            # Load data from file
            self.dataloader = DataLoader(fn, HEADERS)
            self.dataloader.read_file()

            self.initNewUI()

    def xbee_update(self, data_row):
        if self.dataloader is None:
            print("Error: Session has not been started")
            return

        self.dataloader.update(HEADERS, data_row)

        x = "Time"
        y = self.currentPlot

        # Fetch data
        data = self.dataloader.fetch([x, y])

        # Compute new limits
        self.compute_plot_limits(data)

        # Update limits
        self.update_plot_limits()

        # Update plot controls
        self.update_plot_controls(True)

        print("here")

        # Update plot
        self.plot_points(data, x, y)

        # Update text below graph
        self.update_text_vals()

    def select_port(self):
        valid = False
        while not valid:
            port, choice = self.inputdialog("Port", "Input Port (/dev/tty or COM)")
            if not choice:
                print("None selected")
                return
            else:
                print("Chosen port: {}".format(port))
                valid = xbee_communicator.connect(port)
                if not valid:
                    self.warningdialog("Not a valid port, try again.")

    def xbee_start(self):
        valid = xbee_communicator.start(self.xbee_update)
        if not valid:
            self.warningdialog("Not connected; cannot start.")

    def xbee_pause(self):
        valid = xbee_communicator.pause()
        if not valid:
            self.warningdialog("Not connected; cannot pause.")

    def xbee_stop(self):
        valid = xbee_communicator.stop()
        if not valid:
            self.warningdialog("Not connected; cannot stop.")

    def inputdialog(self, title, message):
        return QtWidgets.QInputDialog.getText(self.ui.mainwindow, title, message)

    # update the current plot and replace the previous plot
    def setComboBox(self, index):
        self.currentPlot = self.ui.comboBox.itemText(index)

        if (self.currentPlot == "All"):
            pass
        else:
            data = self.dataloader.fetch(["Time", self.currentPlot])
            warnings.filterwarnings("ignore",module="matplotlib")
            self.plotClear()
            self.plot(data, "Time", self.currentPlot)

            # Compute new limits
            self.compute_plot_limits(data)

            # Update limits
            self.update_plot_limits()

            # Update plot controls
            self.update_plot_controls(True)

    # change the max range of the y axis using slider
    def maxSliderChange(self, slider):
        newMaxValue = slider.value()
        axes = self.ui.figure.gca()

        # slider for Y Max
        if (slider == self.ui.horizontalSlider_2):
            if self.minY < newMaxValue:
                self.maxY = newMaxValue
                self.update_plot_limits()
        # slider for X Max
        else:
            if self.minX < newMaxValue:
                self.maxX = newMaxValue
                self.update_plot_limits()
        warnings.filterwarnings("ignore",module="matplotlib")

        self.update_plot_controls(False)

    # change the min range of the y axis using slider
    def minSliderChange(self, slider):
        newMinValue = slider.value()
        axes = self.ui.figure.gca()

        # Slider for Y Min
        if (slider == self.ui.horizontalSlider):
            if newMinValue < self.maxY:
                self.minY = newMinValue
                self.update_plot_limits()
        # Slider for X Min
        else:
            if newMinValue < self.maxX:
                self.minX = newMinValue
                self.update_plot_limits()

        warnings.filterwarnings("ignore",module="matplotlib")

        self.update_plot_controls(False)

    def getCurrentPlot(self):
        return self.currentPlot

    def snapshot(self):
        if not self.xbee_communicator.snapshot():
            self.warningdialog("Not connected to XBee")

    def warningdialog(self, message):
        msg = QtWidgets.QMessageBox()
        msg.setIcon(QtWidgets.QMessageBox.Warning)
        msg.setText(message)
        msg.setStandardButtons(QtWidgets.QMessageBox.Ok)
        msg.setWindowTitle("WARNING!")
        ret = msg.exec_()

    def yesno_prompt(self, title, msg):
        reply = QtWidgets.QMessageBox.question(self.ui.mainwindow, title,
                 msg, QtWidgets.QMessageBox.Yes, QtWidgets.QMessageBox.No)
        return True if reply == QtWidgets.QMessageBox.Yes else False

    def openFileNameDialog(self):
        options = QtWidgets.QFileDialog.Options()
        options |= QtWidgets.QFileDialog.DontUseNativeDialog
        fileName, _ = QtWidgets.QFileDialog.getOpenFileName(self.ui.mainwindow,"QFileDialog.getOpenFileName()",
            "","All Files (*);;Python Files (*.py)", options=options)
        return fileName

    def saveFileNameDialog(self):
        options = QtWidgets.QFileDialog.Options()
        options |= QtWidgets.QFileDialog.DontUseNativeDialog
        fileName, _ = QtWidgets.QFileDialog.getSaveFileName(self.ui.mainwindow,"QFileDialog.getSaveFileName()",
            "","All Files (*);;Python Files (*.py)", options=options)
        return fileName

    def update_session_name(self, name):
        self.ui.session_label.setText(name)

    def update_text_vals(self):
        label_order = [self.ui.label_11, self.ui.label_9, self.ui.label_6, self.ui.label_5, self.ui.label_4]
        for field, unit, label in zip([v1, v2, v3, v4, v5], [v1_u, v2_u, v3_u, v4_u, v5_u], label_order):
            txt_data = self.dataloader.fetch([field])
            text = str(txt_data[field].iloc[-1]) + unit if len(txt_data[field]) > 1 else "N/A"
            label.setText(text)

# Instantiate UI
if __name__ == "__main__":
    app = QtWidgets.QApplication(sys.argv)
    MainWindow = QtWidgets.QMainWindow()
    ui = Ui_MainWindow()
    ui.setupUi(MainWindow)
    ui.mainwindow = MainWindow

    # Setup necessary components
    # Optional
    dataloader = DataLoader("./data/Data.txt", HEADERS)
    dataloader.read_file()

    xbee_communicator = XBeeCommunicator()

    # Create wrapper around UI Window
    window = Wrapper(ui, xbee_communicator, dataloader)
    window.codeUpdatesToUI()
    window.setUpHandlers()
    window.initNewUI()

    # Show window
    MainWindow.show()

    sys.exit(app.exec_())


import numpy as np
from tiles3 import tiles, IHT
from utils import argmax
from agent import BaseAgent

class MountainCarTileCoder:
    def __init__(self, iht_size, num_tilings, num_tiles):
        '''
        Initializes the MountainCar Tile Coder
        iht_size -- int, the size of the index hash table, typically a power of 2
        num_tilings -- int, the number of tillings
        num_tiles -- int, the number of tiles. Both width and height of the tile coder are the same.

        Class Variables:
        self.iht -- IHT, the index hash table that the tile coder will use
        self.num_tilings -- int, the number of tilings the tile coder will use
        self.num_tiles -- int, the number of tiles the tile coder will use
        '''

        self.iht = IHT(iht_size)
        self.num_tilings = num_tilings
        self.num_tiles = num_tiles

    def get_tiles(self, position, velocity):
        '''
        Takes a position and velocity from the MountainCar Environment and returns a numpy array of active tiles.

        Arguments:
        position -- float, the position of the agent between -1.2 and 0.5
        velocity -- float, the velocity of the agent between -0.007 and 0.007

        returns -- np.array, active tiles
        '''
        # Set the max and min of position and velocity to scale the input
        min_position = -1.2
        max_position = 0.5
        min_velocity = -0.07
        max_velocity = 0.07

        # Scale position and velocity
        scale_position = self.num_tiles / (max_position - min_position)
        scale_velocity = self.num_tiles / (max_velocity - min_velocity)

        active_tiles = tiles(self.iht, self.num_tilings, [scale_position * position, scale_velocity * velocity])
        return np.array(active_tiles)

class ExpectedSarsaAgent(BaseAgent):

    def __init__(self):
        # Initialization of Sarsa Agent. All values are set to None so they can be initialized in the agent_init method.
        self.last_action = None
        self.last_state = None
        self.epsilon = None
        self.gamma = None
        self.iht_size = None
        self.w = None
        self.alpha = None
        self.num_tilings = None
        self.num_tiles = None
        self.mctc = None
        self.initial_weights = None
        self.num_actions = None
        self.previous_tiles = None

    def agent_init(self, agent_info={}):
        # Setup for the agent called when the experiment first starts.
        self.num_tilings = agent_info.get("num_tilings")
        self.num_tiles = agent_info.get("num_tiles")
        self.iht_size = agent_info.get("iht_size")
        self.epsilon = agent_info.get("epsilon")
        self.gamma = agent_info.get("gamma")
        self.alpha = agent_info.get("alpha")
        self.initial_weights = agent_info.get("initial_weights")
        self.num_actions = agent_info.get("num_actions")

        # Initialize self.w to num_actions times the iht_size. Because we need to have one set of weights for each action.
        self.w = np.ones((self.num_actions, self.iht_size)) * self.initial_weights

        # Initialize self.mctc to MountainCarTileCoder
        self.mctc = MountainCarTileCoder(self.iht_size, self.num_tilings, self.num_tiles)

    def select_action(self, tiles):
        # Selects an action using epsilon greedy
        # set action_values to a empty list
        # chosen_action to None
        action_values = []
        chosen_action = None

        # Loop through the weights to get action value for each action
        for action in range(self.num_actions):
            value = 0
            for active_tile in tiles:
                value += self.w[action][active_tile]
            # append action value of each action into action_values list
            action_values.append(value)

        # Use epsilon greedy to select an action
        if np.random.random() < self.epsilon:
            chosen_action = np.random.choice(self.num_actions)
        else:
            chosen_action = argmax(action_values)
        # return chosen_action and the action value of the chosen action
        return chosen_action, action_values[chosen_action], action_values

    def agent_start(self, state):
        # set position and velocity from state
        position, velocity = state

        # use self.mctc to set active_tiles using position and velocity
        active_tiles = self.mctc.get_tiles(position, velocity)

        # set current_action, action_value and action_values(list of action value if all actions) using select_action function with active_tiles
        current_action, action_value, action_values = self.select_action(active_tiles)

        # set self.last_action_value and self.last_action_values
        self.last_action_value = action_value
        self.last_action_values = action_values

        # set self.last_action and self.previous_tiles
        self.last_action = current_action
        self.previous_tiles = np.copy(active_tiles)
        # return self.last_action
        return self.last_action

    def agent_step(self, reward, state):
        # set position and velocity from state
        position, velocity = state

        # use self.mctc to set active_tiles using position and velocity
        active_tiles = self.mctc.get_tiles(position, velocity)
        # set current_action, action_value and action_values(list of action value if all actions) using select_action function with active_tiles
        current_action, action_value, action_values = self.select_action(active_tiles)
        # initalized feature_vector to all zeors
        feature_vector = np.zeros(self.iht_size)
        # loop through self.previous_tiles
        for i in self.previous_tiles:
            # set the value of index = to active tile in feature vecotor to 1 else 0.
            feature_vector[i] = 1
        # initalized target to 0
        target = 0
        # loop through all the actions
        # using epsilon greedy policy to get policy
        for action in range(self.num_actions):
            if (action == np.argmax(action_values)):
                policy = 1 - self.epsilon + (self.epsilon / self.num_actions)
            else:
                policy = self.epsilon / self.num_actions
            # sum the target of all actions
            # target = policy of next state, next action times value of next state next action and weights
            target += policy * action_values[action]
        # updated self.w of currrent action
        self.w[self.last_action] = self.w[self.last_action] + self.alpha * (reward + self.gamma * target - self.last_action_value) * feature_vector
        # set self.last_action
        self.last_action_value = action_value

        # set self.last_action and self.previous_tiles
        self.last_action = current_action
        self.previous_tiles = np.copy(active_tiles)
        return self.last_action

    def agent_end(self, reward):
        # initalized feature_vector to all zeors
        feature_vector = np.zeros(self.iht_size)
        for i in self.previous_tiles:
            feature_vector[i] = 1
        # updated self.w (No next state and next action)
        self.w[self.last_action] = self.w[self.last_action] + self.alpha * (reward - self.last_action_value) * feature_vector

    def agent_cleanup(self):
        pass

    def agent_message(self, message):
        pass


from django import forms

from .models import Comment

class CommentForm(forms.ModelForm):
    author = forms.CharField(label='Ваше имя:', widget=forms.TextInput(attrs={
        'class': 'form-control',
        'placeholder': 'Введите имя',
    }))
    text = forms.CharField(label='Комментарий:', widget=forms.Textarea(attrs={
        'class': 'form-control',
        'placeholder': 'Введите комментарий',
    }))

    class Meta:
        model = Comment
        fields = ('author', 'text')


# -*- coding: utf-8 -*-
"""
Created on Fri Oct 29 10:51:03 2021

@author: Dina
"""

# imports

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from pyrolite.plot import pyroplot
import plotly.express as px
import ternary
from sklearn.metrics import mean_squared_error, r2_score
from sklearn.preprocessing import PolynomialFeatures
from sklearn.feature_extraction import DictVectorizer
# some values are below detection limit; these will be replaced with half of the detection limit
# there are occasional Xs instead of a number; these will be replaced by zeros

def item_cleanup(s):
    if type(s) != str:
        c = pd.to_numeric(s)
    elif s == "X":
        c = 0
    elif "<" in s:
        c = pd.to_numeric(s[1:]) / 2
    else:
        c = pd.to_numeric(s)
    return c

def data_cleanup(s):
    c = s.copy()
    for i in range(s.size):
        c.values[i] = item_cleanup(s.values[i])
    return c

# read the data
df = pd.read_csv('greenstone_geochemistry_rock_selected_groups.csv',encoding='cp1252')
# 1. EDA
# We will look at the Cu, Au, Ni (key ore elements), As and SO3 (maybe pathfinder elements, or elements that are
# contained in the ore minerals, but are not useful by themselves):
els = ['SO3_pct',
'Au_ppb',
'Ni_ppm',
'Cu_ppm',
'As_ppm']

ore_els = df[els].copy()
ore_els = ore_els.dropna()

ore_els = ore_els.apply(data_cleanup) # convert to numeric values

sns.histplot(ore_els['SO3_pct'], bins = 50) # bins is how many bars we have
plt.show()
sns.histplot(ore_els['Au_ppb'], bins = 50) # bins is how many bars we have
plt.show()
sns.histplot(ore_els['Ni_ppm'], bins = 50) # bins is how many bars we have
plt.show()
sns.histplot(ore_els['Cu_ppm'], bins = 50) # bins is how many bars we have
plt.show()
sns.histplot(ore_els['As_ppm'], bins = 50) # bins is how many bars we have
plt.show()

# all long tails except Au (very few significant values) and Ni (sort of bimodal)
# will need to apply log-transform

# 1.2 Basic plots for looking at rock composition (from Scott Halley's tutorial).

# plot the density plots of Si vs Mg, Fe, Cr, Al, Sc, Ti (page 4 of the tutorial)

silica_els = ['SiO2_pct',
                'TiO2_pct',
                'Al2O3_pct',
                'Fe2O3T_pct',
                'MgO_pct',
                'Sc_ppm',
                'Cr_ppm',
                'MnO_pct',
                'CaO_pct',
                'Na2O_pct',
                'K2O_pct',
                'P2O5_pct',
                'SO3_pct']

silica_et_al = df[silica_els].copy()
silica_et_al = silica_et_al.dropna()
silica_et_al = silica_et_al.apply(data_cleanup)

# multiple detection limits! leave for the moment, but for the future this has to be half of the lowest detection limit?

# Convert oxides to els by dividing by the molar weight of oxide and multiplyting by molar weight of element:
# Mg = MgO/40.3044*24.305
# Fe = Fe2O3/159.69*55.845
# Al = Al2O3/101.96*26.9815
# Ti = TiO2/79.866*47.867
silica_et_al['Mg'] = silica_et_al['MgO_pct']/40.3044*24.305
silica_et_al['Fe'] = silica_et_al['Fe2O3T_pct']/159.69*55.845
silica_et_al['Al'] = silica_et_al['Al2O3_pct']/101.96*26.9815
silica_et_al['Ti'] = silica_et_al['TiO2_pct']/79.866*47.867

# plot with pyrolite library
fig, ax = plt.subplots(2, 3, figsize=(14, 8))       # create a composite plot with 6 figures, 3 in 2 rows
ax = ax.flat                                       # this iterates through the figures...

Mg_vs_Si = ['SiO2_pct', 'Mg']
Fe_vs_Si = ['SiO2_pct', 'Fe']
Cr_vs_Si = ['SiO2_pct', 'Cr_ppm']
Al_vs_Si = ['SiO2_pct', 'Al']
Sc_vs_Si = ['SiO2_pct', 'Sc_ppm']
Ti_vs_Si = ['SiO2_pct', 'Ti']

silica_et_al.loc[:, Mg_vs_Si].pyroplot.density(ax=ax[0])
silica_et_al.loc[:, Fe_vs_Si].pyroplot.density(ax=ax[1])
silica_et_al.loc[:, Cr_vs_Si].pyroplot.density(ax=ax[2])
silica_et_al.loc[:, Al_vs_Si].pyroplot.density(ax=ax[3])
silica_et_al.loc[:, Sc_vs_Si].pyroplot.density(ax=ax[4])
silica_et_al.loc[:, Ti_vs_Si].pyroplot.density(ax=ax[5])

titles = ["SiO2,% vs Mg, %", "SiO2,% vs Fe, %", "SiO2,% vs Cr, ppm", "SiO2,% vs Al, %", "SiO2,% vs Sc, ppm", "SiO2,% vs Ti, %"]
for t, a in zip(titles + [i + " (log-log)" for i in titles], ax):    # put headers on the diagrams - leaving this pyroplot cycle for future log-log plots
    a.set_title(t)
plt.tight_layout()
plt.show()

fig, ax = plt.subplots(2, 3, figsize=(14, 8))       # create a composite plot with 6 figures, 3 in 2 rows
ax = ax.flat                                       # this iterates through the figures...
contours = [0.95, 0.7, 0.5, 0.3]  # set the density contours: cumulative 95% ... 30% of points

silica_et_al.loc[:, Mg_vs_Si].pyroplot.density(ax=ax[0], contours=contours)
silica_et_al.loc[:, Fe_vs_Si].pyroplot.density(ax=ax[1], contours=contours)
silica_et_al.loc[:, Cr_vs_Si].pyroplot.density(ax=ax[2], contours=contours)
silica_et_al.loc[:, Al_vs_Si].pyroplot.density(ax=ax[3], contours=contours)
silica_et_al.loc[:, Sc_vs_Si].pyroplot.density(ax=ax[4], contours=contours)
silica_et_al.loc[:, Ti_vs_Si].pyroplot.density(ax=ax[5], contours=contours)

for t, a in zip(titles + [i + " (log-log)" for i in titles], ax):    # put headers on the diagrams - leaving this pyroplot cycle for future log-log plots
    a.set_title(t)
plt.tight_layout()
plt.show()

# 1.2 Jensen plot
# calculate cation percentage for Fe+Ti, Al and Mg, normalize to sum of those three to plot on a ternary diagram
silica_et_al['cat_Si'] = silica_et_al['SiO2_pct']/60.09*1
silica_et_al['cat_Ti'] = silica_et_al['TiO2_pct']/79.9*1
silica_et_al['cat_Al'] = silica_et_al['Al2O3_pct']/101.96*2
silica_et_al['cat_Fe'] = silica_et_al['Fe2O3T_pct']/159.69*2
silica_et_al['cat_Mg'] = silica_et_al['MgO_pct']/40.3*1
silica_et_al['cat_Mn'] = silica_et_al['MnO_pct']/70.94*1
silica_et_al['cat_Ca'] = silica_et_al['CaO_pct']/56.08*1
silica_et_al['cat_Na'] = silica_et_al['Na2O_pct']/61.98*2
silica_et_al['cat_K'] = silica_et_al['K2O_pct']/94.2*2
silica_et_al['cat_P'] = silica_et_al['P2O5_pct']/141.95*2
silica_et_al['cat_S'] = silica_et_al['SO3_pct']/80.06*1

silica_et_al['cat_total'] = silica_et_al['cat_Si'] + silica_et_al['cat_Ti'] + silica_et_al['cat_Al'] + silica_et_al['cat_Fe'] + \
    silica_et_al['cat_Mg'] + silica_et_al['cat_Mn'] + silica_et_al['cat_Ca'] + silica_et_al['cat_Na'] + silica_et_al['cat_K'] + \
    silica_et_al['cat_P'] + silica_et_al['cat_S']

silica_et_al['cat_total_ternary'] = silica_et_al['cat_Al'] + silica_et_al['cat_Mg'] + silica_et_al['cat_Fe'] + silica_et_al['cat_Ti']

silica_et_al['cat_Fe_Ti_ternary'] = (silica_et_al['cat_Ti'] + silica_et_al['cat_Fe'])/silica_et_al['cat_total_ternary']
silica_et_al['cat_Mg_ternary'] = silica_et_al['cat_Mg']/silica_et_al['cat_total_ternary']
silica_et_al['cat_Al_ternary'] = silica_et_al['cat_Al']/silica_et_al['cat_total_ternary']

# make a ternary plot based on: https://colab.research.google.com/github/agile-geoscience/xlines/blob/master/notebooks/12_Ternary_diagrams.ipynb#scrollTo=NuDIFksmNObR

fig, tax = ternary.figure(scale=1)
fig.set_size_inches(5, 4.5)
# tax.horizontal_line(16)
# draw lines (Rollinson, page 62)
# p1 = (22, 8, 10)
# p2 = (2, 22, 16)
# tax.line(p1, p2, linewidth=3., marker='s', color='green', linestyle=":")

tax.scatter(silica_et_al[['cat_Mg_ternary', 'cat_Fe_Ti_ternary', 'cat_Al_ternary']].values, marker='o')
tax.gridlines(multiple=10)
tax.get_axes().axis('off')
# axis labels
fontsize = 12
offset = 0.14
tax.set_title("Jensen plot attempt\n", fontsize=fontsize)
tax.right_corner_label("Mg cation %", fontsize=fontsize)
tax.top_corner_label("Ti + Fe (total) cation %", fontsize=fontsize)
tax.left_corner_label("Al cation %", fontsize=fontsize)

# make a density plot, or a heatmap (color by Cu content)
# draw lines (Rollinson, page 62) - work in progress, doesn't want to plot
# p1 = (22, 8, 10)
# p2 = (2, 22, 16)
# tax.line(p1, p2, linewidth=3., marker='s', color='green', linestyle=":")
# p1, p2 = (90, 10,), (53.5, 28.5)
# tax.line(p1, p2, linewidth=3, color='k', alpha=0.35, linestyle="-")
# p1, p2 = (75, 25, 0), (75, 0, 25)
# tax.line(p1, p2, linewidth=3, color='k', alpha=0.35, linestyle="-")
# 90, 10, 0; 53.5, 28.5, 18; 52.5, 29,
# 18.5; 51.5, 29, 19.5; 50.5, 27.5, 22;
# 50.3, 25, 24.7; 50.8, 20, 29.2; 51.5,
# 12.5, 36

plt.savefig('ternary.jpg', dpi = 300)

# select the data that is likely to be of interest: all major elements, plus some indicator elements
# we test whether we can predict Cu and Ni from other selected elements

# define list of elements for each dataset separately

Cu_and_els = ['SiO2_pct',
         'TiO2_pct',
         'Al2O3_pct',
         'Fe2O3T_pct',
         'MgO_pct',
         'Sc_ppm',
         'Cr_ppm',
         'MnO_pct',
         'CaO_pct',
         'Na2O_pct',
         'K2O_pct',
         'P2O5_pct',
         'SO3_pct',
         'Sr_ppm',
         'Zr_ppm',
         'Y_ppm',
         'V_ppm',
         'Hf_ppm',
         'Ba_ppm',
         'Nb_ppm',
         'Cu_ppm',
         'Preferred chemical classification']

Ni_and_els = ['SiO2_pct',
         'TiO2_pct',
         'Al2O3_pct',
         'Fe2O3T_pct',
         'MgO_pct',
         'Sc_ppm',
         'Cr_ppm',
         'MnO_pct',
         'CaO_pct',
         'Na2O_pct',
         'K2O_pct',
         'P2O5_pct',
         'SO3_pct',
         'Sr_ppm',
         'Zr_ppm',
         'Y_ppm',
         'V_ppm',
         'Hf_ppm',
         'Ba_ppm',
         'Nb_ppm',
         'Ni_ppm']

# clean up data and convert to numerical
df_Cu = df[Cu_and_els].copy().dropna()
for i in Cu_and_els[:-1]:
    df_Cu[i] = df_Cu[i].apply(item_cleanup)

# df_Ni = df[Ni_and_els].copy()
# df_Ni = df_Ni.dropna()
# df_Ni = df_Ni.apply(data_cleanup)

# Split the dataset into training, validation and test sets. We should not really do it like that in
# geology, because the datasets are often not balanced and we need another kind of split that takes into
# account rock types

from sklearn.model_selection import train_test_split
df_full_train, df_test = train_test_split(df_Cu, test_size = 0.2, random_state = 1)

# convert objects to numerical
# we still have "objects" instead of numeric values in the dataset, I have no clue why
# so I convert them here explicitly

df_full_train['Cu_ppm'] = df_full_train['Cu_ppm'].astype(float)
df_full_train['Sc_ppm'] = df_full_train['Sc_ppm'].astype(float)
df_full_train['Cr_ppm'] = df_full_train['Cr_ppm'].astype(float)
df_full_train['MnO_pct'] = df_full_train['MnO_pct'].astype(float)
df_full_train['CaO_pct'] = df_full_train['CaO_pct'].astype(float)
df_full_train['Na2O_pct'] = df_full_train['Na2O_pct'].astype(float)
df_full_train['K2O_pct'] = df_full_train['K2O_pct'].astype(float)
df_full_train['SO3_pct'] = df_full_train['SO3_pct'].astype(float)
df_full_train['Sr_ppm'] = df_full_train['Sr_ppm'].astype(float)
df_full_train['Zr_ppm'] = df_full_train['Zr_ppm'].astype(float)
df_full_train['V_ppm'] = df_full_train['V_ppm'].astype(float)
df_full_train['Hf_ppm'] = df_full_train['Hf_ppm'].astype(float)
df_full_train['Ba_ppm'] = df_full_train['Ba_ppm'].astype(float)
df_full_train['Nb_ppm'] = df_full_train['Nb_ppm'].astype(float)

df_train, df_val = train_test_split(df_full_train, test_size = 0.25)

# reset indices
df_train = df_train.reset_index(drop = True)
df_val = df_train.reset_index(drop = True)
df_test = df_train.reset_index(drop = True)

y_train = df_train.Cu_ppm.values
y_val = df_val.Cu_ppm.values
y_test = df_test.Cu_ppm.values

# just in case our Cu long tail is real, we do a log transform and see if it changes the RMSE

y_train_log = np.log1p(df_train.Cu_ppm).values
y_val_log = np.log1p(df_val.Cu_ppm).values
y_test_log = np.log1p(df_test.Cu_ppm).values

# delete the target variable - could cause overfitting if occasionally used for training

del df_train['Cu_ppm']
del df_val['Cu_ppm']
del df_test['Cu_ppm']
# 1.2 Feature importance

# check correlation
X_columns = ['SiO2_pct',
         'TiO2_pct',
         'Al2O3_pct',
         'Fe2O3T_pct',
         'MgO_pct',
         'Sc_ppm',
         'Cr_ppm',
         'MnO_pct',
         'CaO_pct',
         'Na2O_pct',
         'K2O_pct',
         'P2O5_pct',
         'SO3_pct',
         'Sr_ppm',
         'Zr_ppm',
         'Y_ppm',
         'V_ppm',
         'Hf_ppm',
         'Ba_ppm',
         'Nb_ppm']

# check the correlation between each input element and the output element, Cu
corr_check = df_full_train[X_columns].corrwith(df_full_train.Cu_ppm)

# None of the elements are highly correlated with Cu

# Logistic regression with scikit learn
from sklearn.linear_model import LinearRegression
model_LR = LinearRegression()
X_train = df_train[X_columns]
X_val = df_val[X_columns]
X_test = df_test[X_columns]
model_LR.fit(X_train, y_train)
y_pred = model_LR.predict(X_val)
rmse_LR = np.sqrt(mean_squared_error(y_val, y_pred))

# rmse of 65 for predicting ppms of Cu means that whenever we are predicting Cu values from other elements with
# linear regression, we will be on average 65 ppm off from the true value for the validation dataset
# plot the predictions
plt.figure()
sns.histplot (y_pred, color = 'red', alpha = 0.5, bins = 50)
sns.histplot(y_val, color = 'blue', alpha = 0.5, bins = 50)
plt.xlim(0, 500)
plt.show()
# from the figure we see that the predictions are kind of off, and Cu had a long tail, so we will test on the log-transformed variable

model_LR_log = LinearRegression()
X_train = df_train[X_columns]
X_val = df_val[X_columns]
X_test = df_test[X_columns]
model_LR_log.fit(X_train, y_train_log)
y_pred_log = model_LR_log.predict(X_val)
rmse_LR_log = np.sqrt(mean_squared_error(y_val_log, y_pred_log))
rmse_LR_log_exp = np.exp(rmse_LR_log)

# plot the predictions that resulted from linear regression on the log-transformed values;
# to be comparable with the prediction on the non-log-transformed values, take the exponent
y_pred_LR_log = np.exp(y_pred_log)

plt.figure()
sns.histplot (y_pred_LR_log, color = 'red', alpha = 0.5, bins = 50)
sns.histplot(y_val, color = 'blue', alpha = 0.5, bins = 50)
plt.xlim(0, 500)
plt.show()

# logtransform really means improvement in rmse and in the character of data distribution;
# however, it looks like it was safe to predict low values because of the structure of the dataset,
# so the linear regression got pessimistic

# let's try with polynomial regression on log-transformed data
# we train the model with several order of the polynomials and compare the results
order_polynomial = [2, 3, 4]

for i in order_polynomial:
    polynomial_features = PolynomialFeatures(degree=i)
    X_train_poly = polynomial_features.fit_transform(X_train)
    X_val_poly = polynomial_features.fit_transform(X_val)

    model_polynomial = LinearRegression()
    model_polynomial.fit(X_train_poly, y_train)
    y_poly_pred = model_polynomial.predict(X_val_poly)
    plt.figure()
    sns.histplot (y_poly_pred, color = 'red', alpha = 0.5, bins = 50)
    sns.histplot(y_val, color = 'blue', alpha = 0.5, bins = 50)
    # plt.xlim(0, 500)
    plt.show()
    rmse = np.sqrt(mean_squared_error(y_val, y_poly_pred))
    r2 = r2_score(y_val, y_poly_pred)
    print('polynomial of order', i, 'rmse', rmse)

# I choose the 4th order, although it seems it's overfitting, it reproduces the shape of the data distribution best
# besides, both 2nd and 3rd order polynomials ended up predicting some of the Cu values as negative
# does it make sense to work with polynomials of the log-transformed values?

# Now let's try to use dictvectorizer to code the output variable and predict rock types from chemistry
# first additional manual data clean-up: removing the spaces, checking the rock types for typos
# the dataset is not balanced, but we will see how it gets split;
# an idea is to use only rock types with more than 50 samples in the group

# Logistic regression

from sklearn.linear_model import LogisticRegression
numerical = ['SiO2_pct',
         'TiO2_pct',
         'Al2O3_pct',
         'Fe2O3T_pct',
         'MgO_pct',
         'Sc_ppm',
         'Cr_ppm',
         'MnO_pct',
         'CaO_pct',
         'Na2O_pct',
         'K2O_pct',
         'P2O5_pct',
         'SO3_pct',
         'Sr_ppm',
         'Zr_ppm',
         'Y_ppm',
         'V_ppm',
         'Hf_ppm',
         'Ba_ppm',
         'Nb_ppm']

# choose rock types as y variable

categorical = ['Preferred chemical classification']

# make a new split of the training, validation and test sets trying to keep it balanced with stratify
df_full_train_rocks, df_test_rocks = train_test_split(df_Cu, test_size = 0.2, random_state = 1, stratify=df_Cu[categorical])

# convert objects to numerical
# we still have "objects" instead of numeric values in the dataset, I have no clue why
# so I convert them here explicitly

df_full_train['Cu_ppm'] = df_full_train['Cu_ppm'].astype(float)
df_full_train['Sc_ppm'] = df_full_train['Sc_ppm'].astype(float)
df_full_train['Cr_ppm'] = df_full_train['Cr_ppm'].astype(float)
df_full_train['MnO_pct'] = df_full_train['MnO_pct'].astype(float)
df_full_train['CaO_pct'] = df_full_train['CaO_pct'].astype(float)
df_full_train['Na2O_pct'] = df_full_train['Na2O_pct'].astype(float)
df_full_train['K2O_pct'] = df_full_train['K2O_pct'].astype(float)
df_full_train['SO3_pct'] = df_full_train['SO3_pct'].astype(float)
df_full_train['Sr_ppm'] = df_full_train['Sr_ppm'].astype(float)
df_full_train['Zr_ppm'] = df_full_train['Zr_ppm'].astype(float)
df_full_train['V_ppm'] = df_full_train['V_ppm'].astype(float)
df_full_train['Hf_ppm'] = df_full_train['Hf_ppm'].astype(float)
df_full_train['Ba_ppm'] = df_full_train['Ba_ppm'].astype(float)
df_full_train['Nb_ppm'] = df_full_train['Nb_ppm'].astype(float)

df_train_rocks, df_val_rocks = train_test_split(df_full_train_rocks, test_size = 0.25, stratify=df_full_train_rocks[categorical])

# reset indices
df_train_rocks = df_train_rocks.reset_index(drop = True)
df_val_rocks = df_train_rocks.reset_index(drop = True)
df_test_rocks = df_train_rocks.reset_index(drop = True)

X_train_rocks = df_train_rocks[numerical].values
X_val_rocks = df_val_rocks[numerical].values
X_test_rocks = df_test_rocks[numerical].values

y_train_rocks = df_train_rocks[categorical].values
y_val_rocks = df_val_rocks[categorical].values
y_test_rocks = df_test_rocks[categorical].values

lr = LogisticRegression()

lr.fit(X_train_rocks, y_train_rocks)

y_pred_rocks = lr.predict(X_val_rocks)
accuracy_LR = (y_pred_rocks == y_val_rocks).mean()

# This was done on the dataset with labels "as is", and whoops, we can only guess the rock type in 10% of the cases. A random number generator could do a better job.
# Let's take a smaller subset which is balanced

# increase of accuracy to 12% - where did it all go wrong?

from sklearn.tree import DecisionTreeClassifier

dt = DecisionTreeClassifier(max_depth = 4)
dt.fit(X_train_rocks, y_train_rocks)
dt_predict = dt.predict(X_val_rocks)

# confusion matrix
from sklearn.metrics import confusion_matrix
confusion_matrix(y_val_rocks, dt_predict)
accuracy_dt = (dt_predict == y_val_rocks).mean()

rocks_in_train = np.unique(df_train_rocks[categorical].values)
rocks_in_val = np.unique(df_val_rocks[categorical].values)
rocks_in_test = np.unique(df_test_rocks[categorical].values)

# save the confusion matrix
conf = confusion_matrix(y_val_rocks, dt_predict)
df = pd.DataFrame(conf)

plt.figure(figsize = (15,15))
cm = conf
ax = plt.subplot()
sns.heatmap(cm, annot=True, ax = ax, cmap='Greens'); #annot=True to annotate cells

# labels, title and ticks
ax.set_xlabel('Predicted labels');ax.set_ylabel('True labels');
ax.set_title('Confusion Matrix, decision trees');
ax.xaxis.set_ticklabels(['Basalt High-Ti', 'Dolerite LTB', 'Dolerite/gabbro LTB', 'EMI',
       'Felsic dyke (dacitic)', 'Felsic dyke (rhyolitic)',
       'Felsic vol/volclas (andesitic)', 'Felsic vol/volclas (dacitic)',
       'Felsic vol/volclas (rhyolitic)', 'Granite', 'HTSB', 'ITB',
       'Komatiite', 'Komatiite High-Al', 'LTB', 'LTB High-Nb',
       'Problematic data'], rotation = 90);
ax.yaxis.set_ticklabels(['Basalt High-Ti', 'Dolerite LTB', 'Dolerite/gabbro LTB', 'EMI',
       'Felsic dyke (dacitic)', 'Felsic dyke (rhyolitic)',
       'Felsic vol/volclas (andesitic)', 'Felsic vol/volclas (dacitic)',
       'Felsic vol/volclas (rhyolitic)', 'Granite', 'HTSB', 'ITB',
       'Komatiite', 'Komatiite High-Al', 'LTB', 'LTB High-Nb',
       'Problematic data'], rotation = 45);

plt.savefig('confusion matrix dtree.jpg', dpi = 300)

# I have 17 groups! Why are there 20?
res = pd.DataFrame(dt_predict)
res.to_csv("prediction_results.csv")

# when in doubt, use Gradient Boost

from sklearn.ensemble import GradientBoostingClassifier
gb = GradientBoostingClassifier(n_estimators=100, learning_rate=1.0, max_depth=1, random_state=0).fit(X_train_rocks, y_train_rocks)
gb.fit(X_train_rocks, y_train_rocks)
gb_predict = gb.predict(X_val_rocks)
gb_score = gb.score(X_val_rocks, y_val_rocks)

confusion_matrix(y_val_rocks, gb_predict)
accuracy_gb = (gb_predict == y_val_rocks).mean()

# save the confusion matrix
conf_gb = confusion_matrix(y_val_rocks, gb_predict)
df = pd.DataFrame(conf_gb)

plt.figure(figsize = (15,15))
cm = conf_gb
ax = plt.subplot()
sns.heatmap(cm, annot=True, ax = ax, cmap='Greens'); #annot=True to annotate cells

# labels, title and ticks
ax.set_xlabel('Predicted labels');ax.set_ylabel('True labels');
ax.set_title('Confusion Matrix, Gradient boost');
ax.xaxis.set_ticklabels(['Basalt High-Ti', 'Dolerite LTB', 'Dolerite/gabbro LTB', 'EMI',
       'Felsic dyke (dacitic)', 'Felsic dyke (rhyolitic)',
       'Felsic vol/volclas (andesitic)', 'Felsic vol/volclas (dacitic)',
       'Felsic vol/volclas (rhyolitic)', 'Granite', 'HTSB', 'ITB',
       'Komatiite', 'Komatiite High-Al', 'LTB', 'LTB High-Nb',
       'Problematic data'], rotation = 90);
ax.yaxis.set_ticklabels(['Basalt High-Ti', 'Dolerite LTB', 'Dolerite/gabbro LTB', 'EMI',
       'Felsic dyke (dacitic)', 'Felsic dyke (rhyolitic)',
       'Felsic vol/volclas (andesitic)', 'Felsic vol/volclas (dacitic)',
       'Felsic vol/volclas (rhyolitic)', 'Granite', 'HTSB', 'ITB',
       'Komatiite', 'Komatiite High-Al', 'LTB', 'LTB High-Nb',
       'Problematic data'], rotation = 45);

plt.savefig('confusion matrix gboost.jpg', dpi = 300)

# whatever. Random forest.
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import GridSearchCV
n_estimators = [50,100,200,500,1000]
max_depth = [3,5,10,15,20,None]
##
parameters = {'n_estimators':n_estimators, 'max_depth':max_depth}
classifier = RandomForestClassifier(random_state=27)
modelGS = GridSearchCV(estimator=classifier,param_grid=parameters,cv=10,n_jobs=-1)
modelGS.fit(X_train_rocks, y_train_rocks)

modelGS_predict = modelGS.predict(X_val_rocks)
IGS_score = modelGS.score(X_val_rocks, y_val_rocks)

confusion_matrix(y_val_rocks, modelGS_predict)
accuracy_rf = (modelGS_predict == y_val_rocks).mean()

# save the confusion matrix
conf_rf = confusion_matrix(y_val_rocks, modelGS_predict)
df = pd.DataFrame(conf_rf)

plt.figure(figsize = (15,15))
cm = conf_rf
ax = plt.subplot()
sns.heatmap(cm, annot=True, ax = ax, cmap='Greens'); #annot=True to annotate cells

# labels, title and ticks
ax.set_xlabel('Predicted labels');ax.set_ylabel('True labels');
ax.set_title('Confusion Matrix, random forest');
ax.xaxis.set_ticklabels(['Basalt High-Ti', 'Dolerite LTB', 'Dolerite/gabbro LTB', 'EMI',
       'Felsic dyke (dacitic)', 'Felsic dyke (rhyolitic)',
       'Felsic vol/volclas (andesitic)', 'Felsic vol/volclas (dacitic)',
       'Felsic vol/volclas (rhyolitic)', 'Granite', 'HTSB', 'ITB',
       'Komatiite', 'Komatiite High-Al', 'LTB', 'LTB High-Nb',
       'Problematic data'], rotation = 90);
ax.yaxis.set_ticklabels(['Basalt High-Ti', 'Dolerite LTB', 'Dolerite/gabbro LTB', 'EMI',
       'Felsic dyke (dacitic)', 'Felsic dyke (rhyolitic)',
       'Felsic vol/volclas (andesitic)', 'Felsic vol/volclas (dacitic)',
       'Felsic vol/volclas (rhyolitic)', 'Granite', 'HTSB', 'ITB',
       'Komatiite', 'Komatiite High-Al', 'LTB', 'LTB High-Nb',
       'Problematic data'], rotation = 45);

plt.savefig('confusion random forest.jpg', dpi = 300)

# let's see how it performs on the test set

dt_test_predict = dt.predict(X_test_rocks)
# save the confusion matrix
conf_dt = confusion_matrix(y_test_rocks, dt_test_predict)
df = pd.DataFrame(conf_dt)

plt.figure(figsize = (15,15))
cm = conf_dt
ax = plt.subplot()
sns.heatmap(cm, annot=True, ax = ax, cmap='Greens'); #annot=True to annotate cells

# labels, title and ticks
ax.set_xlabel('Predicted labels');ax.set_ylabel('True labels');
ax.set_title('Confusion Matrix, decision trees, test set');
ax.xaxis.set_ticklabels(['Basalt High-Ti', 'Dolerite LTB', 'Dolerite/gabbro LTB', 'EMI',
       'Felsic dyke (dacitic)', 'Felsic dyke (rhyolitic)',
       'Felsic vol/volclas (andesitic)', 'Felsic vol/volclas (dacitic)',
       'Felsic vol/volclas (rhyolitic)', 'Granite', 'HTSB', 'ITB',
       'Komatiite', 'Komatiite High-Al', 'LTB', 'LTB High-Nb',
       'Problematic data'], rotation = 90);
ax.yaxis.set_ticklabels(['Basalt High-Ti', 'Dolerite LTB', 'Dolerite/gabbro LTB', 'EMI',
       'Felsic dyke (dacitic)', 'Felsic dyke (rhyolitic)',
       'Felsic vol/volclas (andesitic)', 'Felsic vol/volclas (dacitic)',
       'Felsic vol/volclas (rhyolitic)', 'Granite', 'HTSB', 'ITB',
       'Komatiite', 'Komatiite High-Al', 'LTB', 'LTB High-Nb',
       'Problematic data'], rotation = 45);

plt.savefig('confusion matrix dtree test.jpg', dpi = 300)

gb_predict_test = gb.predict(X_test_rocks)
gb_score = gb.score(X_test_rocks, y_test_rocks)

confusion_matrix(y_test_rocks, gb_predict)
accuracy_gb = (gb_predict_test == y_test_rocks).mean()

# save the confusion matrix
conf_gb_test = confusion_matrix(y_test_rocks, gb_predict_test)
df = pd.DataFrame(conf_gb_test)

plt.figure(figsize = (15,15))
cm = conf_gb_test
ax = plt.subplot()
sns.heatmap(cm, annot=True, ax = ax, cmap='Greens'); #annot=True to annotate cells

# labels, title and ticks
ax.set_xlabel('Predicted labels');ax.set_ylabel('True labels');
ax.set_title('Confusion Matrix, Gradient boost, test set');
ax.xaxis.set_ticklabels(['Basalt High-Ti', 'Dolerite LTB', 'Dolerite/gabbro LTB', 'EMI',
       'Felsic dyke (dacitic)', 'Felsic dyke (rhyolitic)',
       'Felsic vol/volclas (andesitic)', 'Felsic vol/volclas (dacitic)',
       'Felsic vol/volclas (rhyolitic)', 'Granite', 'HTSB', 'ITB',
       'Komatiite', 'Komatiite High-Al', 'LTB', 'LTB High-Nb',
       'Problematic data'], rotation = 90);
ax.yaxis.set_ticklabels(['Basalt High-Ti', 'Dolerite LTB', 'Dolerite/gabbro LTB', 'EMI',
       'Felsic dyke (dacitic)', 'Felsic dyke (rhyolitic)',
       'Felsic vol/volclas (andesitic)', 'Felsic vol/volclas (dacitic)',
       'Felsic vol/volclas (rhyolitic)', 'Granite', 'HTSB', 'ITB',
       'Komatiite', 'Komatiite High-Al', 'LTB', 'LTB High-Nb',
       'Problematic data'], rotation = 45);

plt.savefig('confusion matrix gboost test.jpg', dpi = 300)

modelGS_predict_test = modelGS.predict(X_test_rocks)
IGS_score_test = modelGS.score(X_test_rocks, y_test_rocks)

confusion_matrix(y_test_rocks, modelGS_predict)
accuracy_rf = (modelGS_predict == y_test_rocks).mean()

# save the confusion matrix
conf_rf_test = confusion_matrix(y_test_rocks, modelGS_predict_test)
df = pd.DataFrame(conf_rf_test)

plt.figure(figsize = (15,15))
cm = conf_rf_test
ax = plt.subplot()
sns.heatmap(cm, annot=True, ax = ax, cmap='Greens'); #annot=True to annotate cells

# labels, title and ticks
ax.set_xlabel('Predicted labels');ax.set_ylabel('True labels');
ax.set_title('Confusion Matrix, random forest, test set');
ax.xaxis.set_ticklabels(['Basalt High-Ti', 'Dolerite LTB', 'Dolerite/gabbro LTB', 'EMI',
       'Felsic dyke (dacitic)', 'Felsic dyke (rhyolitic)',
       'Felsic vol/volclas (andesitic)', 'Felsic vol/volclas (dacitic)',
       'Felsic vol/volclas (rhyolitic)', 'Granite', 'HTSB', 'ITB',
       'Komatiite', 'Komatiite High-Al', 'LTB', 'LTB High-Nb',
       'Problematic data'], rotation = 90);
ax.yaxis.set_ticklabels(['Basalt High-Ti', 'Dolerite LTB', 'Dolerite/gabbro LTB', 'EMI',
       'Felsic dyke (dacitic)', 'Felsic dyke (rhyolitic)',
       'Felsic vol/volclas (andesitic)', 'Felsic vol/volclas (dacitic)',
       'Felsic vol/volclas (rhyolitic)', 'Granite', 'HTSB', 'ITB',
       'Komatiite', 'Komatiite High-Al', 'LTB', 'LTB High-Nb',
       'Problematic data'], rotation = 45);

plt.savefig('confusion random forest test set.jpg', dpi = 300)


import torch
from torch import nn, optim
from transformers import GPT2TokenizerFast, GPT2Tokenizer, BertTokenizer
from util_fucntions import util_functions
#from lightning.pytorch.trainer import Trainer
#from lightning.pytorch.callbacks.early_stopping import EarlyStopping

MODEL_ITERATION = 1

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "mps")
print(f'Running on {DEVICE}')
INIT_LEARNING_RATE = 1e-5
MIN_LEARNING_RATE = 1e-8

LABEL_KEY_NAME = 'harmonised_desc'
INPUT_KEY_NAME = 'curr_desc'
SAMPLE_IDX_CODE_NAME = 'idx'

PRETRAINED_MODEL_NAME = 'bert-base-cased'

IS_BERT = 'bert' in util_functions.lower_case_and_clear_white_space(PRETRAINED_MODEL_NAME)
MODEL_TOKENIZER = BertTokenizer.from_pretrained(PRETRAINED_MODEL_NAME) if IS_BERT else GPT2Tokenizer.from_pretrained(PRETRAINED_MODEL_NAME)
PADDING_TOKEN = None if IS_BERT else '[PAD]'
if PADDING_TOKEN is not None:
    MODEL_TOKENIZER.add_special_tokens({'pad_token': PADDING_TOKEN})
    padding_token_id = MODEL_TOKENIZER.encode(PADDING_TOKEN)[0]

MAX_LENGTH_TOKEN = 256

MSE_REDUCTION = 'mean'
WEIGHT_DECAY = .0001
SCHEDULED = True
SCHEDULER_PATIENCE = 16
EARLY_STOPPER_PATIENCE = 4
DROPOUT = 0.1
OPTIMISER_EPSILON = 1e-08
#EARLY_STOPPING_CALLBACK = EarlyStopping(monitor="val_mae", patience=PATIENCE, mode="min")
#MY_TRAINER = Trainer(callbacks=[EARLY_STOPPING_CALLBACK])

# Model config
HIDDEN_LAYER_SIZE = 16

NUM_WORKERS = 0
NUM_EPOCHS = 1000
VAL_EPOCH = 20
SAVED_EPOCH = 200
TRAIN_BATCH_SIZE = 18
VAL_BATCH_SIZE = 4
SUB_SAMPLE_SIZE_PER_CLASS = None

RANDOM_SEED = 10
TRAIN_RATIO = .7
VALIDATION_RATIO = .15
TEST_RATIO = .15
MULTI_DATA_FILES = []
for version in range(0, 4):
    MULTI_DATA_FILES.append(f'../../../xlsx_resources/for_trainings/maximo_pm_desc_map_reworked_S/maximo_pm_desc_map_reworked_it{version}.xlsx')

WANDB_PROJECT_NAME = f'Maximo to GAP App harmonised description BERT Base Cased Model Training Iter{MODEL_ITERATION}'
RUNNING_LOG_LOCATION = f'saved_logs/bert_based_cased_run_logs/running_iteration_{MODEL_ITERATION}.txt'
SAVED_MODEL_LOCATION = f'saved_models/bert_based_cased_models_iteration_{MODEL_ITERATION}/'
SAVED_MODEL_FORMAT = '.pt'

SAVED_UNTRAINED_SAMPLE_IDX_LOCATION = f'saved_untrained_wos/bert_based_cased/running_iteration_{MODEL_ITERATION}.json'
SAVED_TRAINED_SAMPLE_IDX_LOCATION = f'saved_trained_wos/bert_based_cased/running_iteration_{MODEL_ITERATION}.json'

# Static class label file
STATIC_CLASS_LABEL_FILE_LOCATION = 'static_classes.json'
STATIC_CLASS_IDX_KEY_NAME = 'class_idx'
STATIC_CLASS_LABEL_KEY_NAME = 'class_label'
STATIC_ENUMERATED_CLASS_DATA_NAME = 'enumerated_class_data'


#!/usr/bin/env python3

"""
A class-based system for creating Donor objects and DonorCollections objects.
"""

import os
import functools

@functools.total_ordering
class Donor:
    """Render Donor objects"""

    def __init__(self, name, amount=0, number=0):
        """Initialize new object with name, amount and number of donations"""

        self.name = name

        if amount >= 0:
            self.donation = amount
        else:
            self.donation = 0

        if isinstance(number, int) and number >= 0:
            self.number = number
        else:
            self.number = 0

    def __str__(self):
        """Return string with Donor name"""
        return "Donor: {}".format(self.name)

    def __eq__(self, other):
        """Define equals operator"""
        return self.donation == other.donation

    def __lt__(self, other):
        """Define less than operator"""
        return self.donation < other.donation

    def add_donation(self, amount):
        """Add donation amount for existing donor (and update number of donations)"""
        if amount >= 0:
            self.donation += amount
            self.add_number(1)

    def add_number(self, number):
        """Add to number of existing donations"""
        if isinstance(number, int) and number >= 0:
            self.number += number

    @property
    def average(self):
        """Return read only attribute for average donation"""
        return float("{:.2f}".format(self.donation / self.number))

    @property
    def write_letter(self):
        """Return string of thank you letter text to donor"""
        letter = """
        Dear {},
        Thank you for your collective contributions of ${:.2f} over the years.
        Your generous donations have been put to good use!
        Sincerely,
        Kyle at Kelby Doggo, Inc\n""".format(self.name, self.donation)
        return letter

    def send_letter(self):
        """Write thank you letter to text file in local directory"""
        text = self.write_letter
        name_list = (self.name.replace(",", "")).split(" ")
        file_name = ("_".join(name_list)) + ".txt"
        with open(file_name, "w") as a_file:
            a_file.write(text)

class DonorCollections:
    """Render custom class object for set of Donor objects"""

    def __init__(self, donors=None):
        """Initialize DonorCollections as sorted list of Donors"""
        try:
            self.list = [x for x in donors]
            self.list.sort(reverse=True)
        except TypeError:
            self.list = []

    def __str__(self):
        """Return string with Donor names in DonorCollections"""
        text = [x.name for x in self.list]
        return "Donor Collection: {}".format(str(text))

    @property
    def get_names(self):
        """Return read only list of Donor Names in DonorCollections"""
        names = []
        for x in self.list:
            names.append(x.name)
        return names

    def add_donor(self, donor):
        """Add new donor to Donor List"""
        self.list.append(donor)
        self.list.sort(reverse=True)

    def format(self):
        """Return list of formatted rows for donors"""
        self.list.sort(reverse=True)
        heading = "| {dn:<20s}\t| {tg:<15s}\t| {ng:<10s}| {ag:<15s}  |".format
        report_list = [heading(dn="Donor Name", tg="Total Given",
                               ng="Num Gifts", ag="Average Gift"), 76 * "*"]
        row = "| {dn:<20s} \t| {ds:<1s} {tg:>14.2f} \t|" \
              "{ng:>10d}\t| {ds2:<1} {ag:>14.2f} |".format
        for i in self.list:
            report_list.append(row(dn=i.name, ds="$", tg=i.donation,
                                   ng=i.number, ds2="$", ag=i.average))
        return report_list

    def send_letters(self):
        """Write thank you letters to text files for all donors in DonorCollections"""
        for d in self.list:
            d.send_letter()


from rest_framework import permissions
from rest_framework import viewsets, mixins

from staff_profile_app.models import StaffProfile
from staff_profile_app.serializers import StaffProfileSerializer
from student_profile_app.serializers import StudentProfessionSerializer, StudentProfileSerializer, FieldReportSerializer
from alumni_profile_app.serializers import AlumniProfileSerializer, AlumniProfessionSerializer, PublishedAlumniSerializer
from student_profile_app.models import StudentProfession, StudentProfile, FieldReport
from alumni_profile_app.models import AlumniProfile, AlumniProfession
from program_app.serializers import ProgramSerializer
from program_app.models import Program
from field_post_app.models import FieldApplication, FieldPost, FieldPostProfession, FieldPostProgram
from field_post_app.serializers import FieldPostProfessionSerializer, FieldApplicationSerializer, FieldPostSerializer, FieldPostProgramSerializer
from internship_post_app.models import InternshipApplication, InternshipPost, InterviewSchedule
from internship_post_app.serializers import InternshipPostSerializer, InternshipApplicationSerializer, InterviewScheduleSerializer
from .models import UserProfile
from .serializers import UserProfileSerializer
from announcement_app.models import Announcement
from announcement_app.serializers import AnnouncementSerializer
from project_app.models import ProjectMember
from project_app.serializers import ProjectMemberSerializer
from questions_app.models import (MultipleChoice, ApplicantScore,
    ApplicantAnswer, Question)
from questions_app.serializers import (MultipleChoiceSerializer,
    ApplicantScoreSerializer, ApplicantAnswerSerializer,
    QuestionSerializer)
from organization_address_app.serializers import (ContractSerializer,
    OrganizationProfileSerializer, InvitationsSerializer)
from organization_address_app.models import Contract, OrganizationProfile, Invitations
from cv_app.serializers import (PersonalInformationSerializer,
    EducationInformationSerializer, ExperienceInformationSerializer)
from cv_app.models import (PersonalInformation, EducationInformation,
    ExperienceInformation)
from certificate_app.models import Certificate
from certificate_app.serializers import CertificateSerializer

class InterviewScheduleByPostViewSet(viewsets.GenericViewSet, mixins.ListModelMixin, mixins.RetrieveModelMixin):
    serializer_class = InterviewScheduleSerializer
    permission_classes = [
        permissions.IsAuthenticated ]

    def get_queryset(self):
        post_id = self.kwargs.get('postId')
        return InterviewSchedule.objects.filter(post=post_id)

class QuestionChoicesViewSet(viewsets.GenericViewSet, mixins.ListModelMixin, mixins.RetrieveModelMixin):
    serializer_class = MultipleChoiceSerializer
    permission_classes = [
        permissions.IsAuthenticated ]

    def get_queryset(self):
        question_id = self.kwargs.get('questionId')
        return MultipleChoice.objects.filter(question=question_id)

class QuestionsByProfessionsViewSet(viewsets.GenericViewSet, mixins.ListModelMixin, mixins.RetrieveModelMixin):
    serializer_class = QuestionSerializer
    permission_classes = [
        permissions.IsAuthenticated ]

    def get_queryset(self):
        profession_id = self.kwargs.get('professionId')
        return Question.objects.filter(profession=profession_id)

class ApplicantMarksViewSet(viewsets.GenericViewSet, mixins.ListModelMixin, mixins.RetrieveModelMixin):
    serializer_class = ApplicantScoreSerializer
    permission_classes = [
        permissions.IsAuthenticated ]

    def get_queryset(self):
        alumni_id = self.kwargs.get('alumniId')
        post_id = self.kwargs.get('postId')
        return ApplicantScore.objects.filter(alumni=alumni_id, post=post_id)

class ApplicantAnswersViewSet(viewsets.GenericViewSet, mixins.ListModelMixin, mixins.RetrieveModelMixin):
    serializer_class = ApplicantAnswerSerializer
    permission_classes = [
        permissions.IsAuthenticated ]

    def get_queryset(self):
        alumni_id = self.kwargs.get('alumniId')
        post_id = self.kwargs.get('postId')
        return ApplicantAnswer.objects.filter(alumni=alumni_id, post=post_id)

class ProjectsByMemberViewSet(viewsets.GenericViewSet, mixins.ListModelMixin, mixins.RetrieveModelMixin):
    serializer_class = ProjectMemberSerializer
    permission_classes = [
        permissions.IsAuthenticated ]

    def get_queryset(self):
        memberId = self.kwargs.get('memberId')
        return ProjectMember.objects.filter(member=memberId)

class ProjectMembersViewSet(viewsets.GenericViewSet, mixins.ListModelMixin, mixins.RetrieveModelMixin):
    serializer_class = ProjectMemberSerializer
    permission_classes = [
        permissions.IsAuthenticated ]

    def get_queryset(self):
        projectId = self.kwargs.get('projectId')
        return ProjectMember.objects.filter(project=projectId)

class AnnouncementsByDesignationsViewSet(viewsets.GenericViewSet, mixins.ListModelMixin, mixins.RetrieveModelMixin):
    serializer_class = AnnouncementSerializer
    permission_classes = [
        permissions.IsAuthenticated ]

    def get_queryset(self):
        designation_id = self.kwargs.get('designationId')
        return Announcement.objects.filter(destination=designation_id)

class UsersProfilesByDesignationIdViewSet(viewsets.GenericViewSet, mixins.ListModelMixin, mixins.RetrieveModelMixin):
    serializer_class = UserProfileSerializer
    permission_classes = [
        permissions.IsAuthenticated ]

    def get_queryset(self):
        designation_id = self.kwargs.get('designationId')
        return UserProfile.objects.filter(designation=designation_id)

class SingleUserProfileViewSet(viewsets.GenericViewSet, mixins.ListModelMixin, mixins.RetrieveModelMixin):
    serializer_class = UserProfileSerializer
    permission_classes = [
        permissions.IsAuthenticated ]

    def get_queryset(self):
        user_id = self.kwargs.get('userId')
        return UserProfile.objects.filter(user=user_id)

class StudentProfessionViewSet(viewsets.GenericViewSet, mixins.ListModelMixin, mixins.RetrieveModelMixin):
    serializer_class = StudentProfessionSerializer
    permission_classes = [
        permissions.IsAuthenticated ]

    def get_queryset(self):
        student_id = self.kwargs.get('studentId')
        return StudentProfession.objects.filter(student=student_id)

class StudentReportViewSet(viewsets.GenericViewSet, mixins.ListModelMixin, mixins.RetrieveModelMixin):
    serializer_class = FieldReportSerializer
    permission_classes = [
        permissions.IsAuthenticated ]

    def get_queryset(self):
        student_id = self.kwargs.get('studentId')
        return FieldReport.objects.filter(student=student_id)

class ProgramsByDepartmentViewSet(viewsets.GenericViewSet, mixins.ListModelMixin, mixins.RetrieveModelMixin):
    serializer_class = ProgramSerializer
    permission_classes = [
        permissions.IsAuthenticated ]

    def get_queryset(self):
        department_id = self.kwargs.get('departmentId')
        return Program.objects.filter(department=department_id)

class StudentsByProgramViewSet(viewsets.GenericViewSet, mixins.ListModelMixin, mixins.RetrieveModelMixin):
    serializer_class = StudentProfileSerializer
    permission_classes = [
        permissions.IsAuthenticated ]

    def get_queryset(self):
        program_id = self.kwargs.get('programId')
        return StudentProfile.objects.filter(program=program_id)

class StudentsByOrganizationViewSet(viewsets.GenericViewSet, mixins.ListModelMixin, mixins.RetrieveModelMixin):
    serializer_class = StudentProfileSerializer
    permission_classes = [
        permissions.IsAuthenticated ]

    def get_queryset(self):
        organization_id = self.kwargs.get('organizationId')
        return StudentProfile.objects.filter(organization=organization_id)

class SingleOrganizationProfileViewSet(viewsets.GenericViewSet, mixins.ListModelMixin, mixins.RetrieveModelMixin):
    serializer_class = OrganizationProfileSerializer
    permission_classes = [
        permissions.IsAuthenticated ]

    def get_queryset(self):
        organization_id = self.kwargs.get('organizationId')
        return OrganizationProfile.objects.filter(organization_id=organization_id)

class StudentsByDepartmentViewSet(viewsets.GenericViewSet, mixins.ListModelMixin, mixins.RetrieveModelMixin):
    serializer_class = StudentProfileSerializer
    permission_classes = [
        permissions.IsAuthenticated ]

    def get_queryset(self):
        department_id = self.kwargs.get('departmentId')
        return StudentProfile.objects.filter(department=department_id)

class StudentsByAcademicSpervisorViewSet(viewsets.GenericViewSet, mixins.ListModelMixin, mixins.RetrieveModelMixin):
    serializer_class = StudentProfileSerializer
    permission_classes = [
        permissions.IsAuthenticated ]

    def get_queryset(self):
        supervisor_id = self.kwargs.get('supervisorId')
        return StudentProfile.objects.filter(academic_supervisor=supervisor_id)

class StudentsByFieldSpervisorViewSet(viewsets.GenericViewSet, mixins.ListModelMixin, mixins.RetrieveModelMixin):
    serializer_class = StudentProfileSerializer
    permission_classes = [
        permissions.IsAuthenticated ]

    def get_queryset(self):
        supervisor_id = self.kwargs.get('supervisorId')
        return StudentProfile.objects.filter(field_supervisor=supervisor_id)

class StudentsByStatusViewSet(viewsets.GenericViewSet, mixins.ListModelMixin, mixins.RetrieveModelMixin):
    serializer_class = StudentProfileSerializer
    permission_classes = [
        permissions.IsAuthenticated ]

    def get_queryset(self):
        status = self.kwargs.get('status')
        return StudentProfile.objects.filter(student_status=status)

class StudentsByYearOfStudyViewSet(viewsets.GenericViewSet, mixins.ListModelMixin, mixins.RetrieveModelMixin):
    serializer_class = StudentProfileSerializer
    permission_classes = [
        permissions.IsAuthenticated ]

    def get_queryset(self):
        year = self.kwargs.get('year')
        return StudentProfile.objects.filter(year_of_study=year)

class ReportedStudentsViewSet(viewsets.GenericViewSet, mixins.ListModelMixin, mixins.RetrieveModelMixin):
    serializer_class = FieldApplicationSerializer
    permission_classes = [
        permissions.IsAuthenticated ]

    def get_queryset(self):
        return FieldApplication.objects.filter(has_reported=True)

class ReportedStudentsProfilesViewSet(viewsets.GenericViewSet, mixins.ListModelMixin, mixins.RetrieveModelMixin):
    serializer_class = StudentProfileSerializer
    permission_classes = [
        permissions.IsAuthenticated ]

    def get_queryset(self):
        return StudentProfile.objects.filter(has_reported=True)

class SingleStudentProfileViewSet(viewsets.GenericViewSet, mixins.ListModelMixin, mixins.RetrieveModelMixin):
    serializer_class = StudentProfileSerializer
    permission_classes = [
        permissions.IsAuthenticated ]

    def get_queryset(self):
        student_id = self.kwargs.get('studentId')
        return StudentProfile.objects.filter(student=student_id)

class AlumniProfessionViewSet(viewsets.GenericViewSet, mixins.ListModelMixin, mixins.RetrieveModelMixin):
    serializer_class = AlumniProfessionSerializer
    permission_classes = [
        permissions.IsAuthenticated ]

    def get_queryset(self):
        alumni_id = self.kwargs.get('alumniId')
        return AlumniProfession.objects.filter(alumni=alumni_id)

class PubllishedAlumniViewSet(viewsets.GenericViewSet, mixins.ListModelMixin, mixins.RetrieveModelMixin):
    serializer_class = PublishedAlumniSerializer
    permission_classes = [
        permissions.IsAuthenticated ]

    def get_queryset(self):
        return AlumniProfile.objects.filter(is_public=True)

class CvPersonalInformationViewSet(viewsets.GenericViewSet, mixins.ListModelMixin, mixins.RetrieveModelMixin):
    serializer_class = PersonalInformationSerializer
    permission_classes = [
        permissions.IsAuthenticated ]

    def get_queryset(self):
        alumni_id = self.kwargs.get('alumniId')
        return PersonalInformation.objects.filter(alumni=alumni_id)

class CvEducationInformationViewSet(viewsets.GenericViewSet, mixins.ListModelMixin, mixins.RetrieveModelMixin):
    serializer_class = EducationInformationSerializer
    permission_classes = [
        permissions.IsAuthenticated ]

    def get_queryset(self):
        alumni_id = self.kwargs.get('alumniId')
        return EducationInformation.objects.filter(alumni=alumni_id)

class CvExperienceInformationViewSet(viewsets.GenericViewSet, mixins.ListModelMixin, mixins.RetrieveModelMixin):
    serializer_class = ExperienceInformationSerializer
    permission_classes = [
        permissions.IsAuthenticated ]

    def get_queryset(self):
        alumni_id = self.kwargs.get('alumniId')
        return ExperienceInformation.objects.filter(alumni=alumni_id)

class AlumniCertificatesViewSet(viewsets.GenericViewSet, mixins.ListModelMixin, mixins.RetrieveModelMixin):
    serializer_class = CertificateSerializer
    permission_classes = [
        permissions.IsAuthenticated ]

    def get_queryset(self):
        alumni_id = self.kwargs.get('alumniId')
        return Certificate.objects.filter(alumni=alumni_id)

class SingleAlumniProfileViewSet(viewsets.GenericViewSet, mixins.ListModelMixin, mixins.RetrieveModelMixin):
    serializer_class = AlumniProfileSerializer
    permission_classes = [
        permissions.IsAuthenticated ]

    def get_queryset(self):
        alumni_id = self.kwargs.get('alumniId')
        return AlumniProfile.objects.filter(alumni=alumni_id)

class SingleStaffProfileViewSet(viewsets.GenericViewSet, mixins.ListModelMixin, mixins.RetrieveModelMixin):
    serializer_class = StaffProfileSerializer
    permission_classes = [
        permissions.IsAuthenticated ]

    def get_queryset(self):
        staff_id = self.kwargs.get('staffId')
        return StaffProfile.objects.filter(staff=staff_id)

class OrganizationContractsViewSet(viewsets.GenericViewSet, mixins.ListModelMixin, mixins.RetrieveModelMixin):
    serializer_class = ContractSerializer
    permission_classes = [
        permissions.IsAuthenticated ]

    def get_queryset(self):
        organization_id = self.kwargs.get('organizationId')
        return Contract.objects.filter(organization=organization_id)

class AlumniByOrganizationViewSet(viewsets.GenericViewSet, mixins.ListModelMixin, mixins.RetrieveModelMixin):
    serializer_class = AlumniProfileSerializer
    permission_classes = [
        permissions.IsAuthenticated ]

    def get_queryset(self):
        organization_id = self.kwargs.get('organizationId')
        return AlumniProfile.objects.filter(organization=organization_id)

class AlumniByCompletionYearViewSet(viewsets.GenericViewSet, mixins.ListModelMixin, mixins.RetrieveModelMixin):
    serializer_class = AlumniProfileSerializer
    permission_classes = [
        permissions.IsAuthenticated ]

    def get_queryset(self):
        year = self.kwargs.get('year')
        return AlumniProfile.objects.filter(completion_year=year)

class AlumniByDegreeProgramViewSet(viewsets.GenericViewSet, mixins.ListModelMixin, mixins.RetrieveModelMixin):
    serializer_class = AlumniProfileSerializer
    permission_classes = [
        permissions.IsAuthenticated ]

    def get_queryset(self):
        program_id = self.kwargs.get('programId')
        return AlumniProfile.objects.filter(program=program_id)

class AlumniByDepartmentViewSet(viewsets.GenericViewSet, mixins.ListModelMixin, mixins.RetrieveModelMixin):
    serializer_class = AlumniProfileSerializer
    permission_classes = [
        permissions.IsAuthenticated ]

    def get_queryset(self):
        department_id = self.kwargs.get('departmentId')
        return AlumniProfile.objects.filter(department=department_id)

class FieldPostByOrganizationViewSet(viewsets.GenericViewSet, mixins.ListModelMixin, mixins.RetrieveModelMixin):
    serializer_class = FieldPostSerializer
    permission_classes = [
        permissions.IsAuthenticated ]

    def get_queryset(self):
        organization_id = self.kwargs.get('organizationId')
        return FieldPost.objects.filter(organization=organization_id)

class FieldProfessionByPostViewSet(viewsets.GenericViewSet, mixins.ListModelMixin, mixins.RetrieveModelMixin):
    serializer_class = FieldPostProfessionSerializer
    permission_classes = [
        permissions.IsAuthenticated ]

    def get_queryset(self):
        post_id = self.kwargs.get('postId')
        return FieldPostProfession.objects.filter(post=post_id)

class FieldProgramByPostViewSet(viewsets.GenericViewSet, mixins.ListModelMixin, mixins.RetrieveModelMixin):
    serializer_class = FieldPostProgramSerializer
    permission_classes = [
        permissions.IsAuthenticated ]

    def get_queryset(self):
        post_id = self.kwargs.get('postId')
        return FieldPostProgram.objects.filter(post=post_id)

class FieldApplicationByPostViewSet(viewsets.GenericViewSet, mixins.ListModelMixin, mixins.RetrieveModelMixin):
    serializer_class = FieldApplicationSerializer
    permission_classes = [
        permissions.IsAuthenticated ]

    def get_queryset(self):
        post_id = self.kwargs.get('postId')
        return FieldApplication.objects.filter(post=post_id)

class FieldApplicationByStudentViewSet(viewsets.GenericViewSet, mixins.ListModelMixin, mixins.RetrieveModelMixin):
    serializer_class = FieldApplicationSerializer
    permission_classes = [
        permissions.IsAuthenticated ]

    def get_queryset(self):
        student_id = self.kwargs.get('studentId')
        return FieldApplication.objects.filter(student=student_id)

class InternshipPostByOrganizationViewSet(viewsets.GenericViewSet, mixins.ListModelMixin, mixins.RetrieveModelMixin):
    serializer_class = InternshipPostSerializer
    permission_classes = [
        permissions.IsAuthenticated ]

    def get_queryset(self):
        organization_id = self.kwargs.get('organizationId')
        return InternshipPost.objects.filter(organization=organization_id)

class InternshipApplicationByOrganizationViewSet(viewsets.GenericViewSet, mixins.ListModelMixin, mixins.RetrieveModelMixin):
    serializer_class = InternshipApplicationSerializer
    permission_classes = [
        permissions.IsAuthenticated ]

    def get_queryset(self):
        organization_id = self.kwargs.get('organizationId')
        return InternshipApplication.objects.filter(organization=organization_id, status="accepted")

class InternshipApplicationByPostViewSet(viewsets.GenericViewSet, mixins.ListModelMixin, mixins.RetrieveModelMixin):
    serializer_class = InternshipApplicationSerializer
    permission_classes = [
        permissions.IsAuthenticated ]

    def get_queryset(self):
        post_id = self.kwargs.get('postId')
        return InternshipApplication.objects.filter(post=post_id)

class InternshipApplicationByAlumniViewSet(viewsets.GenericViewSet, mixins.ListModelMixin, mixins.RetrieveModelMixin):
    serializer_class = InternshipApplicationSerializer
    permission_classes = [
        permissions.IsAuthenticated ]

    def get_queryset(self):
        alumni_id = self.kwargs.get('alumniId')
        return InternshipApplication.objects.filter(alumni=alumni_id)

class AlumniByStatusViewSet(viewsets.GenericViewSet, mixins.ListModelMixin, mixins.RetrieveModelMixin):
    serializer_class = AlumniProfileSerializer
    permission_classes = [
        permissions.IsAuthenticated ]

    def get_queryset(self):
        status = self.kwargs.get('status')
        return AlumniProfile.objects.filter(is_active=status)

class AlumniInvitationsViewSet(viewsets.GenericViewSet, mixins.ListModelMixin, mixins.RetrieveModelMixin):
    serializer_class = InvitationsSerializer
    permission_classes = [
        permissions.IsAuthenticated ]

    def get_queryset(self):
        alumni_id = self.kwargs.get('alumniId')
        return Invitations.objects.filter(alumni=alumni_id)

class OrganizationInvitationsViewSet(viewsets.GenericViewSet, mixins.ListModelMixin, mixins.RetrieveModelMixin):
    serializer_class = InvitationsSerializer
    permission_classes = [
        permissions.IsAuthenticated ]

    def get_queryset(self):
        organization_id = self.kwargs.get('organizationId')
        return Invitations.objects.filter(organization=organization_id)


# -*- coding: utf-8 -*-
"""WordMaker
"""

import json
import json_create

# Opening our TF-IDF sorted files:

# NOUNS:
n = open('nouns.txt', 'r').readlines()

# ADJECTIVES:
a = open('adjf.txt', 'r').readlines()

# Creating json files for each txt file using our main function:

# NOUNS:
json_create.main(n, 'nouns.json')

# ADJECTIVES:
json_create.main(a, 'adjf.json')

# Opening json files:
f = open('nouns.json')
h = open('adjf.json')
json_file1 = json.load(f)
json_file2 = json.load(h)

k = 0
while k != 10:
 emp = json_create.WordMaker(json_file1)
 c = emp.make_a_word(json_create.WordTypes.NOUN, 500)
 if c in n:
     c = c.replace(c, " ")
 emp = json_create.WordMaker(json_file2)
 b = emp.make_a_word(json_create.WordTypes.ADJECTIVE, 500)
 if b in n:
     b = b.replace(b, " ")
 print(b + " " + c)
 k += 1


import datetime
import os
import yaml
from flask import request
from flask_restful import Api, Resource
from sqlalchemy import exc
import sqlalchemy.exc as sqlalchemy_exc
from config import db, basedir
from models.searched_scripts import SearchedScripts
from schemas.searchedscripts_schema import *
import yfinance as yf

class GetAddSearchedScripts(Resource):
    def __init__(self):
        pass
    # SearchedScripts get call

    def get(self):
        try:
            obj = db.session.query(SearchedScripts).order_by(SearchedScripts.id).all()
            if obj:
                schema = GetSearchedScriptsSchema(many=True)
                data = schema.dump(obj).data
                response = {"success": True, "message":"data fetched successfully" ,"data": data}
                return (response)
            else:
                response = {"success": False,
                            "message": "no data found in SearchedScripts"}
                return(response)
        except Exception as e:
            return({"success": False, "message": "server is down try after sometime", "error": str(e.args)})

    # SearchedScripts post call
    def post(self):
        try:
            request_json_data = request.get_json()
            msft = yf.Ticker(request_json_data['stockscript'])
            # print(msft)
            output = msft.history(period="max")
            output.index = output.index.map(str)

            # print(data,type(data))
            output = output.to_dict('index')

            if len(output)==0:
                output = {"success":False}
            else:
                output["success"]=True
            request_json_data["stockscriptoutput"]=output
            print(request_json_data)
            schema = AddSearchedScriptsSchema()
            obj = schema.load(
                request_json_data, session=db.session).data
            db.session.add(obj)
            db.session.commit()
            data = schema.dump(obj).data
            role_response = {"success": True, "message": "data added successfully","data": data}
            return (role_response)
        except KeyError as e :
            return({"success": False, "message": "Incorrect key name", "error":str(e.args)})
        except exc.IntegrityError as e:
            errorInfo = e.orig.args
            message=errorInfo[1]
            return({"success": False, "message": message})
        except sqlalchemy_exc.SQLAlchemyError as e:
            return({"success": False, "message":str(e.args)})
        except Exception as e:
            return({"success": False, "message": str(e.args)})

class GetUserSearchedScripts(Resource):
    def __init__(self):
        pass
    # SearchedScripts get call

    def get(self,id):
        try:
            obj = SearchedScripts.query.filter(SearchedScripts.userId == id).order_by(SearchedScripts.id.desc()).all()
            # obj = db.session.query(SearchedScripts).order_by(SearchedScripts.id).all()
            if obj:
                schema = GetUserSearchedScriptsSchema(many=True)
                data = schema.dump(obj).data
                response = {"success": True, "message":"data fetched successfully" ,"data": data}
                return (response)
            else:
                response = {"success": False,
                            "message": "no data found in SearchedScripts"}
                return(response)
        except Exception as e:
            return({"success": False, "message": "server is down try after sometime", "error": str(e.args)})


# IP定位
key_pan = 'your_key' #在这里输入你的key值
def ip(keywords:str)->dict:
    url='https://restapi.amap.com/v3/ip?parameters'
    param={
        'key':key_pan,
        'output':'json'
    }
    response = requests.get(url,params=param)
    data=response.json()
    return data


#!/usr/bin/env python3
"""Task 6 Regularization"""
import tensorflow.compat.v1 as tf

def dropout_create_layer(prev, n, activation, keep_prob):
    """Creates a layer of a neural network using dropout"""
    weights_init = tf.keras.initializers.VarianceScaling(scale=2.0,
                                                         mode=("fan_avg"))
    layer = tf.layers.Dense(
        units=n,
        activation=activation,
        name="layer",
        kernel_initializer=weights_init
    )
    dropout = tf.layers.Dropout(keep_prob)

    return dropout(layer(prev))


from datetime import datetime

from solar import CommonSearcher
from solar.document import Document

from celery.events.state import Worker as _Worker, Task
from celery.utils.timeutils import maybe_iso8601

from .app import app

SEP = ':'

def try_fromtimestamp(timestamp):
    if timestamp:
        return datetime.fromtimestamp(timestamp)

def get_project_and_name(worker, project=None):
    if worker:
        if SEP in worker.hostname:
            return worker.hostname.split(SEP, 1)
        return worker.hostname, project
    return None, project

def get_worker_id(project, name):
    return u'%s%s%s' % (project, SEP, name)

class Worker(_Worker):
    # increase expire_window cause 200 is too small
    # Solr doesnt have time to do commit
    # 400 means freq * 4
    # if freq == 5 then
    # worker evaluates as alive during 20 seconds after last heartbeat
    expire_window = 400

class WorkerDocument(Document):
    @property
    def instance(self):
        if not hasattr(self, '_instance'):
            self._instance = Worker(**self.to_solr())
            for timestamp in getattr(self, 'heartbeats', []):
                self._instance.on_heartbeat(timestamp=timestamp)
        return self._instance

class WorkerSearcher(CommonSearcher):
    type_value = 'Worker'
    document_cls = WorkerDocument

    @classmethod
    def generate_indexing_docs(cls, workers, project=None):
        docs = []
        for worker in workers:
            project, name = get_project_and_name(worker, project=project)
            if project not in app.config['CELERIUM_PROJECTS']:
                continue
            doc = dict(worker)
            doc['id'] = get_worker_id(project, name)
            doc['project'] = project
            doc['name'] = name
            if worker.heartbeats:
                doc['last_heartbeat'] = try_fromtimestamp(worker.heartbeats[-1])
            docs.append(doc)
        return docs

class TaskSearcher(CommonSearcher):
    type_value = 'Task'

    iso8601_fields = ('eta', 'expires')
    timestamp_fields = (
        'timestamp', 'sent', 'received', 'started', 'succeeded',
        'failed', 'retried', 'revoked', 'expires',
        )

    @classmethod
    def generate_indexing_docs(cls, tasks, project=None):
        docs = []
        for task in tasks:
            worker = task.worker
            project, worker_name = get_project_and_name(worker, project=project)
            if project not in app.config['CELERIUM_PROJECTS']:
                continue
            doc = dict(task)
            doc.pop('worker', None)
            doc['id'] = doc.pop('uuid')
            doc['project'] = project
            if doc['name'].find('.') != -1:
                doc['module'] = '.'.join(doc['name'].split('.')[:-1])
            for field in cls.timestamp_fields:
                doc[field] = try_fromtimestamp(doc[field])
            for field in cls.iso8601_fields:
                doc[field] = maybe_iso8601(doc[field])
            if worker:
                doc['worker_id'] = get_worker_id(project, worker_name)
                doc['worker_name'] = worker_name
            docs.append(doc)
        return docs

worker_searcher = WorkerSearcher(app.config['CELERIUM_SOLR_URL'])
task_searcher = TaskSearcher(app.config['CELERIUM_SOLR_URL'])


""" A neural chatbot using sequence to sequence model with
attentional decoder.

This is based on Google Translate Tensorflow model
https://github.com/tensorflow/models/blob/master/tutorials/rnn/translate/

Sequence to sequence model by Cho et al.(2014)

Created by Chip Huyen as the starter code for assignment 3,
class CS 20SI: "TensorFlow for Deep Learning Research"
cs20si.stanford.edu

This file contains the code to run the model.

See readme.md for instruction on how to run the starter code.

This implementation learns NUMBER SORTING via seq2seq. Number range: 0,1,2,3,4,5,EOS

https://papers.nips.cc/paper/5346-sequence-to-sequence-learning-with-neural-networks.pdf

See README.md to learn what this code has done!

Also SEE https://stackoverflow.com/questions/38241410/tensorflow-remember-lstm-state-for-next-batch-stateful-lstm
for special treatment for this code

Belief Tracker
"""

import numpy as np
import uuid

class Node:

    API_NODE = "category"
    NORMAL_NODE = "normal"
    PROPERTY_NODE = "property"

    RANGE = "range"
    KEY = "key"

    def __init__(self, slot, value, fields, node_type, id):
        '''
        For now only api_node has more than one field
        Args:
            slot: slot serves as edge
            fields: a map indicating fields that are required or not
            node_type: api_node, property_node, normal
        '''
        self.value = value  # value aka slot-value filling
        self.slot = slot  # slot aks slot value filling
        self.ins = set([slot])
        self.parent_node = []
        self.fields = fields  # is a dict, valued as prob.
        self.outs = set(fields.keys())
        self.field_type = dict() # RANGE, KEY,...
        self.children = dict()  # value to chidren nodes
        self.node_type = node_type
        self.slot_to_values_mapper = dict()
        self.value_to_slot_mapper = dict()
        self.fields_trans = dict()
        self.id = id  # globally uniq
        self.level = 0

    def get_node_by_value(self, value):
        return self.children[value]

    def set_node_slot_trans(self, field, trans):
        if field not in self.fields:
            raise ValueError(
                'translation field error: node must have field which is included in self.fields.'
                ' Fix the graph file before continue;', field, self.fields)
        self.fields_trans[field] = trans

    def get_node_slot_trans(self, field):
        if field in self.fields_trans:
            return self.fields_trans[field]
        return None

    def set_field_type(self, field, field_type):
        self.field_type[field] = field_type

    def get_field_type(self, field):
        if field not in self.field_type:
            return Node.KEY
        return self.field_type[field]

    def has_field(self, field):
        return field in self.fields

    def add_node(self, node, edge=None, ignore_field_conflict=False):
        """
        Args:
            node: child_node to append to the current node, there will be a lot other information brought in
            ignore_field_conflict:
            edge: str
        Raises:
            ValueError: child_node slot must be in self.fields.
        """
        node_value = node.value
        node_field = node.slot
        if not edge:
            edge = node.slot
        if node_field not in self.fields and not ignore_field_conflict:
            raise ValueError(
                'error: node must have field which is included in self.fields.'
                ' Fix the graph file before continue;', node_field, self.fields)
        # self.outs.add(node_field)
        self.children[node_value] = node
        node.ins.add(edge)
        node.level = self.level + 1
        node.parent_node.append(self)

        # supply other information
        if node_field not in self.slot_to_values_mapper:
            self.slot_to_values_mapper[node_field] = []
        self.slot_to_values_mapper[node_field].append(node_value)

        if node_value not in self.value_to_slot_mapper:
            self.value_to_slot_mapper[node_value] = node_field

    def has_ancestor_by_value(self, value):
        for node in self.parent_node:
            # print(self.value, '***', node.value, '---', value)
            if node.value == value:
                return True
            if node.has_ancestor_by_value(value):
                return True
        return False

    def is_api_node(self):
        return self.node_type == self.API_NODE

    def is_property_node(self):
        return self.node_type == self.PROPERTY_NODE

    def get_slot_by_value(self, value):
        if value not in self.value_to_slot_mapper:
            return None
        return self.value_to_slot_mapper[value]

    def gen_required_slot_fields(self):
        required_fields = []
        for field, prob in self.fields.items():
            rnd = np.random.uniform(0, 1)
            if rnd < prob:
                required_fields.append(field)
        return required_fields

    def has_child(self, value):
        """
        has direct child
        """
        return value in self.children

    def get_child(self, value):
        return self.children[value]

    def has_posterity(self, value, belief_graph):
        """
        We use belief_graph to achieve faster bottom to up search
        """
        candidate_nodes = belief_graph.get_nodes_by_value(value)
        if self.value == "ROOT".lower():
            return len(candidate_nodes) > 0

        for node in candidate_nodes:
            id = node.id
            a = node
            while a.value != "ROOT".lower():
                if a.value == self.value:
                    return True
                a = a.parent_node
        return False

    def get_posterity_nodes_by_value(self, value, belief_graph):
        candidate_nodes = belief_graph.get_nodes_by_value(value)
        if self.value == "ROOT".lower():
            return candidate_nodes

        posterity = []
        for node in candidate_nodes:
            id = node.id
            a = node
            while a.value != "ROOT".lower():
                if a.value == self.value:
                    posterity.append(belief_graph.get_node_by_id(id))
                    break
                a = a.parent_node
        return posterity

    # def get_siblings(self, keep_slot):
    #     """
    #     This function might be dangerous to be used
    #     Args:
    #         keep_slot: True or False, return only siblings with same slot or not
    #     """
    #     if not self.parent_node:
    #         return []
    #
    #     siblings = []
    #     children = self.parent_node.chidren
    #     for key, value in children.items():
    #         if key != self.value:
    #             sibling = value
    #             if sibling.slot == self.slot or keep_slot == False:
    #                 siblings.append(sibling)
    #     return siblings

    def get_ancestry_values(self):
        anc_values = set()
        for node in self.parent_node:
            while node.value != "ROOT".lower():
                anc_values.add(node.value)
                anc_values.add(node.get_ancestry_values())
        return anc_values

    # def get_sibling_names(self, value, keep_slot):
    #     siblings = self.get_siblings(keep_slot)
    #     sibling_names = []
    #     for node in siblings:
    #         sibling_names.append(node.value)
    #     return sibling_names

    # def get_parent_node(self):
    #     return self.parent_node

    def get_children_names_by_slot(self, slot):
        children_names = []
        for key, value in self.children.items():
            if self.value_to_slot_mapper[key] == slot:
                children_names.append(key)
        return children_names

    def get_children_names(self, max_num=-1, required_only=False):
        children_names = []
        for key, value in self.children.items():
            if required_only == False or self.fields[key] == 1:
                children_names.append(key)
        children_names = np.array(children_names)
        if max_num == -1:
            return children_names
        children_names = np.random.choice(children_names, max_num)
        return children_names

    def is_leaf(self):
        return len(self.children) == 0


def insertAtBottom(stack, item):
	if isEmpty(stack):
		push(stack, item)
	else:
		temp = pop(stack)
		insertAtBottom(stack, item)
		push(stack, temp)

def reverse(stack):
	if not isEmpty(stack):
		temp = pop(stack)
		reverse(stack)
		insertAtBottom(stack, temp)

def createStack():
	stack = []
	return stack

# Function to check if
# the stack is empty
def isEmpty( stack ):
	return len(stack) == 0

# Function to push an
# item to stack
def push( stack, item ):
	stack.append( item )

# Function to pop an
# item from stack
def pop( stack ):

	# If stack is empty
	# then error
	if(isEmpty( stack )):
		print("Stack Underflow ")
		exit(1)

	return stack.pop()

# Function to print the stack
def prints(stack):
	for i in range(len(stack)-1, -1, -1):
		print(stack[i], end = ' ')
	print()

# Driver Code

stack = createStack()
push( stack, str(4) )
push( stack, str(3) )
push( stack, str(2) )
push( stack, str(1) )
print("Original Stack ")
prints(stack)

reverse(stack)

print("Reversed Stack ")
prints(stack)


from typing import Optional
from logger import logger
from pymongo import MongoClient
from pymongo.database import Database
from pymongo.collection import Collection
import pymongo

USE_DB = 'COV2K'
LOCAL_DB_CONNECTION_STRING = "mongodb://localhost:27017"
GECO_DB_CONNECTION_STRING = "mongodb://localhost:23456/gcm_gisaid"
MONGO_DB_CONNECTION_URI = LOCAL_DB_CONNECTION_STRING if USE_DB == 'COV2K' else GECO_DB_CONNECTION_STRING
DB_NAME = 'cov2k_v21_11_30' if USE_DB == 'COV2K' else 'gcm_gisaid'

_client: Optional[MongoClient] = None

def close_conn():
    global _client
    if _client:
        _client.close()
        logger.info(f'Connection with mongoDB {MONGO_DB_CONNECTION_URI} CLOSED')
        _client = None
    else:
        logger.warning("Request to close connection that was never opened.")

def open_conn() -> Database:
    global _client
    if not _client:
        _client = MongoClient(MONGO_DB_CONNECTION_URI)
        logger.info(f"Connection with mongoDB {MONGO_DB_CONNECTION_URI} ESTABLISHED")
    return _client[DB_NAME]

if __name__ == "__main__":
    pass
    # dbname = get_database()

    # get_variants().insert({
    #     "pangolin-id": "B.1.1.7"
    # })
    # for item in get_variants().find():
    #     print(item)