import enum

import geomdl
import math
import numpy as np
import pybullet as p
import scipy.optimize
from sklearn.decomposition import PCA
import torch
import pytorch_kinematics as pk

from scipy.spatial.transform import Rotation as R

def fibonacci_sphere(samples=1):
    points = []
    phi = math.pi * (3. - math.sqrt(5.))  # golden angle in radians

    for i in range(samples):
        y = 1 - (i / float(samples - 1)) * 2  # y goes from 1 to -1
        radius = math.sqrt(1 - y * y)  # radius at y

        theta = phi * i  # golden angle increment

        x = math.cos(theta) * radius
        z = math.sin(theta) * radius

        points.append((x, y, z))

    return np.asarray(points)

def construct_T_from_position_quaterion(translation, quaternion):
    """
    Given translation t and rotation q, construct
    T = [R(q), t; 0, 1]
    :param translation:
    :param quaternion:
    :return:
    """
    T = np.zeros([4, 4])
    T[:3, -1] = translation
    T[:3, :3] = np.array(p.getMatrixFromQuaternion(quaternion)).reshape([3, 3])
    T[-1, -1] = 1
    return T

def construct_T_from_position_matrix(translation, matrix):
    if type(translation) == torch.Tensor:
        T = torch.zeros([4, 4]).type_as(translation)
        T[:3, -1] = translation
        T[:3, :3] = matrix
        T[-1, -1] = torch.tensor(1).type_as(translation)
        T = torch.clone(T)
    else:
        T = np.zeros([4, 4])
        T[:3, -1] = translation
        T[:3, :3] = matrix
        T[-1, -1] = 1
    return T

def compute_joint_T(joint_angle, joint_axis):
    """ This function is used to compute the transformation matrix for a
    rotation operation. You are supposed to make use of the self._theta
    parameter and the sin and cos functions which have already been imported.
    You will have 3 cases depending on which axis the rotation is about, which
    you can obtain from either self.axis or self.axis_index.

    Returns:
      transform: The 4x4 transformation matrix.
    """
    assert(len(joint_axis)==3)
    nonzero_axis = np.nonzero(joint_axis)[0]
    if len(nonzero_axis) == 0:
        # Fixed joint
        if type(joint_angle) == torch.Tensor:
            return torch.eye(4)
        else:
            return np.eye(4)
    assert len(nonzero_axis) == 1
    # Flip the joint angle if joint axis is negative
    if type(joint_angle) == torch.Tensor:
        joint_angle = torch.clone(joint_angle) * joint_axis[nonzero_axis[0]]
        cth, sth = torch.cos(joint_angle), torch.sin(joint_angle)
        tensor_0 = torch.zeros_like(joint_angle)
        tensor_1 = torch.ones_like(joint_angle)
        if nonzero_axis[0] == 0:
            R = torch.stack([
                torch.stack([tensor_1, tensor_0, tensor_0]),
                torch.stack([tensor_0, cth, -sth]),
                torch.stack([tensor_0, sth, cth])]).reshape(3, 3)
        elif nonzero_axis[0] == 1:
            R = torch.stack([
                torch.stack([cth, tensor_0, sth]),
                torch.stack([tensor_0, tensor_1, tensor_0]),
                torch.stack([-sth, tensor_0, cth])]).reshape(3, 3)
        elif nonzero_axis[0] == 2:
            R = torch.stack([
                torch.stack([cth, -sth, tensor_0]),
                torch.stack([sth, cth, tensor_0]),
                torch.stack([tensor_0, tensor_0, tensor_1])]).reshape(3, 3)
        else:
            raise AssertionError
        return torch.block_diag(R, tensor_1).type_as(joint_angle)
    else:
        joint_angle = np.copy(joint_angle) * joint_axis[nonzero_axis[0]]
        cth, sth = np.cos(joint_angle), np.sin(joint_angle)
        H = np.zeros((4, 4))
        H[3, 3] = 1.
        # H[:3, :3] = R.from_euler(self.axis, self._theta).as_matrix()
        if nonzero_axis[0] == 0:
            # axis is x
            H[1, 1] = cth
            H[2, 2] = cth
            H[1, 2] = -sth
            H[2, 1] = sth
            H[0, 0] = 1.
        elif nonzero_axis[0] == 1:
            # axis is y
            H[0, 0] = cth
            H[2, 2] = cth
            H[2, 0] = -sth
            H[0, 2] = sth
            H[1, 1] = 1.
        elif nonzero_axis[0] == 2:
            # axis is z
            H[0, 0] = cth
            H[1, 1] = cth
            H[0, 1] = -sth
            H[1, 0] = sth
            H[2, 2] = 1.
        else:
            raise AssertionError
    return H

def convert_6d_rotation_to_flattened_matrix_torch(orn_6d_vec):
    """
    :param orn_6d_vec: 1D or 2D torch tensor whose last dimension is 6
    :return:
    """
    # Always unsqueeze to 2d
    orn_6d_vec_2d = torch.atleast_2d(orn_6d_vec).reshape(-1,6)
    v1 = orn_6d_vec_2d[..., :3]
    v2 = orn_6d_vec_2d[..., 3:]
    e1 = v1*(torch.torch.nan_to_num(1./torch.norm(v1, dim=-1)).unsqueeze(1))
    u2 = v2-e1*(torch.sum(torch.mul(e1, v2),dim=1).unsqueeze(1))
    # u2 = v2-torch.inner(e1, v2)*e1
    e2 = u2 * (torch.torch.nan_to_num(1./torch.norm(u2, dim=-1)).unsqueeze(1))  # Nx3
    base_rotation_matrix = torch.transpose(
        torch.stack([e1, e2, torch.cross(e1, e2, dim=1)], e1.dim()-1), dim0=-2, dim1=-1)
    return base_rotation_matrix.reshape(-1,9)

def get_basis_from_3_normals(normals):
    '''
    Given three 6D normals vectors with positions of shape 3x(3+3), return a
    matrix with shape 3x4x3. The 4 dim1s are p, n, t1, t2, where p is the position
    component in the input, [n,t1,t2] is an orthonormal basis in 3D space.
    n is the normals component in the input, t1 is the basis that is parallel to
    the plane defined by the 3 p's, n is the basis perpendicular to the 3 p's
    :param normals:
    :return:
    '''
    assert normals.shape == (3, 6)
    if type(normals) == torch.Tensor:
        device = normals.device
        t1 = torch.vstack([normals[:, 4],
                           -normals[:, 3],
                           torch.zeros(3).to(device)]).transpose(0, 1).type_as(normals)
        for i in range(3):
            if torch.allclose(t1[i, :], torch.zeros(3).type_as(t1)):
                t1[i, 0] += 1

        t2 = torch.cross(normals[:, 3:], t1, dim=1)
        t1 = (t1.transpose(0, 1)/torch.linalg.norm(t1, dim=1)).transpose(0, 1)
        t2 = (t2.transpose(0, 1)/torch.linalg.norm(t2, dim=1)).transpose(0, 1)

        ans = torch.stack([normals[:, :3].reshape(3, 1, 3),
                           normals[:, 3:].reshape(3, 1, 3),
                           t1.reshape(3, 1, 3),
                           t2.reshape(3, 1, 3)
                           ], dim=1).reshape([3, 4, 3])
        return ans
    else:
        raise NotImplementedError

def get_cross_matrix(p):
    if isinstance(p, torch.Tensor):
        return torch.tensor([[0., -p[2], p[1]],
                             [p[2], 0., -p[0]],
                             [-p[1], p[0], 0.]], requires_grad=p.requires_grad)
    else:
        return np.array([[0., -p[2], p[1]],
                         [p[2], 0., -p[0]],
                         [-p[1], p[0], 0.]])

def reformulate_to_soft_QP(Q_0, p_0, G_0, h_0, A_0, b_0,
                           inequality_constraint_weight,
                           equality_constraint_weight
                           ):
    '''
    Transforms a QP of the form
    min zᵀQ₀z + p₀ᵀz s.t. A₀z=b₀, G₀z≤h₀
    into a "soft" QP where Az=b is penalized with equality_constraint_weight
    and Gz≤h is penalized with inequality_constraint_weight:
    min z̃ᵀQ̃z̃ + p̃ᵀz̃ s.t. z̃ ≥ 0
    Dropping the constant terms,
    the constraint A₀z=b₀ is transformed to quadratic penalty
    min (A₀z−b₀)ᵀ(A₀z-b₀) → min zᵀA₀ᵀA₀z−2b₀ᵀA₀z
    and the constraint Gz≤h is transformed to quadratic penalty with slack
    variables s
    min (G₀z-h₀+s)ᵀ(G₀z-h₀+s) → min zᵀG₀ᵀG₀z+2sᵀG₀z+sᵀs−2h₀ᵀG₀z−2h₀ᵀs
    Rewriting the optimization by defining
    z̃ = [z; s], len(s) = len(h₀)
    The original costs become
    Q̃₀ = block_diag([Q₀, 0ₛ])
    p̃₀ᵀ = [p₀ᵀ, 0]
    The equality constraint cost weights are
    Q̃ₑ = block_diag([A₀ᵀA₀, 0ₛ])*equality_constraint_weight
    p̃ₑᵀ = [−2b₀ᵀA₀, 0]*equality_constraint_weight
    The inequality constraint cost weights are
    Q̃ᵢₑ = [G₀ᵀG₀, G₀; G₀ᵀ, Iₛ]*inequality_constraint_weight
    p̃ᵢₑᵀ = [-2h₀ᵀG₀, -2h₀ᵀ]*inequality_constraint_weight
    The total optimization becomes
    Q̃=Q̃₀+Q̃ₑ+Q̃ᵢₑ
    p̃ᵀ=p̃₀ᵀ+p̃ₑᵀ+p̃ᵢₑᵀ
    Ã, b̃ = Variable(torch.Tensor())
    G̃ = -I
    h̃ = 0
    :return: The matrices defining the soft QP Q̃, p̃, G̃, h̃, Ã, b̃
    '''
    raise DeprecationWarning
    slack_count = len(h_0)
    zeros_vec_slack_count = torch.zeros(slack_count).type_as(Q_0)
    zeros_mat_slack_count = torch.zeros(
        [slack_count, slack_count]).type_as(Q_0)
    eye_slack_count = torch.eye(slack_count).type_as(Q_0)
    Q_0_tilde = torch.block_diag(Q_0, zeros_mat_slack_count)
    p_0_tilde = torch.hstack([p_0, zeros_vec_slack_count])
    Q_e_tilde = torch.block_diag(A_0.transpose(0, 1)@A_0,
                                 zeros_mat_slack_count) *\
        equality_constraint_weight
    p_e_tilde = torch.hstack([-2.*b_0@A_0,
                              zeros_vec_slack_count])*equality_constraint_weight

    Q_ie_tilde = torch.vstack(
        [torch.hstack([G_0.transpose(0, 1)@G_0, G_0.transpose(0, 1)]),
         torch.hstack([G_0, eye_slack_count])]) *\
        inequality_constraint_weight
    p_ie_tilde = torch.hstack([-2.*h_0@G_0, -2.*h_0]) * \
        inequality_constraint_weight
    Q_tilde = (Q_0_tilde+Q_e_tilde+Q_ie_tilde).type_as(Q_0)
    p_tilde = (p_0_tilde+p_e_tilde+p_ie_tilde).type_as(Q_0)
    G_tilde = -torch.eye(Q_tilde.shape[0]).type_as(Q_0)
    h_tilde = torch.zeros(Q_tilde.shape[0]).type_as(Q_0)
    A_tilde = torch.autograd.Variable(torch.Tensor()).type_as(Q_0)
    b_tilde = torch.autograd.Variable(torch.Tensor()).type_as(Q_0)
    return Q_tilde, p_tilde, G_tilde, h_tilde, A_tilde, b_tilde

def reformulate_eq_to_soft_QP(Q_0, p_0, G_0, h_0, A_0, b_0,
                              equality_constraint_weight
                              ):
    '''
    Transforms a QP of the form
    min zᵀQ₀z + p₀ᵀz s.t. A₀z=b₀, G₀z≤h₀
    into a "soft" QP where Az=b is penalized with equality_constraint_weight
    but Gz≤h is still enforced.
    min zᵀQ̃z + p̃ᵀz s.t. z ≥ 0, G₀z≤h₀
    Dropping the constant terms
    the constraint A₀z=b₀ is transformed to quadratic penalty
    min (A₀z−b₀)ᵀ(A₀z-b₀) → min zᵀA₀ᵀA₀z−2b₀ᵀA₀z
    The equality constraint cost weights are
    Q̃ₑ = block_diag([A₀ᵀA₀])*equality_constraint_weight
    p̃ₑᵀ = [−2b₀ᵀA₀]*equality_constraint_weight
    The total optimization becomes
    Q̃=Q₀+Q̃ₑ
    p̃ᵀ=p₀ᵀ+p̃ₑᵀ
    Ã, b̃ = Variable(torch.Tensor())
    G̃ = [G₀; -I]
    h̃ = [h₀; 0]
    :return: The matrices defining the soft QP Q̃, p̃, G̃, h̃, Ã, b̃
    '''
    equality_constraint_weight = equality_constraint_weight.to(A_0.device)
    # Synced all devices to match Q_0
    p_0, G_0, h_0, A_0, b_0, equality_constraint_weight = p_0.to(Q_0.device), G_0.to(Q_0.device), h_0.to(Q_0.device), A_0.to(Q_0.device), b_0.to(Q_0.device), equality_constraint_weight.to(Q_0.device)
    Q_e_tilde = A_0.transpose(0, 1)@A_0*equality_constraint_weight
    p_e_tilde = -torch.tensor(2.,device=Q_0.device)*b_0@A_0*equality_constraint_weight
    Q_tilde = (Q_0+Q_e_tilde).type_as(Q_0)
    p_tilde = (p_0+p_e_tilde).type_as(Q_0)
    G_tilde = torch.vstack(
        [G_0, -torch.eye(Q_0.shape[0]).type_as(Q_0)]).type_as(Q_0)
    h_tilde = torch.hstack(
        [h_0, torch.zeros(Q_0.shape[0]).type_as(Q_0)]).type_as(Q_0)
    A_tilde = torch.autograd.Variable(torch.Tensor()).type_as(Q_0)
    b_tilde = torch.autograd.Variable(torch.Tensor()).type_as(Q_0)
    return Q_tilde, p_tilde, G_tilde, h_tilde, A_tilde, b_tilde

def polynomial_function_deg3(data, a, b, c, d, e, f, g, h, i, j):
    """
    z = f(x,y) = ax³+by³+cx²y+dxy²+ex²+fy²+gxy+hx+iy+j
    :param data: nx2 data
    :return: nx1
    """
    x = data[:,0]
    y = data[:,1]
    return a*x**3+b*y**3+c*x**2*y+d*x*y**2+e*x**2+f*y**2+g*x*y+h*x+i*y+j

def polynomial_function_deg3_grad(data, a, b, c, d, e, f, g, h, i, j):
    """
    ∂f/∂x = 3ax²+2cxy+dy²+2ex+gy+h
    ∂f/∂y = 3by²+cx²+2dxy+2fy+gx+i
    :param data:
    :return: 2x1 numpy array [∂f/∂x, ∂f/∂y]ᵀ
    """
    x = data[:,0]
    y = data[:,1]
    return np.array([[3*a*x**2+2*c*x*y+d*y**2+2*e*x+g*y+h],
                     [3*b*y**2+c*x**2+2*d*x*y+2*f*y+g*x+i]])

def polynomial_function_deg1(data, a, b, c):
    """
    z = ax+by+c
    :param data:
    :return:
    """
    x = data[:,0]
    y = data[:,1]
    return a*x+b*y+c

def polynomial_function_deg1_grad(data, a, b, c):
    return np.array([[a], [b]])

class SurfaceFitFunction(enum.Enum):
    POLYNOMIAL_DEG3 = (polynomial_function_deg3, polynomial_function_deg3_grad)
    POLYNOMIAL_DEG1 = (polynomial_function_deg1, polynomial_function_deg1_grad)

def get_fitted_fn(fn_type, popt):
    fn, fn_grad = fn_type.value
    def fitted_f(x):
        return fn(x, *popt)
    def fitted_f_grad(x):
        return fn_grad(x, *popt)
    return fitted_f, fitted_f_grad

def compute_principal_components(points):
    """

    :param points:
    :return: 3-tuple
    ndarray of shape (n_components, n_features)
    ndarray of shape (n_features,)
    """
    pca = PCA(n_components=3)
    transformed_points = pca.fit_transform(points)
    return pca.components_, pca.mean_, transformed_points

def fit_surface_to_points(points,
                          fit_function=SurfaceFitFunction.POLYNOMIAL_DEG1,
                          **kwargs):
    """
    Given nx3 points, find the rotation matrix that provide.
    :param points:
    :param fit_function:
    :param kwargs:
    :return: (R, the best rotation matrix that fits the
    """
    assert points.shape[1] == 3
    return scipy.optimize.curve_fit(fit_function.value[0], points[:, :2], points[:, 2])

def transform_and_fit_surface_to_points(
        p_W,
        fit_function=SurfaceFitFunction.POLYNOMIAL_DEG3,
        **kwargs):
    """

    :param p_W:
    :param fit_function:
    :param kwargs:
    :return:
    (R, t, fitted_points, fitted_fn, transformed_fitted_fn)
    Given a point x, y, the predicted function is
    """
    # First do PCA
    components, mean, p_S = compute_principal_components(p_W)
    # Fit to the transformed points
    popt, pcov = fit_surface_to_points(p_S,
                          fit_function,**kwargs)
    # Reconstruct the function
    fitted_fn_S, fitted_fn_grad_S = get_fitted_fn(fit_function, popt)
    R_WS = components.T
    t_WS = mean
    return R_WS, t_WS, p_S, fitted_fn_S, fitted_fn_grad_S

def transform_and_fit_consistent_normal_surface_to_points(
        all_p_W,
        fit_function=SurfaceFitFunction.POLYNOMIAL_DEG3,
        **kwargs):
    """

    :param points: 3D array of shape [3, n, 3]
    :param fit_function:
    :param kwargs:
    :return:
    (R, t, fitted_points, fitted_fn, transformed_fitted_fn)
    Given a point x, y, the predicted function is
    """
    assert len(all_p_W) == 3
    for i in range(len(all_p_W)):
        assert all_p_W[i].shape[1] == 3
    all_components = np.zeros([3, 3, 3])
    all_means = np.zeros([3, 3])
    for idx, points in enumerate(all_p_W):
        # First do PCA
        components, mean, _ = compute_principal_components(points)
        all_components[idx, :, :] = components
        all_means[idx,:] = mean
    all_pavg_W = np.average(all_means, axis=0)
    ans = []
    for idx, points in enumerate(all_p_W):
        # Flip the signs of the basis if necessary
        sign = int(np.dot(all_means[idx,:]-all_pavg_W, all_components[idx, 2, :]) < 0)
        all_components[idx] *= (-1)**sign
        # Make sure the normals are consistent
        # Fit to the transformed points
        # Reconstruct the function
        R_WS = all_components[idx].T
        t_WS = all_means[idx,:]
        p_S = (R_WS.T @ (all_p_W[idx]-t_WS).T).T
        popt, pcov = fit_surface_to_points(p_S,
                              fit_function, **kwargs)
        fitted_fn, fitted_fn_grad = get_fitted_fn(fit_function, popt)
        ans.append((R_WS, t_WS, p_S, fitted_fn, fitted_fn_grad))
    return ans

def compute_nF_S(p_S, fitted_fn_grad):
    """
    Compute the fitted surface normal of a point p_S in the PCA basis space
    See https://en.wikipedia.org/wiki/Normal_(geometry)#Calculating_a_surface_normal for math
    :param p_S:
    :param fitted_fn_grad:
    :return: nF_S
    """
    grad_p = np.squeeze(fitted_fn_grad(p_S[:2,:].reshape(1,-1)))
    nF_S = np.array([-grad_p[0], -grad_p[1], 1.])
    return nF_S / np.linalg.norm(nF_S)

def compute_n_W(p_W, R_WS, t_WS, fitted_fn_grad):
    p_S = np.matmul(R_WS.T, p_W.reshape(3,1) - t_WS.reshape(3,1))
    n_S = compute_nF_S(p_S, fitted_fn_grad)
    return R_WS @ n_S

def rotation_matrix_from_vectors(vec1, vec2):
    """
    Compute the rotation matrix R that can: Rvec1 = vec2,
    If flipped 180, bu default rotate along x axis.
    There may still be some edge cases that can cause problem..
    """
    a, b = (vec1 / np.linalg.norm(vec1)).reshape(3), (vec2 / np.linalg.norm(vec2)).reshape(3)
    v = np.cross(a, b)
    c = np.dot(a, b)
    s = np.linalg.norm(v)
    if s==0:
        if c > 0:
            return np.eye(3)
        else:
            return np.array([[-1,0,0],[0,1,0],[0,0,-1]])
    kmat = np.array([[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]])
    rotation_matrix = np.eye(3) + kmat + kmat.dot(kmat) * ((1 - c) / (s ** 2))
    return rotation_matrix

def minimumDist2dBox(p_client, rigid_box, shape, point):
    """
    axes aligned rect is a rectangle numpy array [4,3]
    if not we need to express the point in rect frame
    """
    pos, orn = p_client.getBasePositionAndOrientation(rigid_box)
    rel_pos = torch.from_numpy(point - np.asarray(pos))
    orn = torch.tensor([orn[3], orn[0], orn[1], orn[2]])
    inv_q = pk.quaternion_invert(orn)
    x,y,z = pk.quaternion_apply(inv_q, rel_pos).tolist()
    # Calculate minimum distance
    dx = max(abs(x)-shape[0], 0)
    dy = max(abs(y)-shape[1], 0)
    dz = max(abs(z)-shape[2], 0)
    return np.linalg.norm([dx, dy, dz])


# -*- coding: utf-8 -*-
"""
Created on Thu Jun  1 00:29:59 2023

@author: 16692
"""

from sqlalchemy import Column, Integer, String, DateTime, Text
from sqlalchemy.orm import declarative_base
from datetime import datetime

Base = declarative_base()
class User(Base):
    __tablename__ = "users"
    user_id = Column(Integer(), primary_key=True)
    name = Column(String(100), nullable=False)
    email = Column(String(100), nullable=False)
    signup_date = Column(DateTime(), default=datetime.now)

    def __init__(self,user_id,name,email,signup_date):
        self.user_id = user_id
        self.name = name
        self.email = email
        self.signup_date =signup_date

from connection import get_connection,get_engine
engine =get_engine()
try:
    Base.metadata.create_all(bind=engine)
    print("database creation done")
except Exception as e:
    print(e)


import time
import torch
import random
import numpy as np
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from tqdm import tqdm, trange
from torchtext import data
from sklearn.metrics import f1_score
from sklearn.metrics import confusion_matrix
from transformers import XLMRobertaModel
from models.MultiHead_BiRU import MultiheadBIRUAttention
from models.bert_cnn_model import BERTCNNSentiment
from transformers import XLMRobertaTokenizer

SEED = 6548
random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)
torch.backends.cudnn.deterministic = True
CUDA_LAUNCH_BLOCKING = 1

max_input_length = 64
BATCH_SIZE = 128
train_name = "training.csv"
test_name = "test.csv"
val_name = "validation.csv"
model_save_names = ['./checkpoint/modelA.txt', "./checkpoint/model2.txt" ]
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

tokenizer = XLMRobertaTokenizer.from_pretrained("xlm-roberta-base")
print('XLM Roberta Tokenizer Loaded...')
init_token_idx = tokenizer.cls_token_id
eos_token_idx = tokenizer.sep_token_id
pad_token_idx = tokenizer.pad_token_id
unk_token_idx = tokenizer.unk_token_id
print("Max input length: %d" %(max_input_length))

def tokenize_and_cut(sentence):
    tokens = tokenizer.tokenize(sentence)
    tokens = tokens[:max_input_length-2]
    return tokens

UID=data.Field(
    sequential=False,
    use_vocab=False,
    pad_token=None
    )

TEXT=data.Field(batch_first=True,
                  use_vocab=False,
                  tokenize=tokenize_and_cut,
                  preprocessing=tokenizer.convert_tokens_to_ids,
                  init_token=init_token_idx,
                  eos_token=eos_token_idx,
                  pad_token=pad_token_idx,
                  unk_token=unk_token_idx
                  )
LABEL=data.LabelField()

fields = [('id',UID),('text', TEXT),('category', LABEL)]
train_data, valid_data, test_data = data.TabularDataset.splits(
                                        path='./',
                                        train=train_name,
                                        test=test_name,
                                        validation=val_name,
                                        format='csv',
                                        fields=fields,
                                        skip_header=True)

print('Data loading complete')
print(f"Number of training examples: {len(train_data)}")
print(f"Number of validation examples: {len(valid_data)}")
print(f"Number of test examples: {len(test_data)}")

LABEL.build_vocab(train_data, valid_data)
train_iterator, valid_iterator, test_iterator=data.BucketIterator.splits(
    (train_data, valid_data, test_data),
    sort_key=lambda x: len(x.text),
    batch_size = BATCH_SIZE,
    device = device,
    )
bert = XLMRobertaModel.from_pretrained('xlm-roberta-base')

OUTPUT_DIM = 5
DROPOUT = 0.3
N_FILTERS = 100
FILTER_SIZES = [2, 3, 5, 7, 9]

model_names = ["A", "B"]

models = [
    BERTCNNSentiment(bert, OUTPUT_DIM, DROPOUT, N_FILTERS, FILTER_SIZES),
    MultiheadBIRUAttention(bert, 128, 100, 100, 4, 768, 5)
]

def clip_gradient(model, clip_value):
    params = list(filter(lambda p: p.grad is not None, model.parameters()))
    for p in params:
        p.grad.data.clamp_(-clip_value, clip_value)

optimizers = [optim.Adam(models[0].parameters()), optim.Adam(models[1].parameters())]
criterion = nn.CrossEntropyLoss()
nll_loss = nn.NLLLoss()
log_softmax = nn.LogSoftmax()

for i in range(2):
    models[i] = models[i].to(device)

criterion=criterion.to(device)
nll_loss=nll_loss.to(device)
log_softmax=log_softmax.to(device)

def categorical_accuracy(preds, y):
    device = y.device
    count0,count1,count2,count3,count4 = torch.zeros(1),torch.zeros(1),torch.zeros(1),torch.zeros(1),torch.zeros(1)
    count0,count1,count2,count3,count4 = count0.to(device) ,count1.to(device) ,count2.to(device) ,count3.to(device) ,count4.to(device)
    total0,total1,total2,total3,total4 = torch.FloatTensor(1),torch.FloatTensor(1),torch.FloatTensor(1),torch.FloatTensor(1),torch.FloatTensor(1)
    max_preds = preds.argmax(dim = 1, keepdim = True)
    correct = max_preds.squeeze(1).eq(y)
    predictions = max_preds.squeeze(1)
    true_correct = [0,0,0,0,0]
    for j,i in enumerate(y.cpu().numpy()):
        true_correct[y.cpu().numpy()[j]]+=1
        if i==0:
            count0+=correct[j]
            total0+=1
        elif i==1:
            count1+=correct[j]
            total1+=1
        elif i==2:
            count2+=correct[j]
            total2+=1
        elif i==3:
            count3+=correct[j]
            total3+=1
        elif i==4:
            count4+=correct[j]
        else:
            total4+=1
    metric=torch.FloatTensor([count0/true_correct[0],count1/true_correct[1],count2/true_correct[2],count3/true_correct[3],count4/true_correct[4],f1_score(y.cpu().numpy(),predictions.cpu().numpy(),average='weighted')])
    metric[metric!=metric] = 0
    hi4 = correct.sum()
    hi5 = torch.FloatTensor([y.shape[0]]).to(device)
    hi1 =  hi4 / hi5
    hi2 = metric
    hi3 = confusion_matrix(y.cpu().numpy(),max_preds.cpu().numpy(),labels=[0,1,2,3,4])
    return hi1,hi2,hi3

def train(model, iterator, optimizer, criterion, i):
    epoch_loss = 0
    epoch_acc = 0
    model.train()
    for batch in iterator:
        optimizer.zero_grad()
        if (i == 0):
            predictions =  model(batch.text).squeeze(1)
        else:
            predictions =  model(batch.text, batch_size = len(batch)).squeeze(1)
        loss = criterion(predictions, batch.category)
        acc,_,_ = categorical_accuracy(predictions, batch.category)
        loss.backward()
        clip_gradient(model, 1e-1)
        optimizer.step()
        epoch_loss += loss.item()
        epoch_acc += acc.item()
    return epoch_loss / len(iterator), epoch_acc / len(iterator)

def evaluate(model, iterator, criterion, i):
    epoch_loss = 0
    epoch_acc = 0
    epoch_all_acc = torch.FloatTensor([0,0,0,0,0,0])
    confusion_mat = torch.zeros((5,5))
    confusion_mat_temp = torch.zeros((5,5))
    model.eval()
    with torch.no_grad():
        for batch in iterator:
            if (i == 0):
                predictions = model(batch.text).squeeze(1)
            else:
                predictions = model(batch.text,batch_size=len(batch)).squeeze(1)
            loss = criterion(predictions, batch.category)
            acc,all_acc,confusion_mat_temp = categorical_accuracy(predictions, batch.category)
            epoch_loss += loss.item()
            epoch_acc += acc.item()
            epoch_all_acc += all_acc
            confusion_mat+=confusion_mat_temp
    return epoch_loss / len(iterator), epoch_acc / len(iterator),epoch_all_acc/len(iterator),confusion_mat

def epoch_time(start_time, end_time):
    elapsed_time = end_time - start_time
    elapsed_mins = int(elapsed_time / 60)
    elapsed_secs = int(elapsed_time - (elapsed_mins * 60))
    return elapsed_mins, elapsed_secs

N_EPOCHS = 60
best_f1 = [-1, -1]
t = trange(N_EPOCHS)

for epoch in t:
    t.set_description('EPOCH %i' % epoch)
    for i in range(2):
        print(model_names[i])
        start_time = time.time()
        train_loss, train_acc = train(models[i], train_iterator, optimizers[i], criterion, i)
        valid_loss, valid_acc,tot,conf = evaluate(models[i], valid_iterator, criterion, i)
        f1 = tot[5]
        end_time = time.time()
        epoch_mins, epoch_secs = epoch_time(start_time, end_time)

        if f1 > best_f1[i]:
            best_f1[i] = f1
            path = model_save_names[i]
            print(path)
            torch.save(models[i].state_dict(), path)

        print(f'Epoch: {epoch+1:02} | Epoch Time: {epoch_mins}m {epoch_secs}s')
        print(f'\tTrain Loss: {train_loss:.3f} | Train Acc: {train_acc*100:.2f}%')
        print(f'\t Val. Loss: {valid_loss:.3f} |  Val. Acc: {valid_acc*100:.2f}%')
        print("Validation F1 : ", f1)


"""
This module defines input sets for CP2K and is a work in progress. The structure/philosophy
of this module is based on the Vasp input sets in Pymatgen. These sets are meant to contain
tested parameters that will result in successful, reproducible, consistent calculations without
need for intervention 99% of the time. 99% of the time, you only need to provide a pymatgen
structure object and let the defaults take over from there.

The sets are intended to be very general, e.g. a set for geometry relaxation, and so most of the
time, if you have specific needs, you can simply specify them via the keyword argument
override_default_params (see Section.update() method). If you have the need to create a new input
set (say for a standardized high throughput calculation) then you can create a new child of the
Cp2kInputSet class.

In order to implement a new Set within the current code structure, follow this 3 step flow:
    (1) Inherit from Cp2kInputSet or one of its children and call the super() constructor
    (2) Create the new sections and insert them into self and its subsections as needed
    (3) Call self.update(override_default_params) in order to allow user settings.
"""

import warnings
from pathlib import Path
from typing import Dict, Union

from pymatgen.core.periodic_table import Element
from pymatgen.core.lattice import Lattice
from pymatgen.core.structure import Molecule, Structure
from pymatgen.io.cp2k.inputs import (
    LDOS,
    PBE,
    PDOS,
    QS,
    XC_FUNCTIONAL,
    Cell,
    Coord,
    Cp2kInput,
    Dft,
    E_Density_Cube,
    ForceEval,
    Global,
    Keyword,
    KeywordList,
    Kind,
    Kpoints,
    Mgrid,
    MO_Cubes,
    OrbitalTransformation,
    Scf,
    Section,
    Smear,
    Subsys,
    V_Hartree_Cube,
)
from pymatgen.io.cp2k.utils import (
    get_aux_basis,
    get_basis_and_potential,
    get_unique_site_indices,
)

__author__ = "Nicholas Winner"
__version__ = "0.2"
__email__ = "nwinner@berkeley.edu"
__date__ = "January 2019"

MODULE_DIR = Path(__file__).resolve().parent

class Cp2kInputSet(Cp2kInput):

    """
    The basic representation of a CP2K input set as a collection of "sections" defining the simulation
    connected to a structure object. At the most basis level, CP2K requires a &GLOBAL section and
    &FORCE_EVAL section. Global sets parameters like "RUN_TYPE" or the overall verbosity. FORCE_EVAL is
    the largest section usually, containing the cell and coordinates of atoms, the DFT settings, and more.
    This top level input set is meant to initialize GLOBAL and FORCE_EVAL based on a structure object and
    and sections that the user provides.

    Like everything that goes into a cp2k input file, this base input set is essentially a section object.
    These sets are distinguished by saving default settings for easy implementation of calculations such
    as relaxation and static calculations. This base set is here to transfer a pymatgen structure object
    into the input format for cp2k and associate the basis set and pseudopotential to use with each
    element in the structure.

    Generally, this class will not be used directly, and instead one of
    its child-classes will be used, which contain more predefined initializations of various sections, and,
    if modifications are required, the user can specify override_default_settings.
    """

    def __init__(
        self,
        structure: Union[Structure, Molecule],
        potential_and_basis: Dict = {},
        multiplicity: int = 0,
        project_name: str = "CP2K",
        override_default_params: Dict = {},
        **kwargs,
    ):
        """
        Args:
            structure: (Structure or Molecule) pymatgen structure or molecule object used to define
                the lattice, coordinates, and elements. This structure object cannot contain "special"
                species like the Dummy species, e.g. X, or fractional occupations, e.g. Fe0.2, etc.
            potential_and_basis: (dict) Specifies what basis set and potential to use. Specify these
                as a dict of the form:
                    { element: {'cardinality': __, 'sr': __, 'q': __},
                     'cardinality': __, 'functional': __}
                Where cardinality and functional are overall specifications (for all elements), while
                <key='element'> specifies the overrides for a specific element. Currently the following
                conventions must be followed:
                    (a) All species of a particular element must have the same potential/basis
            multiplicity: (int) Specify the system's multiplicity if appropriate
            project_name: (str) Specify the project name. This will be used to name the output files
                from a CP2K calculation
            override_default_params: (dict) Specifies user-defined settings to override the settings of any
                input set (See Section.update())
        """
        super().__init__(name="CP2K_INPUT", subsections={})

        # Important CP2K set parameters
        self.structure = structure
        self.charge = structure.charge
        self.potential_and_basis = potential_and_basis
        self.multiplicity = multiplicity  # spin multiplicity = 2s+1
        self.override_default_params = override_default_params
        self.project_name = project_name
        self.kwargs = kwargs

        for s in self.structure.species:
            assert s in Element

        self.insert(ForceEval())  # always present in cp2k
        self.basis_set_file_names = None  # need for dft
        self.potential_file_name = None  # need for dft
        self.create_subsys(self.structure)  # assemble structure with atom types and pseudopotentials assigned

        if self.kwargs.get("print_forces", True):
            self.print_forces()
        if self.kwargs.get("print_motion", True):
            self.print_motion()

        self.update(override_default_params)

    def create_subsys(self, structure: Union[Structure, Molecule]):
        """
        Create the structure for the input
        """
        subsys = Subsys()
        if isinstance(structure, Structure):
            subsys.insert(Cell(structure.lattice))

        # Decide what basis sets/pseudopotentials to use
        basis_and_potential = get_basis_and_potential([str(s) for s in structure.species], self.potential_and_basis)

        # Insert atom kinds by identifying the unique sites (unique element and site properties)
        unique_kinds = get_unique_site_indices(structure)
        for k, v in unique_kinds.items():
            kind = k.split("_")[0]
            kwargs = {}
            if "magmom" in self.structure.site_properties:
                kwargs["magnetization"] = self.structure.site_properties["magmom"][v[0]]

            if "ghost" in self.structure.site_properties:
                kwargs["ghost"] = self.structure.site_properties["ghost"][v[0]]

            if "basis_set" in self.structure.site_properties:
                basis_set = self.structure.site_properties["basis_set"][v[0]]
            else:
                basis_set = basis_and_potential[kind]["basis"]

            if "potential" in self.structure.site_properties:
                potential = self.structure.site_properties["potential"][v[0]]
            else:
                potential = basis_and_potential[kind]["potential"]

            if "aux_basis" in self.structure.site_properties:
                kwargs["aux_basis"] = self.structure.site_properties["aux_basis"][v[0]]

            subsys.insert(Kind(kind, alias=k, basis_set=basis_set, potential=potential, **kwargs))
        coord = Coord(structure, aliases=unique_kinds)
        subsys.insert(coord)
        self["FORCE_EVAL"].insert(subsys)
        self.basis_set_file_names = basis_and_potential["basis_filenames"]
        self.potential_file_name = basis_and_potential["potential_filename"]

    def print_forces(self):
        """
        Print out the forces and stress during calculation
        """
        self["FORCE_EVAL"].insert(Section("PRINT", subsections={}))
        self["FORCE_EVAL"]["PRINT"].insert(Section("FORCES", subsections={}))
        self["FORCE_EVAL"]["PRINT"].insert(Section("STRESS_TENSOR", subsections={}))

    def print_motion(self):
        """
        Print the motion info (trajectory, cell, forces, stress
        """
        if not self.check("MOTION"):
            self.insert(Section("MOTION", subsections={}))
        self["MOTION"].insert(Section("PRINT", subsections={}))
        self["MOTION"]["PRINT"].insert(Section("TRAJECTORY", section_parameters=["ON"], subsections={}))
        self["MOTION"]["PRINT"].insert(Section("CELL", subsections={}))
        self["MOTION"]["PRINT"].insert(Section("FORCES", subsections={}))
        self["MOTION"]["PRINT"].insert(Section("STRESS", subsections={}))

class DftSet(Cp2kInputSet):
    """
    Base for an input set using the Quickstep module (i.e. a DFT calculation). The DFT section is pretty vast
    in CP2K, so this set hopes to make the DFT setup fairly simple. The provided parameters are pretty conservative,
    and so they should not need to be changed very often.
    """

    def __init__(
        self,
        structure: Union[Structure, Molecule],
        ot: bool = True,
        band_gap: float = 0.01,
        eps_default: float = 1e-12,
        eps_scf: float = 1e-7,
        max_scf: Union[int, None] = None,
        minimizer: str = "DIIS",
        preconditioner: str = "FULL_ALL",
        algorithm: str = "STRICT",
        linesearch: str = "2PNT",
        cutoff: int = 1200,
        rel_cutoff: int = 80,
        ngrids: int = 5,
        progression_factor: int = 3,
        override_default_params: Dict = {},
        wfn_restart_file_name: str = None,
        kpoints: Union[Kpoints, None] = None,
        smearing: bool = False,
        **kwargs,
    ):
        """
        Args:
            structure: Pymatgen structure or molecule object
            ot (bool): Whether or not to use orbital transformation method for matrix diagonalization. OT is the
                flagship scf solver of CP2K, and will provide huge speed-ups for this part of the calculation,
                but the system must have a band gap for OT to be used (higher band-gap --> faster convergence).
                Band gap is also used by the preconditioner for OT, and should be set as a value SMALLER than the true
                band gap to get good efficiency. Generally, this parameter does not need to be changed from
                default of 0.01
            band_gap (float): The band gap can also be specified in order to determine if ot should be turned on.
            eps_default (float): Replaces all EPS_XX Keywords in the DFT section (NOT its subsections!) to have this
                value, ensuring an overall accuracy of at least this much.
            eps_scf (float): The convergence criteria for leaving the SCF loop in Hartrees. Default is 1e-7. Should
                ensure reasonable results for all properties. Smaller than 1e-7 is generally not needed unless
                you need very high precision. 1e-6 may be used for difficult systems, and should still give
                reasonable results for most properties.
            max_scf (int): The max number of SCF cycles before terminating the solver. NOTE: With the OT solver, this
                corresponds to the max number of INNER scf loops, and then the outer loops are set with outer_max_scf,
                while with diagnolization it corresponds to the overall (INNER*OUTER) number of SCF steps, with the
                inner loop limit set by
            minimizer (str): The minimization scheme. DIIS can be as much as 50% faster than the more robust conjugate
                gradient method, and so it is chosen as default. Switch to CG if dealing with a difficult system.
            preconditioner (str): Preconditioner for the OT method. FULL_ALL is the most reliable, and is the
                default. Though FULL_SINGLE_INVERSE has faster convergence according to our internal tests. Should
                only change from theses two when simulation cell gets to be VERY large,
                in which case FULL_KINETIC might be preferred.
            cutoff (int): Cutoff energy (in Ry) for the finest level of the multigrid. A high cutoff will allow you to
                have very accurate calculations PROVIDED that REL_CUTOFF is appropriate.
            rel_cutoff (int): This cutoff decides how the Guassians are mapped onto the different levels of the
                multigrid. From CP2K: A Gaussian is mapped onto the coarsest level of the multi-grid, on which the
                    function will cover number of grid points greater than or equal to the number of grid points
                    will cover on a reference grid defined by REL_CUTOFF.
            progression_factor (int): Divisor of CUTOFF to get the cutoff for the next level of the multigrid.

            Takeaway for the cutoffs: https://www.cp2k.org/howto:converging_cutoff
            If CUTOFF is too low, then all grids will be coarse and the calculation may become inaccurate; and if
            REL_CUTOFF is too low, then even if you have a high CUTOFF, all Gaussians will be mapped onto the coarsest
            level of the multi-grid, and thus the effective integration grid for the calculation may still be too
            coarse.
        """

        super().__init__(structure, **kwargs)

        self.structure = structure
        self.ot = ot
        self.band_gap = band_gap
        self.eps_default = eps_default
        self.eps_scf = eps_scf
        self.max_scf = max_scf
        self.minimizer = minimizer
        self.preconditioner = preconditioner
        self.algorithm = algorithm
        self.linesearch = linesearch
        self.cutoff = cutoff
        self.rel_cutoff = rel_cutoff
        self.ngrids = ngrids
        self.progression_factor = progression_factor
        self.override_default_params = override_default_params
        self.wfn_restart_file_name = wfn_restart_file_name
        self.kpoints = kpoints
        self.smearing = smearing
        self.kwargs = kwargs

        # Build the QS Section
        qs = QS(eps_default=eps_default)
        max_scf = max_scf if max_scf else 20 if ot else 400  # If ot, max_scf is for inner loop
        scf = Scf(eps_scf=eps_scf, max_scf=max_scf, subsections={})

        # If there's a band gap, use OT, else use Davidson
        if ot:
            if band_gap <= 0:
                warnings.warn(
                    "Orbital Transformation method is being used for"
                    "a system without a bandgap. OT can have very poor"
                    "convergence for metallic systems, proceed with caution.",
                    UserWarning,
                )
            scf.insert(
                OrbitalTransformation(
                    minimizer=minimizer,
                    preconditioner=preconditioner,
                    energy_gap=band_gap,
                    algorithm=algorithm,
                    linesearch=linesearch,
                )
            )
            scf.insert(
                Section(
                    "OUTER_SCF",
                    subsections={},
                    keywords={
                        "MAX_SCF": Keyword("MAX_SCF", kwargs.get("outer_max_scf", 20)),
                        "EPS_SCF": Keyword("EPS_SCF", kwargs.get("outer_eps_scf", eps_scf)),
                    },
                )
            )
        else:
            scf.insert(Section("DIAGONALIZATION", subsections={}))
            mixing_kwds = {
                "METHOD": Keyword("METHOD", "BROYDEN_MIXING"),
                "ALPHA": Keyword("ALPHA", 0.2),
                "NBUFFER": Keyword("NBUFFER", 5),
            }
            mixing = Section("MIXING", keywords=mixing_kwds, subsections=None)
            scf.insert(mixing)
            davidson_kwds = {"PRECONDITIONER": Keyword("PRECONDITIONER", "FULL_ALL")}
            davidson = Section("DAVIDSON", keywords=davidson_kwds, subsections=None)
            scf["DIAGONALIZATION"].insert(davidson)

        # Create the multigrid for FFTs
        mgrid = Mgrid(
            cutoff=cutoff,
            rel_cutoff=rel_cutoff,
            ngrids=ngrids,
            progression_factor=progression_factor,
        )

        # Set the DFT calculation with global parameters
        dft = Dft(
            MULTIPLICITY=self.multiplicity,
            CHARGE=self.charge,
            basis_set_filenames=self.basis_set_file_names,
            potential_filename=self.potential_file_name,
            subsections={"QS": qs, "SCF": scf, "MGRID": mgrid},
            wfn_restart_file_name=wfn_restart_file_name,
        )

        if kpoints:
            dft.insert(Kpoints.from_kpoints(kpoints))
        if smearing or (band_gap <= 0.0):
            scf.kwargs["ADDED_MOS"] = 100
            scf["ADDED_MOS"] = 100  # TODO: how to grab the appropriate number?
            scf.insert(Smear())

        # Create subsections and insert into them
        self["FORCE_EVAL"].insert(dft)
        xc_functional = XC_FUNCTIONAL(functional=kwargs.get("functional", "PBE"))
        xc = Section("XC", subsections={"XC_FUNCTIONAL": xc_functional})
        self["FORCE_EVAL"]["DFT"].insert(xc)
        self["FORCE_EVAL"]["DFT"].insert(Section("PRINT", subsections={}))

        if isinstance(structure, Molecule):
            self.activate_nonperiodic()

        if kwargs.get("print_pdos", True):
            self.print_pdos()
        if kwargs.get("print_ldos", False):
            self.print_ldos()
        if kwargs.get("print_mo_cubes", True):
            self.print_mo_cubes()
        if kwargs.get("print_hartree_potential", False):
            self.print_hartree_potential()
        if kwargs.get("print_e_density", False):
            self.print_e_density()

        self.update(self.override_default_params)

    def print_pdos(self, nlumo=-1):
        """
        Activate creation of the PDOS file.

        Args:
            nlumo (int): Number of virtual orbitals to be added to the MO set (-1=all).
                CAUTION: Setting this value to be higher than the number of states present may cause a Cholesky error.
        """
        if not self.check("FORCE_EVAL/DFT/PRINT/PDOS"):
            self["FORCE_EVAL"]["DFT"]["PRINT"].insert(PDOS(nlumo=nlumo))

    def print_ldos(self, nlumo=-1):
        """
        Activate the printing of LDOS files, printing one for each atom kind by default

        Args:
            nlumo (int): Number of virtual orbitals to be added to the MO set (-1=all).
                CAUTION: Setting this value to be higher than the number of states present may cause a Cholesky error.
        """
        if not self.check("FORCE_EVAL/DFT/PRINT/PDOS"):
            self["FORCE_EVAL"]["DFT"]["PRINT"].insert(PDOS(nlumo=nlumo))
        for i in range(self.structure.num_sites):
            self["FORCE_EVAL"]["DFT"]["PRINT"]["PDOS"].insert(LDOS(i + 1, alias="LDOS {}".format(i + 1), verbose=False))

    def print_mo_cubes(self, write_cube=False, nlumo=-1, nhomo=-1):
        """
        Activate printing of molecular orbitals.

        Args:
            write_cube (bool): whether to write cube file for the MOs (setting false will just print levels in out file)
            nlumo (int): Controls the number of lumos that are printed and dumped as a cube (-1=all)
            nhomo (int): Controls the number of homos that are printed and dumped as a cube (-1=all)
        """
        if not self.check("FORCE_EVAL/DFT/PRINT/MO_CUBES"):
            self["FORCE_EVAL"]["DFT"]["PRINT"].insert(MO_Cubes(write_cube=write_cube, nlumo=nlumo, nhomo=nhomo))

    def print_mo(self):
        """
        Print molecular orbitals when running non-OT diagonalization
        """
        raise NotImplementedError

    def print_hartree_potential(self, stride=[1, 1, 1]):
        """
        Controls the printing of a cube file with eletrostatic potential generated by the total density
        (electrons+ions). It is valid only for QS with GPW formalism.
        Note that by convention the potential has opposite sign than the expected physical one.
        """
        if not self.check("FORCE_EVAL/DFT/PRINT/V_HARTREE_CUBE"):
            self["FORCE_EVAL"]["DFT"]["PRINT"].insert(V_Hartree_Cube(keywords={"STRIDE": Keyword("STRIDE", *stride)}))

    def print_e_density(self):
        """
        Controls the printing of cube files with the electronic density and, for LSD calculations, the spin density
        """
        if not self.check("FORCE_EVAL/DFT/PRINT/E_DENSITY_CUBE"):
            self["FORCE_EVAL"]["DFT"]["PRINT"].insert(E_Density_Cube())

    def set_charge(self, charge):
        """
        Set the overall charge of the simulation cell
        """
        self["FORCE_EVAL"]["DFT"]["CHARGE"] = Keyword("CHARGE", charge)

    def activate_hybrid(
        self,
        hybrid_functional: str = "PBE0",
        hf_fraction: float = 0.25,
        gga_x_fraction: float = 0.75,
        gga_c_fraction: float = 1,
        max_memory: int = 2000,
        cutoff_radius: float = 8.0,
        potential_type: str = None,
        omega: float = 0.2,
        aux_basis: Union[Dict, None] = None,
        admm: bool = True,
        eps_schwarz: float = 1e-6,
        eps_schwarz_forces: float = 1e-6,
        screen_on_initial_p: bool = True,
        screen_p_forces: bool = True,
    ):
        """
        Basic set for activating hybrid DFT calculation using Auxiliary Density Matrix Method.

        Note 1: When running ADMM with cp2k, memory is very important. If the memory requirements exceed
        what is available (see max_memory), then CP2K will have to calculate the 4-electron integrals
        for HFX during each step of the SCF cycle. ADMM provides a huge speed up by making the memory
        requirements *feasible* to fit into RAM, which means you only need to calculate the integrals
        once each SCF cycle. But, this only works if it fits into memory. When setting up ADMM
        calculations, we recommend doing whatever is possible to fit all the 4EI into memory.

        Note 2: This set is designed for reliable high-throughput calculations, NOT for extreme
        accuracy. Please review the in-line comments in this method if you want more control.

        Args:
            hybrid_functional (str): Type of hybrid functional. This set supports HSE (screened) and PBE0
                (truncated). Default is PBE0, which converges easier in the GPW basis used by
                cp2k.
            hf_fraction (float): fraction of exact HF exchange energy to mix. Default: 0.25
            gga_x_fraction (float): fraction of gga exchange energy to retain. Default: 0.75
            gga_c_fraction (float): fraction of gga correlation energy to retain. Default: 1.0
            max_memory (int): Maximum memory available to each MPI process (in Mb) in the calculation.
                Most modern computing nodes will have ~2Gb per core, or 2048 Mb, but check for
                your specific system. This value should be as large as possible while still leaving
                some memory for the other parts of cp2k. Important: If this value is set larger
                than the memory limits, CP2K will likely seg-fault.
                Default: 2000
            cutoff_radius (float): for truncated hybrid functional (i.e. PBE0), this is the cutoff
                radius. The default is selected as that which generally gives convergence, but
                maybe too low (if you want very high accuracy) or too high (if you want a quick
                screening). Default: 8 angstroms
            potential_type (str): what interaction potential to use for HFX. Available in CP2K are
                COULOMB, GAUSSIAN, IDENTITY, LOGRANGE, MIX_CL, MIX_CL_TRUNC, MIX_LG, SHORTRANGE,
                and TRUNCATED. Default is None, and it will be set automatically depending on the
                named hybrid_functional that you use, but setting it to one of the acceptable
                values will constitute a user-override.
            omega (float): For HSE, this specifies the screening parameter. HSE06 sets this as
                0.2, which is the default.
            aux_basis (dict): If you want to specify the aux basis to use, specify it as a dict of
                the form {'specie_1': 'AUX_BASIS_1', 'specie_2': 'AUX_BASIS_2'}
            admm (bool): Whether or not to use the auxiliary density matrix method for the exact
                HF exchange contribution. Highly recommended. Speed ups between 10x and aaa1000x are
                possible when compared to non ADMM hybrid calculations. Default: True
            eps_schwarz (float): Screening threshold for HFX, in Ha. Contributions smaller than this
                will be screened. The smaller the value, the more accurate, but also the more
                costly. Default value is 1e-6, which is quite aggressive. Aggressive screening
                can also lead to convergence issues. 1e-7 should be a safe value if 1e-6 is too
                aggressive.
            eps_schwarz_forces (float): Same as for eps_schwarz, but for screening contributions to
                forces. Convergence is not as sensitive with respect to eps_schwarz forces as
                compared to eps_schwarz, and so 1e-6 should be good default.
            screen_on_initial_p (bool): If an initial density matrix is provided, in the form of a
                CP2K wfn restart file, then this initial density will be used for screening. This
                is generally very computationally efficient, but, as with eps_schwarz, can lead to
                instabilities if the initial density matrix is poor.
            screen_p_forces (bool): Same as screen_on_initial_p, but for screening of forces.
        """
        if admm:
            aux_basis = aux_basis if aux_basis else {}
            aux_basis = {s: aux_basis[s] if s in aux_basis else None for s in self.structure.symbol_set}
            basis = get_aux_basis(basis_type=aux_basis)
            if isinstance(self["FORCE_EVAL"]["DFT"]["BASIS_SET_FILE_NAME"], KeywordList):
                self["FORCE_EVAL"]["DFT"]["BASIS_SET_FILE_NAME"].extend(
                    [Keyword("BASIS_SET_FILE_NAME", k) for k in ["BASIS_ADMM", "BASIS_ADMM_MOLOPT"]],
                )

            for k, v in self["FORCE_EVAL"]["SUBSYS"].subsections.items():
                if v.name.upper() == "KIND":
                    kind = v["ELEMENT"].values[0]
                    v.keywords["BASIS_SET"] += Keyword("BASIS_SET", "AUX_FIT", basis[kind])

            # Don't change unless you know what you're doing
            # Use NONE for accurate eigenvalues (static calcs)
            aux_matrix_params = {
                "ADMM_PURIFICATION_METHOD": Keyword("ADMM_PURIFICATION_METHOD", "NONE"),
                "METHOD": Keyword("METHOD", "BASIS_PROJECTION"),
            }
            aux_matrix = Section(
                "AUXILIARY_DENSITY_MATRIX_METHOD",
                keywords=aux_matrix_params,
                subsections={},
            )
            self.subsections["FORCE_EVAL"]["DFT"].insert(aux_matrix)

        # Define the GGA functional as PBE
        pbe = PBE("ORIG", scale_c=gga_c_fraction, scale_x=gga_x_fraction)
        xc_functional = XC_FUNCTIONAL("PBE", subsections={"PBE": pbe})

        screening = Section(
            "SCREENING",
            subsections={},
            keywords={
                "EPS_SCHWARZ": Keyword("EPS_SCHWARZ", eps_schwarz),
                "EPS_SCHWARZ_FORCES": Keyword("EPS_SCHWARZ_FORCES", eps_schwarz_forces),
                "SCREEN_ON_INITIAL_P": Keyword("SCREEN_ON_INITIAL_P", screen_on_initial_p),
                "SCREEN_P_FORCES": Keyword("SCREEN_P_FORCES", screen_p_forces),
            },
        )

        ip_keywords = {}
        if hybrid_functional == "HSE06":
            potential_type = potential_type if potential_type else "SHORTRANGE"
            xc_functional.insert(
                Section(
                    "XWPBE",
                    subsections={},
                    keywords={
                        "SCALE_X0": Keyword("SCALE_X0", 1),
                        "SCALE_X": Keyword("SCALE_X", -hf_fraction),
                        "OMEGA": Keyword("OMEGA", omega),
                    },
                )
            )
            ip_keywords.update(
                {
                    "POTENTIAL_TYPE": Keyword("POTENTIAL_TYPE", potential_type),
                    "OMEGA": Keyword("OMEGA", omega),
                }
            )
        elif hybrid_functional == "PBE0":
            potential_type = potential_type if potential_type else "TRUNCATED"
            ip_keywords.update(
                {
                    "POTENTIAL_TYPE": Keyword("POTENTIAL_TYPE", potential_type),
                    "CUTOFF_RADIUS": Keyword("CUTOFF_RADIUS", cutoff_radius),
                    "T_C_G_DATA": Keyword("T_C_G_DATA", "t_c_g.dat"),
                }
            )
        interaction_potential = Section("INTERACTION_POTENTIAL", subsections={}, keywords=ip_keywords)

        # Unlikely for users to override
        load_balance = Section(
            "LOAD_BALANCE",
            keywords={"RANDOMIZE": Keyword("RANDOMIZE", True)},
            subsections={},
        )

        # EPS_STORAGE_SCALING squashes the integrals for efficient storage
        # Unlikely for users to override.
        memory = Section(
            "MEMORY",
            subsections={},
            keywords={
                "EPS_STORAGE_SCALING": Keyword("EPS_STORAGE_SCALING", 0.1),
                "MAX_MEMORY": Keyword("MAX_MEMORY", max_memory),
            },
        )
        hf = Section(
            "HF",
            keywords={"FRACTION": Keyword("FRACTION", hf_fraction)},
            subsections={
                "SCREENING": screening,
                "INTERACTION_POTENTIAL": interaction_potential,
                "LOAD_BALANCE": load_balance,
                "MEMORY": memory,
            },
        )
        xc = Section("XC", subsections={"XC_FUNCTIONAL": xc_functional, "HF": hf})

        self.subsections["FORCE_EVAL"]["DFT"].insert(xc)

    def activate_fast_minimization(self, on):
        """
        Method to modify the set to use fast SCF minimization.
        """
        if on:
            ot = OrbitalTransformation(
                minimizer="DIIS",
                preconditioner="FULL_ALL",
                algorithm="IRAC",
                energy_gap=0.01,
                linesearch="2PNT",
            )
            self.update({"FORCE_EVAL": {"DFT": {"SCF": {"OT": ot}}}})

    def activate_robust_minimization(self):
        """
        Method to modify the set to use more robust SCF minimization technique
        """
        ot = OrbitalTransformation(
            minimizer="CG",
            preconditioner="FULL_ALL",
            algorithm="STRICT",
            energy_gap=0.05,
            linesearch="3PNT",
        )
        self.update({"FORCE_EVAL": {"DFT": {"SCF": {"OT": ot}}}})

    def activate_very_strict_minimization(self):
        """
        Method to modify the set to use very strict SCF minimization scheme
        :return:
        """
        ot = OrbitalTransformation(
            minimizer="CG",
            preconditioner="FULL_ALL",
            algorithm="STRICT",
            energy_gap=0.05,
            linesearch="GOLD",
        )
        self.update({"FORCE_EVAL": {"DFT": {"SCF": {"OT": ot}}}})

    def activate_nonperiodic(self):
        """
        Activates a calculation with non-periodic calculations by turning of PBC and
        changing the poisson solver. Still requires a CELL to put the atoms
        """
        kwds = {
            "POISSON_SOLVER": Keyword("POISSON_SOLVER", "MT"),
            "PERIODIC": Keyword("PERIODIC", "NONE"),
        }
        self["FORCE_EVAL"]["DFT"].insert(Section("POISSON", subsections={}, keywords=kwds))
        if not self.check("FORCE_EVAL/SUBSYS/CELL"):
            x = max([s.coords[0] for s in self.structure.sites])
            y = max([s.coords[1] for s in self.structure.sites])
            z = max([s.coords[2] for s in self.structure.sites])
            self["FORCE_EVAL"]["SUBSYS"].insert(Cell(lattice=Lattice([[x, 0, 0], [0, y, 0], [0, 0, z]])))
        self["FORCE_EVAL"]["SUBSYS"]["CELL"].add(Keyword("PERIODIC", "NONE"))

    def modify_dft_print_iters(self, iters, add_last="no"):
        """
        Modify all DFT print iterations at once. Common use is to set iters to the max
        number of iterations + 1 and then set add_last to numeric. This would have the
        effect of printing only the first and last iteration, which might be useful for
        speeding up/saving space on GEO_OPT or MD runs where you don't need the intermediate
        values.

        Args
            iters (int): print each "iters" iterations.
            add_last (str): Whether to explicitly include the last iteration, and how to mark it.
                numeric: mark last iteration with the iteration number
                symbolic: mark last iteration with the letter "l"
                no: do not explicitly include the last iteration
        """
        assert add_last.lower() in ["no", "numeric", "symbolic"]
        if self.check("FORCE_EVAL/DFT/PRINT"):
            run_type = self["global"].get("run_type", Keyword("run_type", "energy")).values[0]
            for k, v in self["force_eval"]["dft"]["print"].subsections.items():

                if v.name.upper() in [
                    "ACTIVE_SPACE",
                    "BAND_STRUCTURE",
                    "GAPW",
                    "IMPLICIT_PSOLVER",
                    "SCCS",
                    "WFN_MIX",
                ]:
                    continue

                v.insert(
                    Section(
                        "EACH",
                        subsections=None,
                        keywords={run_type: Keyword(run_type, iters)},
                    )
                )
                v.keywords["ADD_LAST"] = Keyword("ADD_LAST", add_last)

class StaticSet(DftSet):

    """
    Basic static energy calculation. Turns on Quickstep module, sets the run_type in global,
    and uses structure object to build the subsystem.
    """

    def __init__(
        self,
        structure: Union[Structure, Molecule],
        project_name: str = "Static",
        run_type: str = "ENERGY_FORCE",
        override_default_params: Dict = {},
        **kwargs,
    ):
        """
        Args:
            structure: Pymatgen structure object
            project_name (str): What to name this cp2k project (controls naming of files printed out)
            run_type (str): Run type. As a static set it should be one of the static aliases, like 'ENERGY_FORCE'
        """
        super().__init__(structure, **kwargs)
        global_section = Global(project_name=project_name, run_type=run_type)
        self.structure = structure
        self.project_name = project_name
        self.run_type = run_type
        self.override_default_params = override_default_params
        self.insert(global_section)
        self.update(override_default_params)
        self.kwargs = kwargs

class RelaxSet(DftSet):

    """
    CP2K input set containing the basic settings for performing geometry optimization. Values are all cp2k
    defaults, and should be good for most systems of interest.
    """

    def __init__(
        self,
        structure: Union[Structure, Molecule],
        max_drift: float = 3e-3,
        max_force: float = 4.5e-3,
        max_iter: int = 200,
        project_name: str = "Relax",
        optimizer: str = "BFGS",
        override_default_params: Dict = {},
        **kwargs,
    ):
        """
        Args:
            structure:
            max_drift: Convergence criterion for the maximum geometry change between the current and the
                last optimizer iteration. This keyword cannot be repeated and it expects precisely one real.
                Default value: 3.00000000E-003
                Default unit: [bohr]
            max_force (float): Convergence criterion for the maximum force component of the current configuration.
                This keyword cannot be repeated and it expects precisely one real.
                Default value: 4.50000000E-004
                Default unit: [bohr^-1*hartree]
            max_iter (int): Specifies the maximum number of geometry optimization steps.
                One step might imply several force evaluations for the CG and LBFGS optimizers.
                This keyword cannot be repeated and it expects precisely one integer.
                Default value: 200
            optimizer (str): Specify which method to use to perform a geometry optimization.
                This keyword cannot be repeated and it expects precisely one keyword. BFGS is a
                quasi-newtonian method, and will best for "small" systems near the minimum. LBFGS
                is a limited memory version that can be used for "large" (>1000 atom) systems when
                efficiency outweights robustness. CG is more robust, especially when you are far from
                the minimum, but it slower.
                Default value: BFGS
        """
        super().__init__(structure, **kwargs)

        self.structure = structure
        self.max_drift = max_drift
        self.max_force = max_force
        self.max_iter = max_iter
        self.project_name = project_name
        self.optimizer = optimizer
        self.override_default_params = override_default_params
        self.kwargs = kwargs

        global_section = Global(project_name=project_name, run_type="GEO_OPT")

        geo_opt_params = {
            "TYPE": Keyword("TYPE", "MINIMIZATION"),
            "MAX_DR": Keyword("MAX_DR", max_drift),
            "MAX_FORCE": Keyword("MAX_FORCE", max_force),
            "RMS_DR": Keyword("RMS_DR", 1.5e-3),
            "MAX_ITER": Keyword("MAX_ITER", max_iter),
            "OPTIMIZER": Keyword("OPTIMIZER", optimizer),
        }
        geo_opt = Section("GEO_OPT", subsections={}, keywords=geo_opt_params)
        if not self.check("MOTION"):
            self.insert(Section("MOTION", subsections={}))
        self["MOTION"].insert(geo_opt)
        self.insert(global_section)
        self.modify_dft_print_iters(max_iter + 1, add_last="numeric")
        self.update(override_default_params)

class CellOptSet(DftSet):

    """
    CP2K input set containing the basic settings for performing geometry optimization. Values are all cp2k
    defaults, and should be good for most systems of interest.
    """

    def __init__(
        self,
        structure: Union[Structure, Molecule],
        project_name: str = "CellOpt",
        override_default_params: Dict = {},
        **kwargs,
    ):
        """
        Args:
            structure: Pymatgen structure object
            max_drift: Convergence criterion for the maximum geometry change between the current and the
                last optimizer iteration. This keyword cannot be repeated and it expects precisely one real.
                Default value: 3.00000000E-003
                Default unit: [bohr]
            max_force (float): Convergence criterion for the maximum force component of the current configuration.
                This keyword cannot be repeated and it expects precisely one real.
                Default value: 4.50000000E-004
                Default unit: [bohr^-1*hartree]
            max_iter (int): Specifies the maximum number of geometry optimization steps.
                One step might imply several force evaluations for the CG and LBFGS optimizers.
                This keyword cannot be repeated and it expects precisely one integer.
                Default value: 200
            optimizer (str): Specify which method to use to perform a geometry optimization.
                This keyword cannot be repeated and it expects precisely one keyword. BFGS is a
                quasi-newtonian method, and will best for "small" systems near the minimum. LBFGS
                is a limited memory version that can be used for "large" (>1000 atom) systems when
                efficiency outweights robustness. CG is more robust, especially when you are far from
                the minimum, but it slower.
                Default value: BFGS
        """
        super().__init__(structure, **kwargs)

        self.structure = structure
        self.project_name = project_name
        self.override_default_params = override_default_params
        self.kwargs = kwargs
        global_section = Global(project_name=project_name, run_type="CELL_OPT")
        self.insert(global_section)
        self.modify_dft_print_iters(self.get("max_iter", 200) + 1, add_last="numeric")
        self.update(override_default_params)

class HybridStaticSet(StaticSet):

    """
    Static calculation using hybrid DFT with the ADMM formalism in Cp2k.
    """

    def __init__(
        self,
        structure: Union[Structure, Molecule],
        hybrid_functional: str = "PBE0",
        hf_fraction: float = 0.25,
        project_name: str = "Hybrid-Static",
        gga_x_fraction: float = 0.75,
        gga_c_fraction: float = 1,
        override_default_params: Dict = {},
        max_memory: int = 2000,
        cutoff_radius: float = 8.0,
        omega: float = 0.2,
        aux_basis: Union[Dict, None] = None,
        admm: bool = True,
        eps_schwarz: float = 1e-6,
        eps_schwarz_forces: float = 1e-6,
        screen_on_initial_p: bool = True,
        screen_p_forces: bool = True,
        **kwargs,
    ):
        """
        Args:
            structure: pymatgen structure object
            method: hybrid dft method to use (currently select between HSE06 and PBE0)
            hf_fraction: percentage of exact HF to mix-in
            project_name: what to call this project
            gga_x_fraction: percentage of gga exchange to use
            gga_c_fraction: percentage of gga correlation to use
            override_default_params: override settings (see above).
        """
        super().__init__(structure, project_name=project_name, **kwargs)

        self.structure = structure
        self.hybrid_functional = hybrid_functional
        self.hf_fraction = hf_fraction
        self.project_name = project_name
        self.gga_x_fraction = gga_x_fraction
        self.gga_c_fraction = gga_c_fraction
        self.override_default_params = override_default_params
        self.max_memory = max_memory
        self.cutoff_radius = cutoff_radius
        self.omega = omega
        self.aux_basis = aux_basis
        self.admm = admm
        self.eps_schwarz = eps_schwarz
        self.eps_schwarz_forces = eps_schwarz_forces
        self.screen_on_initial_p = screen_on_initial_p
        self.screen_p_forces = screen_p_forces
        self.kwargs = kwargs

        self.activate_hybrid(
            hybrid_functional=hybrid_functional,
            hf_fraction=hf_fraction,
            gga_x_fraction=gga_x_fraction,
            gga_c_fraction=gga_c_fraction,
            max_memory=max_memory,
            cutoff_radius=cutoff_radius,
            omega=omega,
            aux_basis=aux_basis,
            admm=admm,
            eps_schwarz=eps_schwarz,
            eps_schwarz_forces=eps_schwarz_forces,
            screen_on_initial_p=screen_on_initial_p,
            screen_p_forces=screen_p_forces,
        )
        self.update(override_default_params)

class HybridRelaxSet(RelaxSet):

    """
    Static calculation using hybrid DFT with the ADMM formalism in Cp2k.
    """

    def __init__(
        self,
        structure: Union[Structure, Molecule],
        hybrid_functional: str = "PBE0",
        hf_fraction: float = 0.25,
        project_name: str = "Hybrid-Relax",
        gga_x_fraction: float = 0.75,
        gga_c_fraction: float = 1,
        override_default_params: Dict = {},
        max_memory: int = 2000,
        cutoff_radius: float = 8.0,
        omega: float = 0.2,
        aux_basis: Union[Dict, None] = None,
        admm: bool = True,
        eps_schwarz: float = 1e-6,
        eps_schwarz_forces: float = 1e-6,
        screen_on_initial_p: bool = True,
        screen_p_forces: bool = True,
        **kwargs,
    ):
        """
        Args:
            structure: pymatgen structure object
            method: hybrid dft method to use (currently select between HSE06 and PBE0)
            hf_fraction: percentage of exact HF to mix-in
            project_name: what to call this project
            gga_x_fraction: percentage of gga exchange to use
            gga_c_fraction: percentage of gga correlation to use
            override_default_params: override settings (see above).
        """
        super().__init__(structure, project_name=project_name, **kwargs)

        self.structure = structure
        self.hybrid_functional = hybrid_functional
        self.hf_fraction = hf_fraction
        self.project_name = project_name
        self.gga_x_fraction = gga_x_fraction
        self.gga_c_fraction = gga_c_fraction
        self.override_default_params = override_default_params
        self.max_memory = max_memory
        self.cutoff_radius = cutoff_radius
        self.omega = omega
        self.aux_basis = aux_basis
        self.admm = admm
        self.eps_schwarz = eps_schwarz
        self.eps_schwarz_forces = eps_schwarz_forces
        self.screen_on_initial_p = screen_on_initial_p
        self.screen_p_forces = screen_p_forces
        self.kwargs = kwargs

        self.activate_hybrid(
            hybrid_functional=hybrid_functional,
            hf_fraction=hf_fraction,
            gga_x_fraction=gga_x_fraction,
            gga_c_fraction=gga_c_fraction,
            max_memory=max_memory,
            cutoff_radius=cutoff_radius,
            omega=omega,
            aux_basis=aux_basis,
            admm=admm,
            eps_schwarz=eps_schwarz,
            eps_schwarz_forces=eps_schwarz_forces,
            screen_on_initial_p=screen_on_initial_p,
            screen_p_forces=screen_p_forces,
        )
        self.update(override_default_params)

class HybridCellOptSet(CellOptSet):

    """
    Static calculation using hybrid DFT with the ADMM formalism in Cp2k.
    """

    def __init__(
        self,
        structure: Union[Structure, Molecule],
        hybrid_functional: str = "PBE0",
        hf_fraction: float = 0.25,
        project_name: str = "Hybrid-CellOpt",
        gga_x_fraction: float = 0.75,
        gga_c_fraction: float = 1,
        override_default_params: Dict = {},
        max_memory: int = 2000,
        cutoff_radius: float = 8.0,
        omega: float = 0.2,
        aux_basis: Union[Dict, None] = None,
        admm: bool = True,
        eps_schwarz: float = 1e-6,
        eps_schwarz_forces: float = 1e-6,
        screen_on_initial_p: bool = True,
        screen_p_forces: bool = True,
        **kwargs,
    ):
        """
        Args:
            structure: pymatgen structure object
            method: hybrid dft method to use (currently select between HSE06 and PBE0)
            hf_fraction: percentage of exact HF to mix-in
            project_name: what to call this project
            gga_x_fraction: percentage of gga exchange to use
            gga_c_fraction: percentage of gga correlation to use
            override_default_params: override settings (see above).
        """
        super().__init__(structure, project_name=project_name, **kwargs)

        self.structure = structure
        self.hybrid_functional = hybrid_functional
        self.hf_fraction = hf_fraction
        self.project_name = project_name
        self.gga_x_fraction = gga_x_fraction
        self.gga_c_fraction = gga_c_fraction
        self.override_default_params = override_default_params
        self.max_memory = max_memory
        self.cutoff_radius = cutoff_radius
        self.omega = omega
        self.aux_basis = aux_basis
        self.admm = admm
        self.eps_schwarz = eps_schwarz
        self.eps_schwarz_forces = eps_schwarz_forces
        self.screen_on_initial_p = screen_on_initial_p
        self.screen_p_forces = screen_p_forces
        self.kwargs = kwargs

        self.activate_hybrid(
            hybrid_functional=hybrid_functional,
            hf_fraction=hf_fraction,
            gga_x_fraction=gga_x_fraction,
            gga_c_fraction=gga_c_fraction,
            max_memory=max_memory,
            cutoff_radius=cutoff_radius,
            omega=omega,
            aux_basis=aux_basis,
            admm=admm,
            eps_schwarz=eps_schwarz,
            eps_schwarz_forces=eps_schwarz_forces,
            screen_on_initial_p=screen_on_initial_p,
            screen_p_forces=screen_p_forces,
        )
        self.update(override_default_params)


import numpy as np
import pandas as pd
import networkx as nx
import sys

def main(argv):
   if len(argv) != 1:
      raise ValueError("Must pass in graph representation input")

   df = pd.read_csv(argv[0], sep="\t")
   assert len(df.index) > 0

   DEBUG = False

   G = nx.from_pandas_edgelist(
      df,
      "FromNodeId",
      "ToNodeId")

   if DEBUG:
      print("G number of nodes: ", G.number_of_nodes())
      print("G max node: ", max(list(G.nodes)))
      print("G number of edges: ", G.number_of_edges())

      print("Number of edges for node 9721: ", G.degree[9721])
      print("9721 neighbors: ", G.edges(9721))
   assert G.number_of_edges() > 0
   d = nx.shortest_path_length(G, source=0)
   d_sorted = dict(sorted(d.items()))
   p = nx.shortest_path(G, source=0)

   f = open("python_shortest_path_outs.txt", "w")
   f.write("source\tdistance\n")
   for (key, value) in d_sorted.items():
      f.write(f"{key}\t{value}\n")

   print(p[1969993][-1])

if __name__ == "__main__":
   main(sys.argv[1:])


from sys import argv
from stenograpi import Stenograpi
from argparse import ArgumentParser

def parse_args(args):
    parser = ArgumentParser(prog='stenograpi.py',
                            description='Document your HTTP API automatically through tests.')

    parser.add_argument('hostname', type=str, help='hostname of Stenograpi')
    parser.add_argument('port', type=int, help='port Stenograpi should listen on')

    parser.add_argument('app-hostname', type=str, help='hostname of your app')
    parser.add_argument('app-port', type=int, help='port your app is listening on')
    return parser.parse_args(args)

if __name__ == '__main__':
    args = parse_args(argv[1:])
    stenograpi = Stenograpi(args.hostname, args.port, getattr(args, 'app-port'))
    stenograpi.listen()
    request = stenograpi.get_latest_request()
    print(request)


import numpy as np
import matplotlib.pyplot as plt
import math

def quad2(x_o=[-9.0,9.0],a=2.0,eta=0.1,threshold=0.01,maxiter=1000,alpha=0,anim = 1):
    it = 0
    x1 = np.linspace(-10,10,21)

    x2 = np.linspace(-10,10,21)

    [X,Y] = np.meshgrid(x1,x2)

    Y = (a*X**2+Y**2)/2

    plt.clf()
    plt.contour(Y,10)
    plt.xticks([0,5,10,15,20],[-10, -5, 0, 5, 10])
    plt.yticks([0,5,10,15,20],[-10, -5, 0, 5, 10])
    ax = plt.gca()
    ax.set_aspect('equal','box')

    plt.tight_layout()

    f = (a*x_o[0]**2+x_o[1]**2)/2

    varx = np.array([0,0])
    ###Gradient Method####
    while it != maxiter:
        fold = f

        grad = np.array([a*x_o[0], x_o[1]])

        varx = alpha*varx+(1-alpha)*grad
        x_old = np.asarray(x_o)

        x_o = np.asarray(x_o-eta*varx)

        try:
            f = (a*x_o[0]**2+x_o[1]**2)/2
            if (f < threshold or fold < threshold):
                break
            else:
                if anim:
                    plt.plot([x_old[0]+10, x_o[0]+10],[x_old[1]+10,x_o[1]+10],'r.-')
                    plt.pause(0.2)
                it += 1
        except:
            print('Diverged!')
            #plt.show()
            break

    if it == maxiter:
        return False, it
    else:
        return True, it+1

if __name__ == "__main__":
    eta = (0.003, 0.01, 0.03, 0.1, 0.3, 1, 3, 10)
    alpha = (0, 0.5, 0.7, 0.9, 0.95)
    for e in eta:
        for a in alpha:
            conv, iters = quad2(a=20.0,eta=e,anim = 0,alpha=a)
            if conv:
                string = f"{iters} iterations"
            else:
                string = "Diverged!"
            print(f"eta: {e}, alpha: {a} ->" + string)


"""
    Author: Emanuele Albini

    Adapter from a counterfactual method to background generator.
"""
from typing import Union
import numpy as np

from ..base import (BaseBackgroundGenerator, BackgroundGenerator, CounterfactualMethod, MultipleCounterfactualMethod,
                    ListOf2DArrays)

from ..utils import (
    get_top_counterfactuals,
    expand_dims_counterfactuals,
)

__all__ = ['CounterfactualMethodBackgroundGeneratorAdapter']

class CounterfactualMethodBackgroundGeneratorAdapter(BaseBackgroundGenerator, BackgroundGenerator):
    """Adapter to make a counterfactual method into a background generator"""
    def __init__(
            self,
            counterfactual_method: CounterfactualMethod,
            n_top: Union[None, int] = None,  # By default: All
    ):
        """

        Args:
            counterfactual_method (CounterfactualMethod): The counterfactual method
            n_top (Union[None, int], optional): Number of top-counterfactuals to select as background. Defaults to None (all).
        """
        self.counterfactual_method = counterfactual_method
        self.n_top = n_top

    def get_backgrounds(self, X: np.ndarray) -> ListOf2DArrays:
        """Generate the background datasets for each query instance

        Args:
            X (np.ndarray): The query instances

        Returns:
            ListOf2DArrays: A list/array of background datasets
                nb_query_intances x nb_background_points x nb_features
        """

        X = self.preprocess(X)

        # If we do not have any background then we compute it
        if isinstance(self.counterfactual_method, MultipleCounterfactualMethod):
            if self.n_top is None:
                return self.counterfactual_method.get_multiple_counterfactuals(X)
            elif self.n_top == 1:
                return expand_dims_counterfactuals(self.counterfactual_method.get_counterfactuals(X))
            else:
                return get_top_counterfactuals(
                    self.counterfactual_method.get_multiple_counterfactuals(X),
                    X,
                    n_top=self.n_top,
                    nan=False,
                )
        elif isinstance(self.counterfactual_method, CounterfactualMethod):
            if self.n_top is not None and self.n_top != 1:
                raise ValueError('Counterfactual methodology do not supportthe generation of multiple counterfactuals.')
            return np.expand_dims(self.counterfactual_method.get_counterfactuals(X), axis=1)
        else:
            raise NotImplementedError('Unsupported CounterfactualMethod.')


from mpl_toolkits.mplot3d import Axes3D
import numpy as np
import matplotlib.pyplot as plt
import gif
import matplotlib.animation as animation
train = np.loadtxt('data/click.csv', delimiter=',', dtype='int', skiprows=1)
train_x = train[:, 0]
train_y = train[:, 1]

mu = train_x.mean()
sigma = train_x.std()

# zscore 标准化，mean变为0，sigama变为1
def standardize(x):
    return (x - mu) / sigma

train_z = standardize(train_x)

# 预测函数 一元一次：f(x)-theta0+theta1*x
def f(x):
    return theta0 + theta1 * x
# 目标函数
def E(x, y):
    return 0.5 * np.sum((y - f(x)) ** 2)

#代价函数
theta1 = np.arange(-10, 10, 0.1)
theta0 = np.arange(-10, 10, 0.1)
theta1, theta0 = np.meshgrid(theta1, theta0)
print(theta1)  # 打印出来瞅瞅
print(theta0)

es = 0
#因为theta1和theta0已经变成了网格矩阵了
#一次性带入全部计算，我们需要一个一个的算
n = len(train_y)
for i in range(n):
	y_pre = theta1*train_x[i]+theta0  # 取出一个样本在网格矩阵上计算，得到一个预测矩阵
	e = (train_y[i]-y_pre)**2  # 标准值减去预测（矩阵）得到方差矩阵
	es += e  # 把单样本上的方差矩阵不断累加到es上
es = es/n  # 求平均值，这样es方差矩阵每个点的位置就是对应的theta1和theta0矩阵每个点位置预测得到的方差
fig = plt.figure()
ax = Axes3D(fig)
ax.set_zlim(0, 500000)

#plot_surface函数绘制曲面
#cmap='rainbow表示彩虹图（用不同的颜色表示不同值）
ax.plot_surface(theta1, theta0, es, cmap='rainbow')
#显示图像
plt.show()

# 学习率
ETA = 1e-3
# 误差的差值
diff = 1
# 更新次数
epoch = 0
# 参数初始化
theta0 = np.random.rand()
theta1 = np.random.rand()
# 直到误差的差值小于 0.01 为止，重复参数更新
error = E(train_z, train_y)
errors = []
errors.append(error)

# ？这里用diff来判断是否停止循环，也就是当收敛时停止循环，我看有些人的代码是用epoch来规定迭代次数，所以机器学习里是用哪一个？
fig = plt.figure()
ims = []
plt.xlim(-3, 3)
plt.ylim(0, 700)
x = np.linspace(-3, 3, 100)
while diff > 1e-2:
    # 更新结果保存到临时变量 (问：为什么要先保存在临时变量中？ 答：为了同时更新参数，不然一个直接改了，另一个参数的更新也在改变)
    tmp_theta0 = theta0 - ETA * np.sum((f(train_z) - train_y))
    tmp_theta1 = theta1 - ETA * np.sum((f(train_z) - train_y) * train_z)

    # 更新参数
    theta0 = tmp_theta0
    theta1 = tmp_theta1

    # 计算与上一次误差的差值
    error = E(train_z, train_y)
    errors.append(error)
    diff = abs(errors[-2]-errors[-1])
    # 输出日志
    epoch += 1
    log = 'Epoch {}  : theta0 = {:.3f}, theta1 = {:.3f}, diff = {:.4f}'
    print(log.format(epoch, theta0, theta1, diff))

    # 动态显示f(x)变化图
    # plt.text(0.8, 1, f'Epoch:{epoch:},Loss:{error:.2f}')
    im = plt.plot(train_z, train_y, 'o', color="blue") + plt.plot(x, f(x))
    ims.append(im)

ani = animation.ArtistAnimation(fig, ims, interval=20, repeat_delay=1000)
ani.save("regression1_linear.gif", writer='pillow')

# # 绘图确认
plt.clf()
plt.plot(train_z, train_y, 'o')
plt.plot(x, f(x))
plt.show()

# # error 不断在下降
x = np.arange(len(errors))
plt.plot(x, errors)
plt.show()


# This Python 3 environment comes with many helpful analytics libraries installed
# It is defined by the kaggle/python docker image: https://github.com/kaggle/docker-python
# For example, here's several helpful packages to load in

import numpy as np # linear algebra
import pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)
import random
# Input data files are available in the "../input/" directory.
# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory

# import os
# for dirname, _, filenames in os.walk('/kaggle/input'):
#     for filename in filenames:
#         print(os.path.join(dirname, filename))

# Any results you write to the current directory are saved as output.

from tensorflow.python.keras.layers import Dense, Dropout, Input, Conv2D, Flatten
from tensorflow.python.keras.models import Sequential, Model
from tensorflow.python.keras.layers.advanced_activations import LeakyReLU
from tensorflow.python.keras.optimizers import adam
from tensorflow.python.keras.layers import Reshape, Conv2DTranspose, BatchNormalization, Activation
import matplotlib.pyplot as plt
import numpy as np
from tqdm import tqdm
import glob
from PIL import Image

num_epochs = 50
batch_size = 32
input_height = 32
input_width = 32
output_height = 256
output_width = 256

train_dir = '/kaggle/input/flower-enhance/train/'
val_dir = '/kaggle/input/flower-enhance/test/'
train_size = len(
    glob.glob(train_dir + "/*-in.jpg"))

small_input_filenames = glob.glob(train_dir + "/*-in.jpg")
large_input_filenames = glob.glob(train_dir + "/*-out.jpg")
# small_images = np.zeros(
#             (train_size, input_width, input_height, 3))
Small_images = np.empty((train_size, 32*32*3), dtype=np.uint8)
Large_images = np.empty((train_size, 256, 256, 3), dtype=np.uint8)

for i in range(train_size):
    small_img = small_input_filenames[i]
    large_img = large_input_filenames[i]
    small_images = np.array(Image.open(small_img))
    small_images = (small_images.astype(np.float32) - 127.5)/127.5
    small_images = small_images.reshape(1, input_width*input_height*3)
    Small_images[i] = small_images
    large_images = np.array(Image.open(large_img))
    large_images = (large_images.astype(np.float32) - 127.5) / 127.5
    # large_images = large_images.reshape(1, input_width * input_height * 3)
    Large_images[i] = large_images

def adam_optimizer():

    return adam(lr=0.0002, beta_1=0.5)

z_dim = input_height*input_width*3

def create_generator():
    model = Sequential()

    # Reshape input into 32x32x256 tensor via a fully connected layer
    model.add(Dense(64 * 32 * 32, input_dim=z_dim))
    model.add(Reshape((32, 32, 64)))
    model.add(Conv2DTranspose(                  # Transposed convolution layer, from 32x32x256 into 64x64x128 tensor
        128, kernel_size=3, strides=2, padding='same'))
    model.add(BatchNormalization())             # Batch normalization
    model.add(LeakyReLU(alpha=0.01))            # Leaky ReLU
    model.add(Conv2DTranspose(                  # Transposed convolution layer, from 64x64x128 to 128x128x64 tensor
        64, kernel_size=3, strides=2, padding='same'))
    model.add(BatchNormalization())             # Batch normalization
    model.add(LeakyReLU(alpha=0.01))            # Leaky ReLU
    model.add(Conv2DTranspose(                  # Transposed convolution layer, from 128x128x64 to 256x256x3 tensor
        3, kernel_size=3, strides=2, padding='same'))
    model.add(Activation('tanh'))               # Tanh activation
    model.compile(loss='binary_crossentropy', optimizer=adam_optimizer())
    model.summary()
    return model

def create_discriminator():

    model = Sequential()
    model.add(Conv2D(32, kernel_size=3, strides=2, padding="same", input_shape=(256, 256, 3)))
    model.add(LeakyReLU(alpha=0.2))
    model.add(Dropout(0.25))
    model.add(BatchNormalization(momentum=0.8))
    model.add(Conv2D(64, kernel_size=3, strides=2, padding="same"))
    model.add(LeakyReLU(alpha=0.2))
    model.add(Dropout(0.25))
    model.add(BatchNormalization(momentum=0.8))
    model.add(Conv2D(128, kernel_size=3, strides=2, padding="same"))
    model.add(LeakyReLU(alpha=0.2))
    model.add(Dropout(0.25))
    model.add(BatchNormalization(momentum=0.8))
    model.add(Flatten())
    model.add(Dense(1, activation='sigmoid'))
    model.compile(loss='binary_crossentropy', optimizer=adam_optimizer())
    model.summary()
    return model

def image_generator(batch_size, img_dir):
    """A generator that returns small images and large images.  DO NOT ALTER the validation set"""
    input_filenames = glob.glob(img_dir + "/*-in.jpg")
    counter = 0
    random.shuffle(input_filenames)
    while True:
        small_images = np.zeros(
            (batch_size, 32, 32, 3))
        large_images = np.zeros(
            (batch_size, 256, 256, 3))
        if counter+batch_size >= len(input_filenames):
            counter = 0
        for i in range(batch_size):
            img = input_filenames[counter + i]
            small_images[i] = np.array(Image.open(img)) / 255.0
            large_images[i] = np.array(
                Image.open(img.replace("-in.jpg", "-out.jpg"))) / 255.0
        small_images = small_images.reshape(batch_size, 32*32*3)
        yield (small_images, large_images)
        counter += batch_size

def create_gan(discriminator, generator):
    discriminator.trainable = False
    gan_input = Input(shape=32*32*3)
    x = generator(gan_input)
    gan_output = discriminator(x)
    gan = Model(inputs=gan_input, outputs=gan_output)
    gan.compile(loss='binary_crossentropy', optimizer='adam')
    return gan

def plot_generated_images(train_dir, epoch, generator, examples=100, dim=(10,10), figsize=(10,10)):
    img_gen_object = image_generator(100, train_dir)
    small_images, large_images = next(img_gen_object)
    noise = small_images[np.random.randint(low=0, high=small_images.shape[0], size=examples)]
    generated_images = generator.predict(noise)*255
    generated_images = generated_images.reshape(100, 256, 256, 3)
    plt.figure(figsize=figsize)
    for i in range(generated_images.shape[0]):
        plt.subplot(dim[0], dim[1], i+1)
        plt.imshow(generated_images[i], interpolation='nearest')
        plt.axis('off')
    plt.tight_layout()
    plt.savefig('gan_generated_image %d.png' %epoch)

def training(train_dir, epochs=1, batch_size=32):
    # Loading the data
    # (X_train, y_train, X_test, y_test) = load_data()
    # batch_count = X_train.shape[0] / batch_size

    # Creating GAN
    generator = create_generator()
    discriminator = create_discriminator()
    gan = create_gan(discriminator, generator)

    for e in range(1, epochs + 1):
        img_gen_object = image_generator(batch_size, train_dir)
        small_images, large_images = next(img_gen_object)
        print("Epoch %d" % e)
        for _ in tqdm(range(batch_size)):
            # generate  random noise as an input  to  initialize the  generator
            noise = small_images[np.random.randint(low=0, high=small_images.shape[0], size=batch_size)]

            # Generate fake MNIST images from noised input
            generated_images = generator.predict(noise)

            # Get a random set of  real images
            image_batch = large_images[np.random.randint(low=0, high=large_images.shape[0], size=batch_size)]

            # Construct different batches of  real and fake data
            X = np.concatenate([image_batch, generated_images])

            # Labels for generated and real data
            y_dis = np.zeros(2 * batch_size)
            y_dis[:batch_size] = 0.9

            # Pre train discriminator on  fake and real data  before starting the gan.
            discriminator.trainable = True
            discriminator.train_on_batch(X, y_dis)

            # Tricking the noised input of the Generator as real data
            noise = small_images[np.random.randint(low=0, high=small_images.shape[0], size=batch_size)]
            y_gen = np.ones(batch_size)

            # During the training of gan,
            # the weights of discriminator should be fixed.
            # We can enforce that by setting the trainable flag
            discriminator.trainable = False

            # training  the GAN by alternating the training of the Discriminator
            # and training the chained GAN model with Discriminator’s weights freezed.
            gan.train_on_batch(noise, y_gen)

        if e == 1 or e % 20 == 0:
            plot_generated_images(train_dir, e, generator)

training(train_dir, 50, 32)


# -*- coding: utf-8 -*-
"""{{ cookiecutter.repo_name }} URL Configuration

https://docs.djangoproject.com/en/1.8/topics/http/urls/

"""
from django.conf import settings
from django.conf.urls import include, url
from django.conf.urls.static import static
from django.contrib import admin
from django.contrib.staticfiles.storage import staticfiles_storage
from django.http import HttpResponse
from django.views.defaults import (
    server_error
)
from django.views.generic.base import RedirectView

from wagtail.contrib.wagtailapi import urls as wagtailapi_urls
from wagtail.wagtailcore import urls as wagtail_urls
from wagtail.wagtailadmin import urls as wagtailadmin_urls
from wagtail.wagtaildocs import urls as wagtaildocs_urls
from wagtail.wagtailsearch import urls as wagtailsearch_urls

urlpatterns = [

    # www.example.com/500.html
    url(
        regex=r'^500/$',
        view=server_error),
    # www.example.com/django-admin
    url(
        regex=r'^favicon.ico$',
        view=RedirectView.as_view(
            url=staticfiles_storage.url('favicon.ico'),
            permanent=False),
        name="favicon"
    ),
    url(
        regex=r'^django-admin/',
        view=include(admin.site.urls)),
    # www.example.com/admin
    url(
        regex=r'^admin/',
        view=include(wagtailadmin_urls)),
    # www.example.com/search
    url(
        regex=r'^search/',
        view=include(wagtailsearch_urls)),
    # www.example.com/documents
    url(
        regex=r'^documents/',
        view=include(wagtaildocs_urls)),
    # www.example.com/api/vi/{pages, images, documents}
    url(
        regex=r'^api/',
        view=include(wagtailapi_urls)),
    # www.example.com
    url(
        regex=r'',
        view=include(wagtail_urls)),
    # www.example.com/robots.txt
    url(
        regex=r'^robots.txt$',
        view=lambda r: HttpResponse("User-agent: *\nDisallow: /",
                                    content_type="text/plain")),
] + static(settings.MEDIA_URL, document_root=settings.MEDIA_ROOT)

# url pattern for django debug toolbar
if settings.DEBUG:
    import debug_toolbar
    urlpatterns += [
        url(r'^__debug__/', include(debug_toolbar.urls)),
    ]


import zipfile
import os

import torchvision.transforms as transforms

# once the images are loaded, how do we pre-process them before being passed into the network
# by default, we resize the images to 64 x 64 in size
# and normalize them to mean = 0 and standard-deviation = 1 based on statistics collected from
# the training set

mean_nums = [0.485, 0.456, 0.406]
std_nums = [0.229, 0.224, 0.225]

resnet_train_transforms = transforms.Compose([
    transforms.Resize(size=224),
    transforms.CenterCrop(size=224),
    transforms.RandomRotation(degrees=10),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize(mean_nums, std_nums)
])

resnet_test_transforms = transforms.Compose([
    transforms.Resize(224),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean_nums, std_nums)
])

inceptionv3_train_transforms = transforms.Compose([
    transforms.RandomRotation(degrees=10),
    transforms.RandomHorizontalFlip(),
    transforms.Resize(299),
    transforms.CenterCrop(299),
    transforms.ToTensor(),
    transforms.Normalize(mean_nums, std_nums)
])

inceptionv3_test_transforms = transforms.Compose([
    transforms.Resize(299),
    transforms.CenterCrop(299),
    transforms.ToTensor(),
    transforms.Normalize(mean_nums, std_nums)
])

inceptionv3_params = {
    'batch_size': 32,
    'momentum': 0.9,
    'lr': 0.045,
    'misc': 'RMSProp(0.9, eps=1.0), exp lr decay 0.94 every 2 epochs, grad clipping for stability',
}


from django.contrib.auth.models import User
from django.db import models

class Location(models.Model):
    name = models.CharField(max_length=255)
    world = models.CharField(max_length=255)
    description = models.TextField()
    image = models.ImageField(upload_to='photos/')
    author = models.ForeignKey(User, on_delete=models.CASCADE, blank=True, null=True)

    def __str__(self):
        return self.name

class Character(models.Model):
    name = models.CharField(max_length=255)
    birth_location = models.ForeignKey(Location, on_delete=models.CASCADE, related_name='birth_location')
    location = models.ForeignKey(Location, on_delete=models.CASCADE, related_name='current_location')
    description = models.TextField()
    image = models.ImageField(upload_to='photos/')
    GENDER = (
        ('male', 'Мужской'),
        ('female', 'Женский'),
        ('unknown', 'Неизвестно')
    )
    gender = models.CharField(max_length=10, choices=GENDER)
    race = models.CharField(max_length=55)
    DEAD = (
        ('dead', 'Мертвый'),
        ('alive', 'Живой')
    )
    dead = models.CharField(max_length=10, choices=DEAD)

    def __str__(self):
        return self.name

class Episode(models.Model):
    name = models.CharField(max_length=255)
    number = models.IntegerField()
    number_season = models.IntegerField()
    image = models.ImageField(upload_to='photos/')
    description = models.TextField()
    premiere = models.DateField()
    character = models.ManyToManyField(Character, related_name='char_to_ep')

    def __str__(self):
        return f"{self.number_season} + {self.number}"