import os
import numpy as np
import matplotlib.pyplot as pyplot
from simtk import unit
from foldamers.cg_model.cgmodel import CGModel
from foldamers.parameters.reweight import *
from foldamers.thermo.calc import calculate_heat_capacity
from foldamers.utilities.util import random_positions
from cg_openmm.build.cg_build import build_topology
from cg_openmm.simulation.rep_exch import *
from foldamers.ensembles.ens_build import *
from foldamers.thermo.calc import *

# Job settings
top_directory = "output"
if not os.path.exists(top_directory):
    os.mkdir(top_directory)

# OpenMM simulation settings
print_frequency = 5  # Number of steps to skip when printing output
total_simulation_time = 5.0 * unit.nanosecond  # Units = picoseconds
simulation_time_step = 5.0 * unit.femtosecond
total_steps = round(total_simulation_time.__div__(simulation_time_step))

# Yank (replica exchange) simulation settings
output_data = str(str(top_directory) + "/output.nc")
number_replicas = 30
min_temp = 5.0 * unit.kelvin
max_temp = 100.0 * unit.kelvin
temperature_list = get_temperature_list(min_temp, max_temp, number_replicas)
print("Using " + str(len(temperature_list)) + " replicas.")
if total_steps > 10000:
    exchange_attempts = round(total_steps / 1000)
else:
    exchange_attempts = 10

###
#
# Global coarse grained model definitions
#
###
backbone_length = 1
sidechain_length = 1
sidechain_positions = 0
include_bond_forces = False
include_bond_angle_forces = True
include_nonbonded_forces = True
include_torsion_forces = True
constrain_bonds = True

# Particle properties
mass = 100.0 * unit.amu
masses = {"backbone_bead_masses": mass, "sidechain_bead_masses": mass}

# Bonded interaction properties
bond_length = 7.5 * unit.angstrom
bond_lengths = {
    "bb_bb_bond_length": bond_length,
    "bb_sc_bond_length": bond_length,
    "sc_sc_bond_length": bond_length,
}
bond_force_constant = 1250 * unit.kilojoule_per_mole / unit.nanometer / unit.nanometer
bond_force_constants = {
    "bb_bb_bond_k": bond_force_constant,
    "bb_sc_bond_k": bond_force_constant,
    "sc_sc_bond_k": bond_force_constant,
}

# Bond angle properties
bond_angle_force_constant = 0.0002 * unit.kilojoule_per_mole / unit.radian / unit.radian
bond_angle_force_constants = {
    "bb_bb_bb_angle_k": bond_angle_force_constant,
    "bb_bb_sc_angle_k": bond_angle_force_constant,
    "bb_sc_sc_angle_k": bond_angle_force_constant,
    "sc_sc_sc_angle_k": bond_angle_force_constant,
    "sc_bb_sc_angle_k": bond_angle_force_constant,
    "sc_sc_bb_angle_k": bond_angle_force_constant,
}
equil_bond_angle = 92
equil_bond_angles = {
    "bb_bb_bb_angle_0": equil_bond_angle,
    "bb_bb_sc_angle_0": equil_bond_angle,
    "bb_sc_sc_angle_0": equil_bond_angle,
    "sc_sc_sc_angle_0": equil_bond_angle,
    "sc_bb_sc_angle_0": equil_bond_angle,
    "sc_sc_bb_angle_0": equil_bond_angle,
}

# Torsion properties
torsion_force_constant = 0.0002
torsion_force_constants = {
    "bb_bb_bb_bb_torsion_k": torsion_force_constant,
    "bb_bb_bb_sc_torsion_k": torsion_force_constant,
    "bb_bb_sc_sc_torsion_k": torsion_force_constant,
    "bb_sc_sc_sc_torsion_k": torsion_force_constant,
    "sc_bb_bb_sc_torsion_k": torsion_force_constant,
    "bb_sc_sc_bb_torsion_k": torsion_force_constant,
    "sc_sc_sc_sc_torsion_k": torsion_force_constant,
    "sc_bb_bb_bb_torsion_k": torsion_force_constant,
}
equil_torsion_angle = 52
equil_torsion_angles = {
    "bb_bb_bb_bb_torsion_0": equil_torsion_angle,
    "bb_bb_bb_sc_torsion_0": equil_torsion_angle,
    "bb_bb_sc_sc_torsion_0": equil_torsion_angle,
    "bb_sc_sc_sc_torsion_0": equil_torsion_angle,
    "sc_bb_bb_sc_torsion_0": equil_torsion_angle,
    "bb_sc_sc_bb_torsion_0": equil_torsion_angle,
    "sc_sc_sc_sc_torsion_0": equil_torsion_angle,
    "sc_bb_bb_bb_torsion_0": equil_torsion_angle,
}

sigma = 2.0 * bond_length
sigmas = {"bb_bb_sigma": sigma, "sc_sc_sigma": sigma}
epsilon = 0.2 * unit.kilocalorie_per_mole
epsilons = {"bb_bb_eps": epsilon, "bb_sc_eps": epsilon, "sc_sc_eps": epsilon}

# Heteropolymer definitions
heteropolymer = True

# Define individual monomer properties
monomer_name = "A"
num_beads = backbone_length + sidechain_length
sigmas = {"bb_bb_sigma": sigma, "bb_sc_sigma": sigma, "sc_sc_sigma": 2.0 * sigma}
A = {
    "monomer_name": monomer_name,
    "backbone_length": backbone_length,
    "sidechain_length": sidechain_length,
    "sidechain_positions": sidechain_positions,
    "num_beads": num_beads,
    "bond_lengths": bond_lengths,
    "epsilons": epsilons,
    "sigmas": sigmas,
}

monomer_name = "B"
num_beads = backbone_length + sidechain_length
sigmas = {"bb_bb_sigma": sigma, "bb_sc_sigma": sigma, "sc_sc_sigma": 0.8 * sigma}
B = {
    "monomer_name": monomer_name,
    "backbone_length": backbone_length,
    "sidechain_length": sidechain_length,
    "sidechain_positions": sidechain_positions,
    "num_beads": num_beads,
    "bond_lengths": bond_lengths,
    "epsilons": epsilons,
    "sigmas": sigmas,
}

monomer_types = [A, B]

sequence = [A, A, A, B, A, A, A, B, A, A, A, B]
polymer_length = len(sequence)

cgmodel = CGModel(
    polymer_length=polymer_length,
    bond_force_constants=bond_force_constants,
    bond_angle_force_constants=bond_angle_force_constants,
    torsion_force_constants=torsion_force_constants,
    equil_bond_angles=equil_bond_angles,
    equil_torsion_angles=equil_torsion_angles,
    include_nonbonded_forces=include_nonbonded_forces,
    include_bond_forces=include_bond_forces,
    include_bond_angle_forces=include_bond_angle_forces,
    include_torsion_forces=include_torsion_forces,
    constrain_bonds=constrain_bonds,
    heteropolymer=True,
    monomer_types=monomer_types,
    sequence=sequence,
)

if not os.path.exists(output_data):
    replica_energies, replica_positions, replica_states = run_replica_exchange(
        cgmodel.topology,
        cgmodel.system,
        cgmodel.positions,
        temperature_list=temperature_list,
        simulation_time_step=simulation_time_step,
        total_simulation_time=total_simulation_time,
        print_frequency=print_frequency,
        output_data=output_data,
    )
    make_replica_pdb_files(cgmodel.topology, replica_positions)
else:
    replica_energies, replica_positions, replica_states = read_replica_exchange_data(
        system=cgmodel.system,
        topology=cgmodel.topology,
        temperature_list=temperature_list,
        output_data=output_data,
        print_frequency=print_frequency,
    )

C_v, dC_v, new_temperature_list = get_heat_capacity(
    replica_energies, temperature_list, num_intermediate_states=1
)

exit()


"""
Description:
Author: Jiaqi Gu (jqgu@utexas.edu)
Date: 2021-12-27 02:35:47
LastEditors: Jiaqi Gu (jqgu@utexas.edu)
LastEditTime: 2021-12-27 02:45:29
"""
import numpy as np
import torch
import torch.nn.functional as F
from pyutils.quantize import input_quantize_fn, weight_quantize_fn
from torch import nn
from torch.nn import Parameter, init
from torch.types import Device

__all__ = ["MLGConv2d"]

class MLGConv2d(nn.Module):
    """
    description: Conv2d layer with memory-efficient multi-level generation (MLG)
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: int = 3,
        stride: int = 1,
        padding: int = 1,
        bias: bool = False,
        w_bit: int = 16,
        in_bit: int = 16,
        device: Device = torch.device("cuda"),
    ):
        super().__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.stride = stride
        self.padding = padding

        self.w_bit = w_bit
        self.qb = w_bit
        self.qu = w_bit
        self.qv = w_bit
        self.in_bit = in_bit
        self.device = device

        ### allocate parameters
        self.weight = None
        ### build trainable parameters
        self.build_parameters()

        ### quantization tool
        self.input_quantizer = input_quantize_fn(self.in_bit, alg="normal", device=self.device)

        self.weight_quantizer = weight_quantize_fn(self.w_bit, alg="dorefa_sym")
        self.basis_quantizer = weight_quantize_fn(self.qb, alg="dorefa_sym")
        self.coeff_in_quantizer = weight_quantize_fn(self.qu, alg="dorefa_sym")
        self.coeff_out_quantizer = weight_quantize_fn(self.qv, alg="dorefa_sym")

        ### default set to slow forward
        self.disable_fast_forward()
        ### default disable dynamic weight generation
        self.disable_dynamic_weight()
        self.eye_b = None
        self.eye_v = None

        if bias:
            self.bias = Parameter(torch.Tensor(out_channels).to(self.device))
        else:
            self.register_parameter("bias", None)

    def build_parameters(self):
        self.weight = Parameter(
            torch.Tensor(self.out_channels, self.in_channels, self.kernel_size, self.kernel_size)
            .to(self.device)
            .float()
        )

    def reset_parameters(self):
        init.kaiming_normal_(self.weight.data, mode="fan_out", nonlinearity="relu")

        if self.bias is not None:
            fan_in = self.in_channels * self.kernel_size ** 2
            bound = 1 / np.sqrt(fan_in)
            init.uniform_(self.bias, -bound, bound)

    def build_weight(self):
        if self.w_bit < 16:
            if self.dynamic_weight_flag:
                if self.coeff_in is not None:
                    coeff_in = self.coeff_in_quantizer(self.coeff_in)
                    self.basis = self.weight[
                        : self.base_out, : self.base_in, ...
                    ]  # [base_out, base_in, k, k]
                else:
                    coeff_in = None
                    self.basis = self.weight[: self.base_out, ...]  # [base_out, inc, k, k]
                basis = self.basis_quantizer(self.basis)
                if self.coeff_out is not None:
                    coeff_out = self.coeff_out_quantizer(self.coeff_out)
                else:
                    coeff_out = None
                weight = self.weight_generation(basis, coeff_in, coeff_out)
            else:
                weight = self.weight_quantizer(self.weight)
        else:
            weight = self.weight
            if self.dynamic_weight_flag:
                if self.coeff_in is not None:
                    self.basis = weight[: self.base_out, : self.base_in, ...]  # [base_out, base_in, k, k]
                else:
                    self.basis = weight[: self.base_out, ...]  # [base_out, inc, k, k]

                weight = self.weight_generation(self.basis, self.coeff_in, self.coeff_out)

        return weight

    def enable_fast_forward(self):
        self.fast_forward_flag = True

    def disable_fast_forward(self):
        self.fast_forward_flag = False

    def load_parameters(self, param_dict):
        """
        description: update parameters based on this parameter dictionary\\
        param param_dict {dict of dict} {layer_name: {param_name: param_tensor, ...}, ...}
        """
        for name, param in param_dict.items():
            getattr(self, name).data.copy_(param)

    def enable_dynamic_weight(self, base_in, base_out):
        ### multi-level weight generation
        self.base_in = base_in  # input channel base
        self.base_out = base_out  # output channel base

        if base_out == 0:  ## disable cross-kernel generation
            self.base_out = self.out_channels  ## maximum
        elif min(self.out_channels, self.in_channels * self.kernel_size ** 2) > self.base_out > 0:
            ### enable generation
            self.base_out = base_out
        else:
            ### base_out is too large, cannot save param, then disable it
            self.base_out = self.out_channels

        ### only when base_in < min(in_channel, kernel_size**2), will intra-kernel generation save #params.
        if min(self.in_channels, self.kernel_size ** 2) > self.base_in > 0:
            self.coeff_in = Parameter(
                torch.Tensor(self.base_out, self.in_channels, self.base_in).to(self.device)
            )
            # init.xavier_normal_(self.coeff_in)
            init.kaiming_normal_(self.coeff_in, mode="fan_out", nonlinearity="relu")
        else:
            ### base_in >= min(in_channel, kernel_size**2), will use the original weight
            self.coeff_in = None
        ### onlt when base_out < min(out_channel, in_channel*kernel_size**2), will cross-kernel generation save #params.
        if min(self.out_channels, self.in_channels * self.kernel_size ** 2) > self.base_out > 0:
            self.coeff_out = Parameter(torch.Tensor(self.out_channels, self.base_out).to(self.device))
            init.kaiming_normal_(self.coeff_out, mode="fan_out", nonlinearity="relu")
        else:
            self.coeff_out = None

        self.dynamic_weight_flag = True if self.coeff_in is not None or self.coeff_out is not None else False
        if self.dynamic_weight_flag:
            if self.coeff_in is not None:
                self.basis = self.weight[: self.base_out, : self.base_in, ...]
            else:
                self.basis = self.weight[: self.base_out, ...]
        else:
            self.basis = None

    def disable_dynamic_weight(self):
        self.dynamic_weight_flag = False

    def weight_generation(self, basis, coeff_in, coeff_out):
        ### Level 1
        if coeff_in is not None:
            # weight_1 [base_out, inc, k^2]
            # coeff_in x basis = [bo, inc, bi] x [bo, bi, k^2]
            basis = basis.view(basis.size(0), basis.size(1), -1)
            weight_1 = torch.matmul(coeff_in, basis)
        else:
            weight_1 = basis

        ### Level 2
        if coeff_out is not None:
            # weight_2 [outc, inc*k*k]
            weight_1 = weight_1.view(weight_1.size(0), -1)
            weight_2 = torch.matmul(coeff_out, weight_1)
            weight_2 = weight_2.view(self.out_channels, self.in_channels, self.kernel_size, self.kernel_size)
        else:
            ## do not use self.out_channel, since for dwconv, we should set out_channel to 1 here.
            weight_2 = weight_1.view(
                self.weight.size(0), self.in_channels, self.kernel_size, self.kernel_size
            )
        return weight_2

    def get_output_dim(self, img_height, img_width):
        h_out = (img_height - self.kernel_size + 2 * self.padding) / self.stride + 1
        w_out = (img_width - self.kernel_size + 2 * self.padding) / self.stride + 1
        return (int(h_out), int(w_out))

    def get_num_params(self, fullrank=False):
        if (self.dynamic_weight_flag == True) and (fullrank == False):
            total = self.basis.numel()
            if self.coeff_in is not None:
                total += self.coeff_in.numel()
            if self.coeff_out is not None:
                total += self.coeff_out.numel()
        else:
            total = self.weight.numel()
        if self.bias is not None:
            total += self.bias.numel()

        return total

    def get_param_size(self, fullrank=False):
        total = 0
        if (self.dynamic_weight_flag == True) and (fullrank == False):
            total += self.basis.numel() * self.qb / 8
            if self.coeff_in is not None:
                total += self.coeff_in.numel() * self.qu / 8
            if self.coeff_out is not None:
                total += self.coeff_out.numel() * self.qv / 8
        else:
            total += self.weight.numel() * 4
        if self.bias is not None:
            total += self.bias.numel() * 4
        return total

    def get_ortho_loss(self):
        ### we want row vectors in the basis to be orthonormal
        if self.dynamic_weight_flag:  ### at least one-level generation
            ## basis ortho loss always exists !
            if self.coeff_in is not None and self.coeff_in.size(2) > 1:
                if self.basis.size(1) > 1:
                    ### only penalize when there are at least two row/column vectors
                    basis = self.basis.view(self.basis.size(0), self.basis.size(1), -1)  # [bo, bi, k^2]
                    dot_b = torch.matmul(
                        basis, basis.permute([0, 2, 1])
                    )  # [bo, bi, k^2] x [bo, k^2, bi] = [bo, bi, bi]
                else:
                    dot_b = None
                ## U
                coeff_in = self.coeff_in / (
                    self.coeff_in.data.norm(p=2, dim=1, keepdim=True) + 1e-8
                )  # normalization
                dot_u = torch.matmul(
                    coeff_in.permute(0, 2, 1), coeff_in
                )  # [bo, bi, ci-bi] x [bo, ci-bi, bi] = [bo, bi, bi]
            else:
                dot_u = None

            if self.coeff_out is not None:
                if self.coeff_in is None:
                    ### if there is no intra-kernel generation, only cross-kernel generation, e.g., conv1x1, we have to treat basis as a matrix [bo, cin*k*k] and encourage it to have bo orthogonal rows
                    basis = self.basis.view(self.basis.size(0), -1)  # [bo, ci*k^2]
                    dot_b = torch.matmul(
                        basis, basis.permute([1, 0])
                    )  # [bo, ci*k^2] x [ci*k^2, bo] = [bo, bo]
                # V
                coeff_out = self.coeff_out / (
                    self.coeff_out.data.norm(p=2, dim=0, keepdim=True) + 1e-8
                )  # normalization
                dot_v = torch.matmul(coeff_out.t(), coeff_out)  # [bo, co-bo] x [co-bo, bo] = [bo, bo]
            else:
                dot_v = None
            if self.basis is not None and self.eye_b is None:
                self.eye_b = torch.eye(dot_b.size(-1), dtype=dot_b.dtype, device=dot_b.device)
                if dot_b.ndim > 2:
                    self.eye_b = self.eye_b.unsqueeze(0).repeat(basis.size(0), 1, 1)
            if self.coeff_out is not None and self.eye_v is None:
                self.eye_v = torch.eye(dot_v.size(-1), dtype=dot_v.dtype, device=dot_v.device)
            loss = 0
            if dot_b is not None:
                loss = loss + F.mse_loss(dot_b, self.eye_b)
            if dot_u is not None:
                loss = loss + F.mse_loss(dot_u, self.eye_b)
            if dot_v is not None:
                loss = loss + F.mse_loss(dot_v, self.eye_v)
        else:
            loss = 0

        return loss

    def assign_separate_weight_bit(self, qb, qu, qv, quant_ratio_b=1, quant_ratio_u=1, quant_ratio_v=1):
        qb, qu, qv = min(qb, 32), min(qu, 32), min(qv, 32)
        self.qb, self.qu, self.qv = qb, qu, qv
        self.basis_quantizer = weight_quantize_fn(qb, alg="dorefa_sym")
        self.coeff_in_quantizer = weight_quantize_fn(qu, alg="dorefa_sym")
        self.coeff_out_quantizer = weight_quantize_fn(qv, alg="dorefa_sym")
        self.basis_quantizer.set_quant_ratio(quant_ratio_b)
        self.coeff_in_quantizer.set_quant_ratio(quant_ratio_u)
        self.coeff_out_quantizer.set_quant_ratio(quant_ratio_v)

    def set_quant_ratio(self, quant_ratio_b=1, quant_ratio_u=1, quant_ratio_v=1, quant_ratio_in=1):
        if hasattr(self, "basis_quantizer"):
            self.basis_quantizer.set_quant_ratio(quant_ratio_b)
        if hasattr(self, "coeff_in_quantizer"):
            self.coeff_in_quantizer.set_quant_ratio(quant_ratio_u)
        if hasattr(self, "coeff_out_quantizer"):
            self.coeff_out_quantizer.set_quant_ratio(quant_ratio_v)
        self.input_quantizer.set_quant_ratio(quant_ratio_in)

    def forward(self, x):
        if self.in_bit < 16:
            x = self.input_quantizer(x)
        if not self.fast_forward_flag or self.weight is None:
            weight = self.build_weight()
        else:
            weight = self.weight
        #### record weight_2
        self.weight_2 = weight

        out = F.conv2d(x, weight, bias=self.bias, stride=self.stride, padding=self.padding)

        return out

    def extra_repr(self):
        s = (
            "{in_channels}, {out_channels}, kernel_size={kernel_size}"
            ", stride={stride}, padding={padding}, wb={w_bit}, ib={in_bit}"
        )
        if self.bias is None:
            s += ", bias=False"
        if self.dynamic_weight_flag:
            s += "base_in={base_in}, base_out={base_out}"
        return s.format(**self.__dict__)


# -*- coding: utf-8 -*-
# Author: XuMing <xuming624@qq.com>
# Brief: 字符到字符的基本seq2seq模型
# input:hello; output:ehllo

import numpy as np
import tensorflow as tf
from tensorflow.python.layers.core import Dense

# params
epochs = 60
batch_size = 128
rnn_size = 50
num_layers = 2
encoding_embedding_size = 15
decoding_embedding_size = 15
learning_rate = 0.001
checkpoint = 'model.ckpt'
display_step = 50
source_data_path = '../data/letters_source.txt'
target_data_path = '../data/letters_target.txt'

def get_corpus(data_path):
    with open(data_path, 'r', encoding='utf-8') as f:
        text = f.read().lower()
    return text

def extract_char_vocab(data):
    """
    mapping dict
    :param data:
    :return:
    """
    special_words = ['<PAD>', '<UNK>', '<BEGIN>', '<END>']
    set_chars = list(set([char for line in data.split() for char in line]))
    # add four special words to mapping dict
    indices_char = {i: c for i, c in enumerate(special_words + set_chars)}
    char_indices = {c: i for i, c in indices_char.items()}
    return indices_char, char_indices

def get_input():
    """
    input tensor
    :return:
    """
    inputs = tf.placeholder(tf.int32, [None, None], name='inputs')
    targets = tf.placeholder(tf.int32, [None, None], name='targets')
    learning_rate = tf.placeholder(tf.float32, name='learning_rate')

    # get target sequence maxlen
    target_sequence_len = tf.placeholder(tf.int32, (None,), name='target_sequence_len')
    target_sequence_maxlen = tf.reduce_max(target_sequence_len, name='target_sequence_maxlen')
    source_sequence_len = tf.placeholder(tf.int32, (None,), name='source_sequence_len')

    return inputs, targets, learning_rate, target_sequence_len, target_sequence_maxlen, source_sequence_len

def get_encoder_layer(input_data, rnn_size, num_layers,
                      source_sequence_len, source_vocab_size,
                      encoding_embedding_size):
    """
    encoder layer
    :param intput_data:
    :param rnn_size:
    :param num_layers:
    :param source_sequence_len:
    :param source_vocab_size:
    :param encodeing_embedding_size:
    :return:
    """
    # encoder embedding
    encoder_embed_input = tf.contrib.layers.embed_sequence(input_data,
                                                           source_vocab_size,
                                                           encoding_embedding_size)

    # RNN cell
    def get_lstm_cell(rnn_size):
        lstm_cell = tf.contrib.rnn.LSTMCell(rnn_size,
                                            initializer=tf.random_uniform_initializer(-0.1, 0.1, seed=2))
        return lstm_cell

    cell = tf.contrib.rnn.MultiRNNCell([get_lstm_cell(rnn_size) for i in range(num_layers)])
    encoder_output, encoder_state = tf.nn.dynamic_rnn(cell,
                                                      encoder_embed_input,
                                                      sequence_length=source_sequence_len,
                                                      dtype=tf.float32)
    return encoder_output, encoder_state

def process_deocder_input(data, vocab_indices, batch_size):
    """
    target sequence process: add <BEGIN>, and del last <END>
    :param data:
    :param vocab_indices:
    :param batch_size:
    :return:
    """
    # cut last char
    ending = tf.strided_slice(data, [0, 0], [batch_size, -1], [1, 1])
    decoder_input = tf.concat([tf.fill([batch_size, 1], vocab_indices['<BEGIN>']), ending], 1)
    return decoder_input

def decoding_layer(target_char_indices, decoding_embedding_size, num_layers, rnn_size,
                   target_sequence_len, target_sequence_maxlen, encoder_state, decoder_input,
                   batch_size=128):
    """
    decode layer
    :param target_char_indices:
    :param decoding_embedding_size:
    :param num_layers:
    :param rnn_size:
    :param target_sequence_len:
    :param target_sequence_maxlen:
    :param encoder_state:
    :param decoder_input:
    :return:
    """
    # embedding
    target_vocab_size = len(target_char_indices)
    decoder_embeddings = tf.Variable(tf.random_uniform([target_vocab_size, decoding_embedding_size]))
    decoder_embed_input = tf.nn.embedding_lookup(decoder_embeddings, decoder_input)

    # build decoder RNN cell
    def get_decoder_cell(rnn_size):
        decoder_cell = tf.contrib.rnn.LSTMCell(rnn_size, initializer=tf.random_uniform_initializer(-0.1, 0.1, seed=2))
        return decoder_cell

    cell = tf.contrib.rnn.MultiRNNCell([get_decoder_cell(rnn_size) for i in range(num_layers)])

    # output fc layer
    output_layer = Dense(target_vocab_size, kernel_initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.1))

    # training decoder
    with tf.variable_scope('decode'):
        training_helper = tf.contrib.seq2seq.TrainingHelper(inputs=decoder_embed_input,
                                                            sequence_length=target_sequence_len,
                                                            time_major=False)
        training_decoder = tf.contrib.seq2seq.BasicDecoder(cell,
                                                           training_helper,
                                                           encoder_state,
                                                           output_layer)
        training_decoder_output, _, _ = tf.contrib.seq2seq.dynamic_decode(training_decoder,
                                                                          impute_finished=True,
                                                                          maximum_iterations=target_sequence_maxlen)
    # predict decoder, share params with training decoder
    with tf.variable_scope('decode', reuse=True):
        start_tokens = tf.tile(tf.constant([target_char_indices['<BEGIN>']], dtype=tf.int32), [batch_size],
                               name='start_tokens')
        predicting_helper = tf.contrib.seq2seq.GreedyEmbeddingHelper(decoder_embeddings, start_tokens,
                                                                     target_char_indices['<END>'])
        predicting_decoder = tf.contrib.seq2seq.BasicDecoder(cell,
                                                             predicting_helper,
                                                             encoder_state,
                                                             output_layer)
        predicting_decoder_output, _, _ = tf.contrib.seq2seq.dynamic_decode(predicting_decoder,
                                                                            impute_finished=True,
                                                                            maximum_iterations=target_sequence_maxlen)
    return training_decoder_output, predicting_decoder_output

def seq2seq(input_data, targets, lr, target_sequence_len,
            target_sequence_maxlen, source_sequence_len,
            source_vocab_size, target_vocab_size,
            encoder_embedding_size, decoder_embedding_size,
            rnn_size, num_layers, target_char_indices, batch_size=128):
    """
    seq2seq model
    :param input_data:
    :param targets:
    :param lr:
    :param target_sequence_len:
    :param target_sequence_maxlen:
    :param source_sequence_len:
    :param source_vocab_size:
    :param target_vocab_size:
    :param encoder_embedding_size:
    :param decoder_embedding_size:
    :param rnn_size:
    :param num_layers:
    :return:
    """
    print('build model...')
    # get state output of encoder
    _, encoder_state = get_encoder_layer(input_data,
                                         rnn_size, num_layers,
                                         source_sequence_len, source_vocab_size,
                                         encoder_embedding_size)
    # input of decoder
    decoder_input = process_deocder_input(targets, target_char_indices, batch_size=batch_size)
    # decoder
    training_decoder_output, predicting_decoder_output = decoding_layer(target_char_indices,
                                                                        decoder_embedding_size,
                                                                        num_layers,
                                                                        rnn_size,
                                                                        target_sequence_len,
                                                                        target_sequence_maxlen,
                                                                        encoder_state,
                                                                        decoder_input)
    return training_decoder_output, predicting_decoder_output

def pad_sentence_batch(sentence_batch, pad_int):
    """
    pad the batch sequence, make sure every batch has same sequence_length
    :param sentence_batch:
    :param pad_int:
    :return:
    """
    max_sentence = max([len(sentence) for sentence in sentence_batch])
    return [sentence + [pad_int] * (max_sentence - len(sentence)) for sentence in sentence_batch]

def get_batches(targets, sources, batch_size, source_pad_int, target_pad_int):
    """
    get batch by generator
    :param targets:
    :param sources:
    :param batch_size:
    :param source_pad_int:
    :param target_pad_int:
    :return:
    """
    for batch_i in range(0, len(sources) // batch_size):
        start_i = batch_i * batch_size
        sources_batch = sources[start_i:start_i + batch_size]
        targets_batch = targets[start_i:start_i + batch_size]
        # pad sequence
        pad_sources_batch = np.array(pad_sentence_batch(sources_batch, source_pad_int))
        pad_targets_batch = np.array(pad_sentence_batch(targets_batch, target_pad_int))

        # get sentence length
        targets_lengths = []
        for target in targets_batch:
            targets_lengths.append(len(target))

        sources_lengths = []
        for source in sources_batch:
            sources_lengths.append(len(source))

        yield pad_targets_batch, pad_sources_batch, targets_lengths, sources_lengths

def train():
    source_data = get_corpus(source_data_path)
    target_data = get_corpus(target_data_path)
    print('corpus length:', len(source_data))

    # see sample data
    print(source_data.split('\n')[:10])
    print(target_data.split('\n')[:10])

    # get mapping dict
    source_indices_char, source_char_indices = extract_char_vocab(source_data)
    target_indices_char, target_char_indices = extract_char_vocab(target_data)

    # chars index
    source_indices = [[source_char_indices.get(c, source_char_indices['<UNK>']) for c in line]
                      for line in source_data.split('\n')]
    target_indices = [
        [target_char_indices.get(c, target_char_indices['<UNK>']) for c in line] + [target_char_indices['<END>']]
        for line in target_data.split('\n')]

    # see sample source indices data
    print(source_indices[:10])
    print(target_indices[:10])

    # split data to train and validation
    train_source, valid_source = source_indices[batch_size:], source_indices[:batch_size]
    train_target, valid_target = target_indices[batch_size:], target_indices[:batch_size]

    (valid_targets_batch, valid_sources_batch,
     valid_targets_lengths, valid_sources_lengths) = next(get_batches(valid_target,
                                                                      valid_source,
                                                                      batch_size,
                                                                      source_char_indices['<PAD>'],
                                                                      target_char_indices['<PAD>']))
    train_graph = tf.Graph()
    with train_graph.as_default():
        # get inputs
        input_data, targets, learning_rate, target_sequence_len, target_sequence_maxlen, source_sequence_len = get_input()
        training_decoder_output, predicting_decoder_output = seq2seq(input_data, targets,
                                                                     learning_rate, target_sequence_len,
                                                                     target_sequence_maxlen, source_sequence_len,
                                                                     len(source_char_indices), len(target_char_indices),
                                                                     encoding_embedding_size, decoding_embedding_size,
                                                                     rnn_size, num_layers,
                                                                     target_char_indices, batch_size)
        training_logits = tf.identity(training_decoder_output.rnn_output, 'logits')
        predicting_logits = tf.identity(predicting_decoder_output.sample_id, name='predictions')

        masks = tf.sequence_mask(target_sequence_len, target_sequence_maxlen, dtype=tf.float32, name='masks')
        with tf.name_scope('optimization'):
            # loss
            cost = tf.contrib.seq2seq.sequence_loss(training_logits, targets, masks)
            # optimizer
            optimizer = tf.train.AdamOptimizer(learning_rate)
            # gradient clipping
            gradients = optimizer.compute_gradients(cost)
            capped_gradients = [(tf.clip_by_value(grad, -5., 5.), var) for grad, var in gradients if grad is not None]
            train_op = optimizer.apply_gradients(capped_gradients)

    with tf.Session(graph=train_graph) as sess:
        sess.run(tf.global_variables_initializer())
        for epoch_i in range(1, epochs + 1):
            batches = get_batches(train_target, train_source, batch_size,
                                  source_char_indices['<PAD>'], target_char_indices['<PAD>'])
            for batch_i, (targets_batch, sources_batch, targets_length, sources_lengths) in enumerate(batches):
                _, loss = sess.run([train_op, cost],
                                   {input_data: sources_batch,
                                    targets: targets_batch,
                                    learning_rate: learning_rate,
                                    target_sequence_len: targets_length,
                                    source_sequence_len: sources_lengths})
                if batch_i % display_step == 0:
                    validation_loss = sess.run([cost],
                                               {input_data: valid_sources_batch,
                                                targets: valid_targets_batch,
                                                learning_rate: learning_rate,
                                                target_sequence_len: valid_targets_lengths,
                                                source_sequence_len: valid_sources_lengths})
                    print('Epoch {:>3}/{} Batch {:>4}/{} - Training Loss: {:>6.3f} - Validation Loss: {:>6.3f}'.format(
                        epoch_i,
                        epochs,
                        batch_i,
                        len(train_source) // batch_size,
                        loss,
                        validation_loss[0]
                    ))
        # save model
        saver = tf.train.Saver()
        saver.save(sess, checkpoint)
        print('Model trained and saved %s' % checkpoint)

def source_2_seq(text, source_char_indices):
    """
    change source data to sequence
    :param text:
    :param source_char_indices:
    :return:
    """
    sequence_len = 7
    return [source_char_indices.get(char, source_char_indices['<UNK>']) for char in text] + \
           [source_char_indices['<PAD>']] * (sequence_len - len(text))

def infer():
    source_data = get_corpus(source_data_path)
    target_data = get_corpus(target_data_path)

    # get mapping dict
    source_indices_char, source_char_indices = extract_char_vocab(source_data)
    target_indices_char, target_char_indices = extract_char_vocab(target_data)

    input_word = 'hello'
    text = source_2_seq(input_word, source_char_indices)
    loaded_graph = tf.Graph()
    with tf.Session(graph=loaded_graph) as sess:
        loader = tf.train.import_meta_graph(checkpoint + '.meta')
        loader.restore(sess, checkpoint)

        input_data = loaded_graph.get_tensor_by_name('inputs:0')
        logits = loaded_graph.get_tensor_by_name('predictions:0')
        source_sequence_len = loaded_graph.get_tensor_by_name('source_sequence_len:0')
        target_sequence_len = loaded_graph.get_tensor_by_name('target_sequence_len:0')
        answer_logits = sess.run(logits, {input_data: [text] * batch_size,
                                          target_sequence_len: [len(input_word)] * batch_size,
                                          source_sequence_len: [len(input_word)] * batch_size})[0]
    pad = source_char_indices['<PAD>']
    print('raw input:', input_word)
    print('\nSource')
    print(' Word 编号:    {}'.format([i for i in text]))
    print(' Input Words:    {}'.format(' '.join([source_indices_char[i] for i in text])))

    print('\nTarget')
    print(' Word 编号:    {}'.format([i for i in answer_logits if i != pad]))
    print(' Response Words: {}'.format(' '.join([target_indices_char[i] for i in answer_logits if i != pad])))

if __name__ == '__main__':
    # train()
    infer()


#!/usr/bin/env python
# Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

'''
Restricted CCSD implementation for real integrals.  Permutation symmetry for
the 4-index integrals (ij|kl) = (ij|lk) = (ji|kl) are assumed.

Note MO integrals are treated in chemist's notation
'''

import time
from functools import reduce
import numpy
from pyscf import lib
from pyscf.lib import logger
from pyscf.cc import ccsd
from pyscf.cc import _ccsd

# t2,l2 as ijab
def kernel(mycc, eris=None, t1=None, t2=None, l1=None, l2=None,
           max_cycle=50, tol=1e-8, verbose=logger.INFO,
           fintermediates=None, fupdate=None):
    if eris is None: eris = mycc.ao2mo()
    cput0 = (time.clock(), time.time())
    log = logger.new_logger(mycc, verbose)

    if t1 is None: t1 = mycc.t1
    if t2 is None: t2 = mycc.t2
    if l1 is None: l1 = t1
    if l2 is None: l2 = t2
    if fintermediates is None:
        fintermediates = make_intermediates
    if fupdate is None:
        fupdate = update_lambda

    imds = fintermediates(mycc, t1, t2, eris)

    if isinstance(mycc.diis, lib.diis.DIIS):
        adiis = mycc.diis
    elif mycc.diis:
        adiis = lib.diis.DIIS(mycc, mycc.diis_file, incore=mycc.incore_complete)
        adiis.space = mycc.diis_space
    else:
        adiis = None
    cput0 = log.timer('CCSD lambda initialization', *cput0)

    conv = False
    for istep in range(max_cycle):
        l1new, l2new = fupdate(mycc, t1, t2, l1, l2, eris, imds)
        normt = numpy.linalg.norm(mycc.amplitudes_to_vector(l1new, l2new) -
                                  mycc.amplitudes_to_vector(l1, l2))
        l1, l2 = l1new, l2new
        l1new = l2new = None
        l1, l2 = mycc.run_diis(l1, l2, istep, normt, 0, adiis)
        log.info('cycle = %d  norm(lambda1,lambda2) = %.6g', istep+1, normt)
        cput0 = log.timer('CCSD iter', *cput0)
        if normt < tol:
            conv = True
            break
    return conv, l1, l2

# l2, t2 as ijab
def make_intermediates(mycc, t1, t2, eris):
    log = logger.Logger(mycc.stdout, mycc.verbose)
    nocc, nvir = t1.shape
    nov = nocc * nvir
    nvir_pair = nvir*(nvir+1)//2
    foo = eris.fock[:nocc,:nocc]
    fov = eris.fock[:nocc,nocc:]
    fvo = eris.fock[nocc:,:nocc]
    fvv = eris.fock[nocc:,nocc:]

    class _IMDS: pass
    imds = _IMDS()
    #TODO: mycc.incore_complete
    imds.ftmp = lib.H5TmpFile()
    imds.woooo = imds.ftmp.create_dataset('woooo', (nocc,nocc,nocc,nocc), 'f8')
    imds.wvooo = imds.ftmp.create_dataset('wvooo', (nvir,nocc,nocc,nocc), 'f8')
    imds.wVOov = imds.ftmp.create_dataset('wVOov', (nvir,nocc,nocc,nvir), 'f8')
    imds.wvOOv = imds.ftmp.create_dataset('wvOOv', (nvir,nocc,nocc,nvir), 'f8')
    imds.wvvov = imds.ftmp.create_dataset('wvvov', (nvir,nvir,nocc,nvir), 'f8')

    w1 = fvv - numpy.einsum('ja,jb->ba', fov, t1)
    w2 = foo + numpy.einsum('ib,jb->ij', fov, t1)
    w3 = numpy.einsum('kc,jkbc->bj', fov, t2) * 2 + fov.T
    w3 -= numpy.einsum('kc,kjbc->bj', fov, t2)
    w3 += lib.einsum('kc,kb,jc->bj', fov, t1, t1)
    w4 = fov.copy()

    time1 = time.clock(), time.time()
    unit = nocc*nvir**2*6
    max_memory = max(0, mycc.max_memory - lib.current_memory()[0])
    blksize = min(nvir, max(ccsd.BLKMIN, int((max_memory*.95e6/8-nocc**4-nvir*nocc**3)/unit)))
    log.debug1('ccsd lambda make_intermediates: block size = %d, nvir = %d in %d blocks',
               blksize, nvir, int((nvir+blksize-1)//blksize))

    fswap = lib.H5TmpFile()
    for istep, (p0, p1) in enumerate(lib.prange(0, nvir, blksize)):
        eris_ovvv = eris.get_ovvv(slice(None), slice(p0,p1))
        fswap['vvov/%d'%istep] = eris_ovvv.transpose(2,3,0,1)

    woooo = 0
    wvooo = numpy.zeros((nvir,nocc,nocc,nocc))
    for p0, p1 in lib.prange(0, nvir, blksize):
        eris_ovvv = eris.get_ovvv(slice(None), slice(p0,p1))
        eris_vvov = numpy.empty(((p1-p0),nvir,nocc,nvir))
        for istep, (q0, q1) in enumerate(lib.prange(0, nvir, blksize)):
            eris_vvov[:,:,:,q0:q1] = fswap['vvov/%d'%istep][p0:p1]

        w1 += numpy.einsum('jcba,jc->ba', eris_ovvv, t1[:,p0:p1]*2)
        w1[:,p0:p1] -= numpy.einsum('jabc,jc->ba', eris_ovvv, t1)
        theta = t2[:,:,:,p0:p1] * 2 - t2[:,:,:,p0:p1].transpose(1,0,2,3)
        w3 += lib.einsum('jkcd,kdcb->bj', theta, eris_ovvv)
        theta = None
        wVOov = lib.einsum('jbcd,kd->bjkc', eris_ovvv, t1)
        wvOOv = lib.einsum('cbjd,kd->cjkb', eris_vvov,-t1)
        g2vovv = eris_vvov.transpose(0,2,1,3) * 2 - eris_vvov.transpose(0,2,3,1)
        for i0, i1 in lib.prange(0, nocc, blksize):
            tau = t2[:,i0:i1] + numpy.einsum('ia,jb->ijab', t1, t1[i0:i1])
            wvooo[p0:p1,i0:i1] += lib.einsum('cibd,jkbd->ckij', g2vovv, tau)
        g2vovv = tau = None

        # Watch out memory usage here, due to the t2 transpose
        wvvov  = lib.einsum('jabd,jkcd->abkc', eris_ovvv, t2) * -1.5
        wvvov += eris_vvov.transpose(0,3,2,1) * 2
        wvvov -= eris_vvov

        g2vvov = eris_vvov * 2 - eris_ovvv.transpose(1,2,0,3)
        for i0, i1 in lib.prange(0, nocc, blksize):
            theta = t2[i0:i1] * 2 - t2[i0:i1].transpose(0,1,3,2)
            vackb = lib.einsum('acjd,kjbd->ackb', g2vvov, theta)
            wvvov[:,:,i0:i1] += vackb.transpose(0,3,2,1)
            wvvov[:,:,i0:i1] -= vackb * .5
        g2vvov = eris_ovvv = eris_vvov = theta = None

        eris_ovoo = _cp(eris.ovoo[:,p0:p1])
        w2 += numpy.einsum('kbij,kb->ij', eris_ovoo, t1[:,p0:p1]) * 2
        w2 -= numpy.einsum('ibkj,kb->ij', eris_ovoo, t1[:,p0:p1])
        theta = t2[:,:,p0:p1].transpose(1,0,2,3) * 2 - t2[:,:,p0:p1]
        w3 -= lib.einsum('lckj,klcb->bj', eris_ovoo, theta)

        tmp = lib.einsum('lc,jcik->ijkl', t1[:,p0:p1], eris_ovoo)
        woooo += tmp
        woooo += tmp.transpose(1,0,3,2)
        theta = tmp = None

        wvOOv += lib.einsum('lbjk,lc->bjkc', eris_ovoo, t1)
        wVOov -= lib.einsum('jbkl,lc->bjkc', eris_ovoo, t1)
        wvooo[p0:p1] += eris_ovoo.transpose(1,3,2,0) * 2
        wvooo[p0:p1] -= eris_ovoo.transpose(1,0,2,3)
        wvooo -= lib.einsum('klbc,iblj->ckij', t2[:,:,p0:p1], eris_ovoo*1.5)

        g2ovoo = eris_ovoo * 2 - eris_ovoo.transpose(2,1,0,3)
        theta = t2[:,:,:,p0:p1]*2 - t2[:,:,:,p0:p1].transpose(1,0,2,3)
        vcjik = lib.einsum('jlcb,lbki->cjki', theta, g2ovoo)
        wvooo += vcjik.transpose(0,3,2,1)
        wvooo -= vcjik*.5
        theta = g2ovoo = None

        eris_voov = _cp(eris.ovvo[:,p0:p1]).transpose(1,0,3,2)
        tau = t2[:,:,p0:p1] + numpy.einsum('ia,jb->ijab', t1[:,p0:p1], t1)
        woooo += lib.einsum('cijd,klcd->ijkl', eris_voov, tau)
        tau = None

        g2voov = eris_voov*2 - eris_voov.transpose(0,2,1,3)
        tmpw4 = numpy.einsum('ckld,ld->kc', g2voov, t1)
        w1 -= lib.einsum('ckja,kjcb->ba', g2voov, t2[:,:,p0:p1])
        w1[:,p0:p1] -= numpy.einsum('ja,jb->ba', tmpw4, t1)
        w2 += lib.einsum('jkbc,bikc->ij', t2[:,:,p0:p1], g2voov)
        w2 += numpy.einsum('ib,jb->ij', tmpw4, t1[:,p0:p1])
        w3 += reduce(numpy.dot, (t1.T, tmpw4, t1[:,p0:p1].T))
        w4[:,p0:p1] += tmpw4

        wvOOv += lib.einsum('bljd,kd,lc->bjkc', eris_voov, t1, t1)
        wVOov -= lib.einsum('bjld,kd,lc->bjkc', eris_voov, t1, t1)

        VOov  = lib.einsum('bjld,klcd->bjkc', g2voov, t2)
        VOov -= lib.einsum('bjld,kldc->bjkc', eris_voov, t2)
        VOov += eris_voov
        vOOv = lib.einsum('bljd,kldc->bjkc', eris_voov, t2)
        vOOv -= _cp(eris.oovv[:,:,p0:p1]).transpose(2,1,0,3)
        wVOov += VOov
        wvOOv += vOOv
        imds.wVOov[p0:p1] = wVOov
        imds.wvOOv[p0:p1] = wvOOv
        wOVov = wOvOv = None

        ov1 = vOOv*2 + VOov
        ov2 = VOov*2 + vOOv
        vOOv = VOov = None
        wvooo -= lib.einsum('jb,bikc->ckij', t1[:,p0:p1], ov1)
        wvooo += lib.einsum('kb,bijc->ckij', t1[:,p0:p1], ov2)
        w3 += numpy.einsum('ckjb,kc->bj', ov2, t1[:,p0:p1])

        wvvov += lib.einsum('ajkc,jb->abkc', ov1, t1)
        wvvov -= lib.einsum('ajkb,jc->abkc', ov2, t1)

        eris_ovoo = _cp(eris.ovoo[:,p0:p1])
        g2ovoo = eris_ovoo * 2 - eris_ovoo.transpose(2,1,0,3)
        tau = t2 + numpy.einsum('ia,jb->ijab', t1, t1)
        wvvov += lib.einsum('laki,klbc->abic', g2ovoo, tau)
        imds.wvvov[p0:p1] = wvvov
        wvvov = ov1 = ov2 = g2ovoo = None

    woooo += _cp(eris.oooo).transpose(0,2,1,3)
    imds.woooo[:] = woooo
    imds.wvooo[:] = wvooo
    woooo = wvooo = None

    w3 += numpy.einsum('bc,jc->bj', w1, t1)
    w3 -= numpy.einsum('kj,kb->bj', w2, t1)

    fswap = None

    imds.w1 = w1
    imds.w2 = w2
    imds.w3 = w3
    imds.w4 = w4
    imds.ftmp.flush()
    return imds

# update L1, L2
def update_lambda(mycc, t1, t2, l1, l2, eris=None, imds=None):
    if imds is None: imds = make_intermediates(mycc, t1, t2, eris)
    time1 = time0 = time.clock(), time.time()
    log = logger.Logger(mycc.stdout, mycc.verbose)
    nocc, nvir = t1.shape
    nov = nocc * nvir
    fov = eris.fock[:nocc,nocc:]

    theta = t2*2 - t2.transpose(0,1,3,2)
    mba = lib.einsum('klca,klcb->ba', l2, theta)
    mij = lib.einsum('ikcd,jkcd->ij', l2, theta)
    theta = None
    mba1 = numpy.einsum('jc,jb->bc', l1, t1) + mba
    mij1 = numpy.einsum('kb,jb->kj', l1, t1) + mij
    mia1 = t1 + numpy.einsum('kc,jkbc->jb', l1, t2) * 2
    mia1 -= numpy.einsum('kc,jkcb->jb', l1, t2)
    mia1 -= reduce(numpy.dot, (t1, l1.T, t1))
    mia1 -= numpy.einsum('bd,jd->jb', mba, t1)
    mia1 -= numpy.einsum('lj,lb->jb', mij, t1)

    l2new = mycc._add_vvvv(None, l2, eris, with_ovvv=False, t2sym='jiba')
    l1new  = numpy.einsum('ijab,jb->ia', l2new, t1) * 2
    l1new -= numpy.einsum('jiab,jb->ia', l2new, t1)
    l2new *= .5  # *.5 because of l2+l2.transpose(1,0,3,2) in the end
    tmp = tmp1 = None

    l1new += fov
    l1new += numpy.einsum('ib,ba->ia', l1, imds.w1)
    l1new -= numpy.einsum('ja,ij->ia', l1, imds.w2)
    l1new -= numpy.einsum('ik,ka->ia', mij, imds.w4)
    l1new -= numpy.einsum('ca,ic->ia', mba, imds.w4)
    l1new += numpy.einsum('ijab,bj->ia', l2, imds.w3) * 2
    l1new -= numpy.einsum('ijba,bj->ia', l2, imds.w3)

    l2new += numpy.einsum('ia,jb->ijab', l1, imds.w4)
    l2new += lib.einsum('jibc,ca->jiba', l2, imds.w1)
    l2new -= lib.einsum('jk,kiba->jiba', imds.w2, l2)

    eris_ovoo = _cp(eris.ovoo)
    l1new -= numpy.einsum('iajk,kj->ia', eris_ovoo, mij1) * 2
    l1new += numpy.einsum('jaik,kj->ia', eris_ovoo, mij1)
    l2new -= lib.einsum('jbki,ka->jiba', eris_ovoo, l1)
    eris_ovoo = None

    tau = _ccsd.make_tau(t2, t1, t1)
    l2tau = lib.einsum('ijcd,klcd->ijkl', l2, tau)
    tau = None
    l2t1 = lib.einsum('jidc,kc->ijkd', l2, t1)

    max_memory = max(0, mycc.max_memory - lib.current_memory()[0])
    unit = nocc*nvir**2*5
    blksize = min(nocc, max(ccsd.BLKMIN, int(max_memory*.95e6/8/unit)))
    log.debug1('block size = %d, nocc = %d is divided into %d blocks',
               blksize, nocc, int((nocc+blksize-1)/blksize))

    l1new -= numpy.einsum('jb,jiab->ia', l1, _cp(eris.oovv))
    for p0, p1 in lib.prange(0, nvir, blksize):
        eris_ovvv = eris.get_ovvv(slice(None), slice(p0,p1))
        l1new[:,p0:p1] += numpy.einsum('iabc,bc->ia', eris_ovvv, mba1) * 2
        l1new -= numpy.einsum('ibca,bc->ia', eris_ovvv, mba1[p0:p1])
        l2new[:,:,p0:p1] += lib.einsum('jbac,ic->jiba', eris_ovvv, l1)
        m4 = lib.einsum('ijkd,kadb->ijab', l2t1, eris_ovvv)
        l2new[:,:,p0:p1] -= m4
        l1new[:,p0:p1] -= numpy.einsum('ijab,jb->ia', m4, t1) * 2
        l1new -= numpy.einsum('ijab,ia->jb', m4, t1[:,p0:p1]) * 2
        l1new[:,p0:p1] += numpy.einsum('jiab,jb->ia', m4, t1)
        l1new += numpy.einsum('jiab,ia->jb', m4, t1[:,p0:p1])
        eris_ovvv = m4buf = m4 = None

        eris_voov = _cp(eris.ovvo[:,p0:p1].transpose(1,0,3,2))
        l1new[:,p0:p1] += numpy.einsum('jb,aijb->ia', l1, eris_voov) * 2
        l2new[:,:,p0:p1] += eris_voov.transpose(1,2,0,3) * .5
        l2new[:,:,p0:p1] -= lib.einsum('bjic,ca->jiba', eris_voov, mba1)
        l2new[:,:,p0:p1] -= lib.einsum('bjka,ik->jiba', eris_voov, mij1)
        l1new[:,p0:p1] += numpy.einsum('aijb,jb->ia', eris_voov, mia1) * 2
        l1new -= numpy.einsum('bija,jb->ia', eris_voov, mia1[:,p0:p1])
        m4 = lib.einsum('ijkl,aklb->ijab', l2tau, eris_voov)
        l2new[:,:,p0:p1] += m4 * .5
        l1new[:,p0:p1] += numpy.einsum('ijab,jb->ia', m4, t1) * 2
        l1new -= numpy.einsum('ijba,jb->ia', m4, t1[:,p0:p1])

        saved_wvooo = _cp(imds.wvooo[p0:p1])
        l1new -= lib.einsum('ckij,jkca->ia', saved_wvooo, l2[:,:,p0:p1])
        saved_wvovv = _cp(imds.wvvov[p0:p1])
        # Watch out memory usage here, due to the l2 transpose
        l1new[:,p0:p1] += lib.einsum('abkc,kibc->ia', saved_wvovv, l2)
        saved_wvooo = saved_wvovv = None

        saved_wvOOv = _cp(imds.wvOOv[p0:p1])
        tmp_voov = _cp(imds.wVOov[p0:p1]) * 2
        tmp_voov += saved_wvOOv
        tmp = l2.transpose(0,2,1,3) - l2.transpose(0,3,1,2)*.5
        l2new[:,:,p0:p1] += lib.einsum('iakc,bjkc->jiba', tmp, tmp_voov)
        tmp = tmp1 = tmp_ovov = None

        tmp = lib.einsum('jkca,bikc->jiba', l2, saved_wvOOv)
        l2new[:,:,p0:p1] += tmp
        l2new[:,:,p0:p1] += tmp.transpose(1,0,2,3) * .5
        saved_wvOOv = tmp = None

    saved_woooo = _cp(imds.woooo)
    m3 = lib.einsum('ijkl,klab->ijab', saved_woooo, l2)
    l2new += m3 * .5
    l1new += numpy.einsum('ijab,jb->ia', m3, t1) * 2
    l1new -= numpy.einsum('ijba,jb->ia', m3, t1)
    saved_woooo = m3 = None
    #time1 = log.timer_debug1('lambda pass [%d:%d]'%(p0, p1), *time1)

    mo_e = eris.fock.diagonal()
    eia = lib.direct_sum('i-a->ia', mo_e[:nocc], mo_e[nocc:])
    l1new /= eia
    l1new += l1

#    l2new = l2new + l2new.transpose(1,0,3,2)
#    l2new /= lib.direct_sum('ia+jb->ijab', eia, eia)
#    l2new += l2
    ij = 0
    for i in range(nocc):
        if i > 0:
            l2new[i,:i] += l2new[:i,i].transpose(0,2,1)
            l2new[i,:i] /= lib.direct_sum('a,jb->jab', eia[i], eia[:i])
            l2new[:i,i] = l2new[i,:i].transpose(0,2,1)
        l2new[i,i] = l2new[i,i] + l2new[i,i].T
        l2new[i,i] /= lib.direct_sum('a,b->ab', eia[i], eia[i])
    l2new += l2

    time0 = log.timer_debug1('update l1 l2', *time0)
    return l1new, l2new

def _cp(a):
    return numpy.array(a, copy=False, order='C')

if __name__ == '__main__':
    from pyscf import gto
    from pyscf import scf
    from pyscf import ao2mo
    from pyscf.cc import ccsd

    mol = gto.M()
    mf = scf.RHF(mol)

    mcc = ccsd.CCSD(mf)

    numpy.random.seed(12)
    nocc = 5
    nmo = 12
    nvir = nmo - nocc
    eri0 = numpy.random.random((nmo,nmo,nmo,nmo))
    eri0 = ao2mo.restore(1, ao2mo.restore(8, eri0, nmo), nmo)
    fock0 = numpy.random.random((nmo,nmo))
    fock0 = fock0 + fock0.T + numpy.diag(range(nmo))*2
    t1 = numpy.random.random((nocc,nvir))
    t2 = numpy.random.random((nocc,nocc,nvir,nvir))
    t2 = t2 + t2.transpose(1,0,3,2)
    l1 = numpy.random.random((nocc,nvir))
    l2 = numpy.random.random((nocc,nocc,nvir,nvir))
    l2 = l2 + l2.transpose(1,0,3,2)

    eris = ccsd._ChemistsERIs()
    eris.oooo = eri0[:nocc,:nocc,:nocc,:nocc].copy()
    eris.ovoo = eri0[:nocc,nocc:,:nocc,:nocc].copy()
    eris.oovv = eri0[:nocc,:nocc,nocc:,nocc:].copy()
    eris.ovvo = eri0[:nocc,nocc:,nocc:,:nocc].copy()
    idx = numpy.tril_indices(nvir)
    eris.ovvv = eri0[:nocc,nocc:,nocc:,nocc:][:,:,idx[0],idx[1]].copy()
    eris.vvvv = ao2mo.restore(4,eri0[nocc:,nocc:,nocc:,nocc:],nvir)
    eris.fock = fock0

    imds = make_intermediates(mcc, t1, t2, eris)
    l1new, l2new = update_lambda(mcc, t1, t2, l1, l2, eris, imds)
    print(lib.finger(l1new) - -6699.5335665027187)
    print(lib.finger(l2new) - -514.7001243502192 )
    print(abs(l2new-l2new.transpose(1,0,3,2)).sum())

    mcc.max_memory = 0
    imds = make_intermediates(mcc, t1, t2, eris)
    l1new, l2new = update_lambda(mcc, t1, t2, l1, l2, eris, imds)
    print(lib.finger(l1new) - -6699.5335665027187)
    print(lib.finger(l2new) - -514.7001243502192 )
    print(abs(l2new-l2new.transpose(1,0,3,2)).sum())

    mol = gto.Mole()
    mol.verbose = 0
    mol.atom = [
        [8 , (0. , 0.     , 0.)],
        [1 , (0. , -0.757 , 0.587)],
        [1 , (0. , 0.757  , 0.587)]]

    mol.basis = 'cc-pvdz'
    mol.build()
    rhf = scf.RHF(mol)
    rhf.conv_tol = 1e-16
    rhf.scf()

    mcc = ccsd.CCSD(rhf)
    mcc.conv_tol = 1e-12
    ecc, t1, t2 = mcc.kernel()

    nmo = rhf.mo_energy.size
    fock0 = numpy.diag(rhf.mo_energy)
    nocc = mol.nelectron // 2
    nvir = nmo - nocc

    eris = mcc.ao2mo()
    conv, l1, l2 = kernel(mcc, eris, t1, t2, tol=1e-8)
    print(numpy.linalg.norm(l1)-0.0132626841292)
    print(numpy.linalg.norm(l2)-0.212575609057)

    from pyscf.cc import ccsd_rdm
    dm1 = ccsd_rdm.make_rdm1(mcc, t1, t2, l1, l2)
    dm2 = ccsd_rdm.make_rdm2(mcc, t1, t2, l1, l2)
    h1 = reduce(numpy.dot, (rhf.mo_coeff.T, rhf.get_hcore(), rhf.mo_coeff))
    eri = ao2mo.full(rhf._eri, rhf.mo_coeff)
    eri = ao2mo.restore(1, eri, nmo).reshape((nmo,)*4)
    e1 = numpy.einsum('pq,pq', h1, dm1)
    e2 = numpy.einsum('pqrs,pqrs', eri, dm2) * .5
    print(e1+e2+mol.energy_nuc() - rhf.e_tot - ecc)


"""Implements self supervised semantic shift functions
It uses poisoning attacks to learn landmarks in a self-supervised way
At each iteration, generate perturbation on the data, generating positive
and negative samples
Learn the separation between them (using any classifier)
Apply the classifier to the original (non-perturbated) data
Negatives -> landmarks
Positives -> semantically changed
We can begin by aligning on all words, and then learn better landmarks from
there. Alternatively, one can start from random landmarks."""

# Third party modules
import numpy as np
import tensorflow as tf
from tensorflow import keras
from sklearn.svm import SVC, LinearSVC
from sklearn.metrics import accuracy_score, log_loss
from scipy.spatial.distance import cosine, euclidean
import matplotlib.pyplot as plt
import seaborn as sns

# Local modules
from WordVectors import WordVectors, intersection
from alignment import align

# Initialize random seeds
np.random.seed(1)
tf.random.set_seed(1)

def negative_samples(words, size, p=None):
    """
    Returns negative samples of semantic change
    May use distribution of cosine distance as sampling distribution
    """
    neg_samples = np.random.choice(words, size, p=p)

    return neg_samples

def inject_change_single(wv, w, words, v_a, alpha, replace=False,
                         max_tries=50):
    """
    Injects change to word w in wv by randomly selecting a word t in wv
    and injecting the sense of t in to w.
    The modified vector of w must have a higher cosine distance to v_a
    than its original version. This is done by sampling t while the cosine of
    w is not greater than that of v_a and wv(w) or until a max_tries.
    v_a is the vector of word w in the parallel corpus (not wv).

    Arguments:
            wv      -   WordVectors of the corpus to be modified
            w       -   (str) Word to be modified
            words   -   (list) Pool of words to sample from, injecting sense
            v_a     -   (np.ndarray) word vector of w in the source parallel to wv
            alpha   -   (float) Rate of injected change
            replace -   (bool) Whether to replace w with t instead of 'moving' w towards t
    Returns:
            x       -   (np.ndarray) modified vector of w
    """
    cos_t = cosine(v_a, wv[w])  # cosine distance threshold we want to surpass

    c = 0
    tries = 0
    w_id = wv.word_id[w]
    v_b = np.copy(wv.vectors[w_id])
    while c < cos_t and tries < max_tries:
        tries += 1
        selected = np.random.choice(words)  # select word with new sense
        if not replace:
            b = wv[w] + alpha*wv[selected]
            v_b = b
        else:
            v_b = wv[selected]

        c = cosine(v_a, v_b)

    return v_b

def inject_change_batch(wv, changes, alpha, replace=True):
    """
    Given a WordVectors object and a list of words, perform fast injection
    of semantic change by using the update rule from Word2Vec
    wv - WordVectors (input)
    changes - list of n tuples (a, b) that drives the change such that b->a
          i.e.: simulates using b in the contexts of a
    alpha - degree in which to inject the change
              if scalar: apply same alpha to every pair
              if array-like: requires size n, specifies individual alpha values
                              for each pair
    replace  - (bool) if True, words are replaced instead of moved
                e.g.: if pair is (dog, car), then v_car <- v_dog
    Returns a WordVectors object with the change
    """
    wv_new = WordVectors(words=wv.words, vectors=np.copy(wv.vectors))
    for i, pair in enumerate(changes):
        t, w = pair
        t_i = wv.word_id[t]  # target word
        w_i = wv.word_id[w]  # modified word
        # Update vector with alpha and score
        # Higher score means vectors are already close, thus apply less change
        # Alpha controls the rate of change
        if not replace:
            b = wv_new[w] + alpha*(1)*wv[t]
            wv_new.vectors[w_i] = b
        else:
            wv_new.vectors[w_i] = wv[t]
        # print("norm b", np.linalg.norm(b))
    return wv_new

def get_features(x, names=["cos"]):
    """
    Compute features given input training data (concatenated vectors)
    Default features is cosine. Accepted features: cosine (cos).
    Attributes:
            x   - size n input training data as concatenated word vectors
            names - size d list of features to compute
    Returns:
            n x d feature matrix (floats)
    """
    x_out = np.zeros((len(x), len(names)), dtype=float)
    for i, p in enumerate(x):
        for j, feat in enumerate(names):
            if feat == "cos":
                x_ = cosine(p[:len(p)//2], p[len(p)//2:])
                x_out[i][j] = x_
    return x_out

def build_sklearn_model():
    """
    Build SVM using sklearn model
    The model uses an RBF kernel and the features are given by difference
    between input vectors u-v.
    Return: sklearn SVC
    """
    model = SVC(random_state=0, probability=True)
    return model

def build_keras_model(dim):
    """
    Builds the keras model to be used in self-supervision.
    Return: Keras-Tensorflow2 model
    """
    h1_dim = 100
    h2_dim = 100
    model = keras.Sequential([
                             keras.layers.Input(shape=(dim)),
                             keras.layers.Dense(h1_dim, activation="relu",
                                                activity_regularizer=keras.regularizers.l2(1e-2)),
                             # keras.layers.Dense(h2_dim, activation="relu",
                             #                    activity_regularizer=keras.regularizers.l2(1e-2)),
                             keras.layers.Dense(1, activation="sigmoid")
                            ])
    model.compile(optimizer="rmsprop",
                  loss="binary_crossentropy",
                  metrics=["accuracy"])
    return model

def threshold_crossvalidation(wv1, wv2, iters=100,
                                        n_fold=1,
                                        n_targets=100,
                                        n_negatives=100,
                                        fast=True,
                                        rate=0.5,
                                        t=0.5,
                                        landmarks=None,
                                        t_overlap=1,
                                        debug=False):
    """
    Runs crossvalidation over self-supervised samples, carrying out a model
    selection to determine the best cosine threshold to use in the final
    prediction.

    Arguments:
        wv1, wv2    - input WordVectors - required to be intersected and ALIGNED before call
        plot        - 1: plot functions in the end 0: do not plot
        iters       - max no. of iterations
        n_fold      - n-fold crossvalidation (1 - leave one out, 10 - 10-fold cv, etc.)
        n_targets   - number of positive samples to generate
        n_negatives - number of negative samples
        fast        - use fast semantic change simulation
        rate        - rate of semantic change injection
        t           - classificaiton threshold (0.5)
        t_overlap   - overlap threshold for (stop criterion)
        landmarks   - list of words to use as landmarks (classification only)
        debug       - toggles debugging mode on/off. Provides reports on several metrics. Slower.
    Returns:
        t - selected cosine threshold t
    """

    wv2_original = WordVectors(words=wv2.words, vectors=wv2.vectors.copy())
    landmark_set = set(landmarks)
    non_landmarks = [w for w in wv1.words if w not in landmark_set]

    for iter in range(iters):

        replace = dict()  # replacement dictionary
        pos_samples = list()
        pos_vectors = dict()

        # Randomly sample words to inject change to
        # If no word is flagged as non_landmarks, sample from all words
        # In practice, this should never occur when selecting landmarks
        # but only for classification when aligning on all words
        if len(non_landmarks) > 0:
            targets = np.random.choice(non_landmarks, n_targets)
            # Make targets deterministic
            #targets = non_landmarks
        else:
            targets = np.random.choice(wv1.words, n_targets)

        for target in targets:

            # Simulate semantic change in target word
            v = inject_change_single(wv2_original, target, wv1.words,
                                     wv1[target], rate)

            pos_vectors[target] = v

            pos_samples.append(target)
        # Convert to numpy array
        pos_samples = np.array(pos_samples)
        # Get negative samples from landmarks
        neg_samples = negative_samples(landmarks, n_negatives, p=None)
        neg_vectors = {w: wv2_original[w] for w in neg_samples}
        # Create dictionary of supervision samples (positive and negative)
        # Mapping word -> vector
        sup_vectors = {**neg_vectors, **pos_vectors}

        # Prepare training data
        words_train = np.concatenate((pos_samples, neg_samples))
        # assign labels to positive and negative samples
        y_train = [1] * len(pos_samples) + [0] * len(neg_samples)

        # Stack columns to shuffle data and labels together
        train = np.column_stack((words_train, y_train))
        # Shuffle batch
        np.random.shuffle(train)
        # Detach data and labels
        words_train = train[:, 0]
        y_train = train[:, -1].astype(int)

        # Calculate cosine distance of training samples
        x_train = np.array([cosine(wv1[w], sup_vectors[w]) for w in words_train])

        # t_pool = [0.2, 0.7]
        t_pool = np.arange(0.2, 1, 0.1)

        best_acc = 0
        best_t = 0
        for t_ in t_pool:
            acc = 0
            for i in range(0, len(x_train), n_fold):
                x_cv = x_train[i:i+n_fold]
                y_true = y_train[i:i+n_fold]
                y_hat = x_cv > t_
                acc += sum(y_hat == y_true)/len(x_cv)
            acc = acc/(len(x_train)//n_fold)
            if acc > best_acc:
                best_acc = acc
                best_t = t_
                print("- New best t", t_, acc)

    return best_t

def s4(wv1, wv2, verbose=0, plot=0, cls_model="nn",
                              iters=100,
                              n_targets=10,
                              n_negatives=10,
                              fast=True,
                              rate=0,
                              t=0.5,
                              t_overlap=1,
                              landmarks=None,
                              update_landmarks=True,
                              return_model=False,
                              debug=False):
    """
    Performs self-supervised learning of semantic change.
    Generates negative samples by sampling from landmarks.
    Generates positive samples via simulation of semantic change on random non-landmark words.
    Trains a classifier, fine-tune it across multiple iterations.
    If update_landmarks is True, then it learns landmarks from that step. In this case,
    the returned values are landmarks, non_landmarks, Q (transform matrix)
    Otherwise, landmarks are fixed from a starting set and the returned value
    is the learned classifier - landmarks must be passed.
    Arguments:
        wv1, wv2    - input WordVectors - required to be intersected before call
        verbose     - 1: display log, 0: quiet
        plot        - 1: plot functions in the end 0: do not plot
        cls_model   - classification model to use {"nn", "svm_auto", "svm_features"}
        iters       - max no. of iterations
        n_targets   - number of positive samples to generate
        n_negatives - number of negative samples
        fast        - use fast semantic change simulation
        rate        - rate of semantic change injection
        t           - classificaiton threshold (0.5)
        t_overlap   - overlap threshold for (stop criterion)
        landmarks   - list of words to use as landmarks (classification only)
        update_landmarks - if True, learns landmarks. Otherwise, learns classification model.
        debug       - toggles debugging mode on/off. Provides reports on several metrics. Slower.
    Returns:
        if update_landmarks is True:
            landmarks - list of landmark words
            non_landmarks - list of non_landmark words
            Q           - transformation matrix for procrustes alignment
        if update_landmarks is False:
            model       - binary classifier
    """

    # Define verbose prints
    if verbose==1:
        def verbose_print(*s, end="\n"):
            print(*s, end=end)
    elif verbose==0:
        def verbose_print(*s, end="\n"):
            return None

    wv2_original = WordVectors(words=wv2.words, vectors=wv2.vectors.copy())

    avg_window = 0  # number of iterations to use in running average

    # Begin alignment
    if update_landmarks:
        # Check if landmarks is initialized
        if landmarks == None:
            wv1, wv2, Q = align(wv1, wv2)  # start form global alignment
            landmark_dists = [euclidean(u, v) for u, v in zip(wv1.vectors, wv2.vectors)]
            landmark_args = np.argsort(landmark_dists)
            landmarks = [wv1.words[i] for i in landmark_args[:int(len(wv1.words)*0.5)]]
            # landmarks = np.random.choice(wv1.words, int(len(wv1)*0.5))
        landmark_set = set(landmarks)
        non_landmarks = np.array([w for w in wv1.words if w not in landmark_set])
    else:
        landmark_set = set(landmarks)
        non_landmarks = [w for w in wv1.words if w not in landmark_set]

    wv1, wv2, Q = align(wv1, wv2, anchor_words=landmarks)

    if cls_model == "nn":
        model = build_keras_model(wv1.dimension*2)
    elif cls_model == "svm_auto" or cls_model == "svm_features":
        model = build_sklearn_model()  # get SVC

    landmark_hist = list()  # store no. of landmark history
    loss_hist = list()  # store self-supervision loss history
    alignment_loss_hist = list()  # store landmark alignment loss
    alignment_out_hist = list()  # store alignment loss outside of lm
    alignment_all_hist = list()

    cumulative_out_hist = list()
    cumulative_alignment_hist = list()  # store cumulative loss alignment
    overlap_hist = list()  # store landmark overlap history
    cumulative_overlap_hist = list()  # mean overlap history
    cumulative_loss = 0

    # History of cosines
    cos_loss_in_hist = list()
    cos_loss_out_hist = list()
    cumulative_cos_in = list()
    cumulative_cos_out = list()

    prev_landmarks = set(landmarks)
    for iter in range(iters):

        replace = dict()  # replacement dictionary
        pos_samples = list()
        pos_vectors = dict()

        # Randomly sample words to inject change to
        # If no word is flagged as non_landmarks, sample from all words
        # In practice, this should never occur when selecting landmarks
        # but only for classification when aligning on all words
        if len(non_landmarks) > 0:
            targets = np.random.choice(non_landmarks, n_targets)
            # Make targets deterministic
            #targets = non_landmarks
        else:
            targets = np.random.choice(wv1.words, n_targets)

        for target in targets:

            # Simulate semantic change in target word
            v = inject_change_single(wv2_original, target, wv1.words,
                                     wv1[target], rate)

            pos_vectors[target] = v

            pos_samples.append(target)
        # Convert to numpy array
        pos_samples = np.array(pos_samples)
        # Get negative samples from landmarks
        neg_samples = negative_samples(landmarks, n_negatives, p=None)
        neg_vectors = {w: wv2_original[w] for w in neg_samples}
        # Create dictionary of supervision samples (positive and negative)
        # Mapping word -> vector
        sup_vectors = {**neg_vectors, **pos_vectors}

        # Prepare training data
        words_train = np.concatenate((pos_samples, neg_samples))
        # assign labels to positive and negative samples
        y_train = [1] * len(pos_samples) + [0] * len(neg_samples)

        # Stack columns to shuffle data and labels together
        train = np.column_stack((words_train, y_train))
        # Shuffle batch
        np.random.shuffle(train)
        # Detach data and labels
        words_train = train[:, 0]
        y_train = train[:, -1].astype(int)

        x_train = np.array([np.append(wv1[w], sup_vectors[w]) for w in words_train])

        # Append history
        landmark_hist.append(len(landmarks))
        v1_land = np.array([wv1[w] for w in landmarks])
        v2_land = np.array([wv2_original[w] for w in landmarks])
        v1_out = np.array([wv1[w] for w in non_landmarks])
        v2_out = np.array([wv2_original[w] for w in non_landmarks])

        alignment_loss = np.linalg.norm(v1_land-v2_land)**2/len(v1_land)
        alignment_loss_hist.append(alignment_loss)
        cumulative_alignment_hist.append(np.mean(alignment_loss_hist[-avg_window:]))

        # out loss
        alignment_out_loss = np.linalg.norm(v1_out-v2_out)**2/len(v1_out)
        alignment_out_hist.append(alignment_out_loss)
        cumulative_out_hist.append(np.mean(alignment_out_hist[-avg_window:]))

        # all loss
        alignment_all_loss = np.linalg.norm(wv1.vectors-wv2_original.vectors)**2/len(wv1.words)
        alignment_all_hist.append(alignment_all_loss)

        if debug:
        # cosine loss
            cos_in = np.mean([cosine(u, v) for u, v in zip (v1_land, v2_land)])
            cos_out = np.mean([cosine(u, v) for u, v in zip(v1_out, v2_out)])
            cos_loss_in_hist.append(cos_in)
            cos_loss_out_hist.append(cos_out)
            cumulative_cos_in.append(np.mean(cos_loss_in_hist))
            cumulative_cos_out.append(np.mean(cos_loss_out_hist))

        # Begin training of neural network
        if cls_model == "nn":
            history = model.train_on_batch(x_train, y_train, reset_metrics=False)
            # history = model.fit(x_train, y_train, epochs=5, verbose=0)
            # history = [history.history["loss"][0]]
        elif cls_model == "svm_auto":
            model.fit(x_train, y_train)
            pred_train = model.predict_proba(x_train)
            history = [log_loss(y_train, pred_train)]
        elif cls_model == "svm_features":
            x_train_ = get_features(x_train)  # retrieve manual features
            model.fit(x_train_, y_train)
            pred_train = model.predict_proba(x_train_)
            y_hat_t = (pred_train[:, 0] > 0.5)
            acc_t = accuracy_score(y_train, y_hat_t)
            history = [log_loss(y_train, pred_train), acc_t]

        loss_hist.append(history[0])

        # Apply model on original data to select landmarks
        x_real = np.array([np.append(u, v) for u, v
                            in zip(wv1.vectors, wv2_original.vectors)])
        if cls_model == "nn":
            predict_real = model.predict(x_real)
        elif cls_model == "svm_auto":
            predict_real = model.predict_proba(x_real)
            predict_real = predict_real[:, 1]
        elif cls_model == "svm_features":
            x_real_ = get_features(x_real)
            predict_real = model.predict_proba(x_real_)
            predict_real = predict_real[:, 1]

        y_predict = (predict_real>t)

        if update_landmarks:
            landmarks = [wv1.words[i] for i in range(len(wv1.words)) if predict_real[i]<t]
            non_landmarks = [wv1.words[i] for i in range(len(wv1.words)) if predict_real[i]>t]

        # Update landmark overlap using Jaccard Index
        isect_ab = set.intersection(prev_landmarks, set(landmarks))
        union_ab = set.union(prev_landmarks, set(landmarks))
        j_index = len(isect_ab)/len(union_ab)
        overlap_hist.append(j_index)

        cumulative_overlap_hist.append(np.mean(overlap_hist[-avg_window:]))  # store mean

        prev_landmarks = set(landmarks)

        verbose_print("> %3d | L %4d | l(in): %.2f | l(out): %.2f | loss: %.2f | overlap %.2f | acc: %.2f" %
                        (iter, len(landmarks), cumulative_alignment_hist[-1],
                         cumulative_out_hist[-1], history[0], cumulative_overlap_hist[-1], history[1]),
                         end="\r")

        wv1, wv2_original, Q = align(wv1, wv2_original, anchor_words=landmarks)

        # Check if overlap difference is below threhsold
        if np.mean(overlap_hist) > t_overlap:
            break

    # Print new line
    verbose_print()

    if plot == 1:
        iter += 1  # add one to iter for plotting
        plt.plot(range(iter), landmark_hist, label="landmarks")
        plt.hlines(len(wv1.words), 0, iter, colors="red")
        plt.ylabel("No. of landmarks")
        plt.xlabel("Iteration")
        plt.show()
        plt.plot(range(iter), loss_hist, c="red", label="loss")
        plt.ylabel("Loss (binary crossentropy)")
        plt.xlabel("Iteration")
        plt.legend()
        plt.show()
        plt.plot(range(iter), cumulative_alignment_hist, label="in (landmarks)")
        plt.plot(range(iter), cumulative_out_hist, label="out")
        plt.plot(range(iter), alignment_all_hist, label="all")
        plt.ylabel("Alignment loss (MSE)")
        plt.xlabel("Iteration")
        plt.legend()
        plt.show()

        if debug:
            plt.plot(range(iter), cumulative_cos_in, label="cos in")
            plt.plot(range(iter), cumulative_cos_out, label="cos out")
            plt.legend()
            plt.show()

        plt.plot(range(iter), cumulative_overlap_hist, label="overlap")

        plt.ylabel("Jaccard Index", fontsize=16)
        plt.xlabel("Iteration", fontsize=16)
        plt.xticks(fontsize=16)
        plt.yticks(fontsize=16)
        # plt.legend()
        plt.tight_layout()
        plt.savefig("overlap.pdf", format="pdf")
        #plt.show()

    if update_landmarks:
        if not return_model:
            return landmarks, non_landmarks, Q
        else:
            return landmarks, non_landmarks, Q, model
    else:
        return model

def main():
    """
    Runs main experiments using self supervised alignment.
    """
    # wv_source = "wordvectors/latin/corpus1/0.vec"
    # wv_target = "wordvectors/latin/corpus2/0.vec"
    # wv_source = "wordvectors/source/theguardianuk.vec"
    # wv_target = "wordvectors/source/thenewyorktimes_1.vec"
    wv_source = "wordvectors/semeval/latin-corpus1.vec"
    wv_target = "wordvectors/semeval/latin-corpus2.vec"
    # wv_source = "wordvectors/usuk/bnc.vec"
    # wv_target = "wordvectors/usuk/coca_mag.vec"
    # wv_source = "wordvectors/artificial/NYT-0.vec"
    # wv_target = "wordvectors/artificial/NYT-500_random.vec"
    plt.style.use("seaborn")

    # Read WordVectors
    normalized = False
    wv1 = WordVectors(input_file=wv_source, normalized=normalized)
    wv2 = WordVectors(input_file=wv_target, normalized=normalized)

    wv1, wv2 = intersection(wv1, wv2)

    landmarks, non_landmarks, Q = s4(wv1, wv2,
                                                            cls_model="nn",
                                                            n_targets=100,
                                                            n_negatives=100,
                                                            rate=1,
                                                            t=0.5,
                                                            iters=100,
                                                            verbose=1,
                                                            plot=1)
    wv1, wv2, Q = align(wv1, wv2, anchor_words=landmarks)
    d_l = [cosine(wv1[w], wv2[w]) for w in landmarks]
    d_n = [cosine(wv1[w], wv2[w]) for w in non_landmarks]
    sns.distplot(d_l, color="blue")
    sns.distplot(d_n, color="red")
    plt.legend()
    plt.show()

if __name__ == "__main__":
    main()


# STUMPY
# Copyright 2019 TD Ameritrade. Released under the terms of the 3-Clause BSD license.  # noqa: E501
# STUMPY is a trademark of TD Ameritrade IP Company, Inc. All rights reserved.

import numpy as np
import scipy.signal

try:
    from numba.cuda.cudadrv.driver import _raise_driver_not_found
except ImportError:
    pass

def driver_not_found():  # pragma: no cover
    """
    Helper function to raise CudaSupportError driver not found error
    """

    _raise_driver_not_found()

def get_pkg_name():  # pragma: no cover
    """
    Return package name
    """

    return __name__.split(".")[0]

def rolling_window(a, window):
    """
    Use strides to generate rolling/sliding windows for a numpy array

    Parameters
    ----------
    a : ndarray
        numpy array

    window : int
        Size of the rolling window

    Returns
    -------
    output : ndarray
        This will be a new view of the original input array.
    """

    shape = a.shape[:-1] + (a.shape[-1] - window + 1, window)
    strides = a.strides + (a.strides[-1],)

    return np.lib.stride_tricks.as_strided(a, shape=shape, strides=strides)

def z_norm(a, axis=0):
    """
    Calculate the z-normalized input array `a` by subtracting the mean and
    dividing by the standard deviation along a given axis.

    Parameters
    ----------
    a : ndarray
        numpy array

    axis : int
        numpy axis

    Returns
    -------
    output : ndarray
        An ndarray with z-normalized values computed along a specified axis.
    """

    return (a - np.mean(a, axis, keepdims=True)) / np.std(a, axis, keepdims=True)

def check_dtype(a, dtype=np.floating):  # pragma: no cover
    """
    Check if the array type of `a` is of type specified by `dtype` parameter.

    Raises
    ------
    TypeError
        If the array type does not match `dtype`
    """

    if not np.issubdtype(a.dtype, dtype):
        msg = f"{dtype} type expected but found {a.dtype}"
        raise TypeError(msg)

    return True

def transpose_dataframe(a):  # pragma: no cover
    """
    Check if the input is a column-wise Pandas `DataFrame`. If `True`, return a
    transpose dataframe since stumpy assumes that each row represents data from a
    different dimension while each column represents data from the same dimension.
    If `False`, return `a` unchanged. Pandas `Series` do not need to be transposed.

    Note that this function has zero dependency on Pandas (not even a soft dependency).

    Parameters
    ----------
    a : ndarray
        First argument.

    Returns
    -------
    output : a
        If a is a Pandas `DataFrame` then return `a.T`. Otherwise, return `a`
    """

    if type(a).__name__ == "DataFrame":
        return a.T

    return a

def are_arrays_equal(a, b):  # pragma: no cover
    """
    Check if two arrays are equal; first by comparing memory addresses,
    and secondly by their values.

    Parameters
    ----------
    a : ndarray
        First argument.

    b : ndarray
        Second argument.

    Returns
    -------
    output : bool
        Returns `True` if the arrays are equal and `False` otherwise.
    """

    if id(a) == id(b):
        return True

    return np.array_equal(a, b)

def are_distances_too_small(a, threshold=10e-6):  # pragma: no cover
    """
    Check the distance values from a matrix profile.

    If the values are smaller than the threshold (i.e., less than 10e-6) then
    it could suggest that this is a self-join.

    Parameters
    ----------
    a : ndarray
        First argument.

    threshold : float
        Minimum value in which to compare the matrix profile to

    Returns
    -------
    output : bool
        Returns `True` if the matrix profile distances are all below the
        threshold and `False` if they are all above the threshold.
    """

    if a.mean() < threshold or np.all(a < threshold):
        return True

    return False

def check_window_size(m):
    if m <= 2:
        raise ValueError(
            "All window sizes must be greater than or equal to three",
            """A window size that is less than or equal to two is meaningless when
            it comes to computing the z-normalized Euclidean distance. In the case of
            `m=1` produces a standard deviation of zero. In the case of `m=2`, both
            the mean and standard deviation for any given subsequence are identical
            and so the z-normalization for any sequence will either be [-1., 1.] or
            [1., -1.]. Thus, the z-normalized Euclidean distance will be (very likely)
            zero between any subsequence and its nearest neighbor (assuming that the
            time series is large enough to contain both scenarios).
            """,
        )

def sliding_dot_product(Q, T):
    """
    Use FFT convolution to calculate the sliding window dot product.

    Parameters
    ----------
    Q : ndarray
        Query array or subsequence

    T : ndarray
        Time series or sequence

    Returns
    -------
    output : ndarray
        Sliding dot product between `Q` and `T`.

    Notes
    -----
    Calculate the sliding dot product

    `DOI: 10.1109/ICDM.2016.0179 \
    <https://www.cs.ucr.edu/~eamonn/PID4481997_extend_Matrix%20Profile_I.pdf>`__

    See Table I, Figure 4

    Following the inverse FFT, Fig. 4 states that only cells [m-1:n]
    contain valid dot products

    Padding is done automatically in fftconvolve step
    """

    n = T.shape[0]
    m = Q.shape[0]
    Qr = np.flipud(Q)  # Reverse/flip Q
    QT = convolution(Qr, T)

    return QT.real[m - 1 : n]

def compute_mean_std(T, m):
    """
    Compute the sliding mean and standard deviation for the array `T` with
    a window size of `m`

    Parameters
    ----------
    T : ndarray
        Time series or sequence

    m : int
        Window size

    Returns
    -------
    M_T : ndarray
        Sliding mean

    Σ_T : ndarray
        Sliding standard deviation

    Notes
    -----
    `DOI: 10.1109/ICDM.2016.0179 \
    <https://www.cs.ucr.edu/~eamonn/PID4481997_extend_Matrix%20Profile_I.pdf>`__

    See Table II

    DOI: 10.1145/2020408.2020587

    See Page 2 and Equations 1, 2

    DOI: 10.1145/2339530.2339576

    See Page 4

    http://www.cs.unm.edu/~mueen/FastestSimilaritySearch.html

    Note that Mueen's algorithm has an off-by-one bug where the
    sum for the first subsequence is omitted and we fixed that!
    """
    n = T.shape[0]

    cumsum_T = np.empty(len(T) + 1)
    np.cumsum(T, out=cumsum_T[1:])  # store output in cumsum_T[1:]
    cumsum_T[0] = 0

    cumsum_T_squared = np.empty(len(T) + 1)
    np.cumsum(np.square(T), out=cumsum_T_squared[1:])
    cumsum_T_squared[0] = 0

    subseq_sum_T = cumsum_T[m:] - cumsum_T[: n - m + 1]
    subseq_sum_T_squared = cumsum_T_squared[m:] - cumsum_T_squared[: n - m + 1]
    M_T = subseq_sum_T / m
    Σ_T = np.abs((subseq_sum_T_squared / m) - np.square(M_T))
    Σ_T = np.sqrt(Σ_T)

    return M_T, Σ_T

def calculate_distance_profile(m, QT, μ_Q, σ_Q, M_T, Σ_T):
    """
    Compute the distance profile

    Parameters
    ----------
    m : int
        Window size

    QT : ndarray
        Dot product between `Q` and `T`

    μ_Q : ndarray
        Mean of `Q`

    σ_Q : ndarray
        Standard deviation of `Q`

    M_T : ndarray
        Sliding mean of `T`

    Σ_T : ndarray
        Sliding standard deviation of `T`

    Returns
    -------
    output : ndarray
        Distance profile

    Notes
    -----
    `DOI: 10.1109/ICDM.2016.0179 \
    <https://www.cs.ucr.edu/~eamonn/PID4481997_extend_Matrix%20Profile_I.pdf>`__

    See Equation on Page 4
    """

    denom = m * σ_Q * Σ_T
    denom[denom == 0] = 1e-10  # Avoid divide by zero
    D_squared = np.abs(2 * m * (1.0 - (QT - m * μ_Q * M_T) / denom))
    return np.sqrt(D_squared)

def mueen_calculate_distance_profile(Q, T):
    """
    Compute the mueen distance profile

    Parameters
    ----------
    Q : ndarray
        Query array or subsequence

    T : ndarray
        Time series or sequence

    Returns
    -------
    output : ndarray
        Distance profile

    Notes
    -----
    `DOI: 10.1109/ICDM.2016.0179 \
    <https://www.cs.ucr.edu/~eamonn/PID4481997_extend_Matrix%20Profile_I.pdf>`__

    See Table II

    DOI: 10.1145/2020408.2020587

    See Page 2 and Equations 1, 2

    DOI: 10.1145/2339530.2339576

    See Page 4

    http://www.cs.unm.edu/~mueen/FastestSimilaritySearch.html

    Note that Mueen's algorithm has an off-by-one bug where the
    sum for the first subsequence is omitted and we fixed that!
    """
    n = T.shape[0]
    m = Q.shape[0]

    μ_Q = np.mean(Q, keepdims=True)
    σ_Q = np.std(Q, keepdims=True)
    Q_norm = (Q - μ_Q) / σ_Q
    QT = sliding_dot_product(Q_norm, T)

    cumsum_T = np.empty(len(T) + 1)  # Add one element, fix off-by-one
    np.cumsum(T, out=cumsum_T[1:])  # store output in cumsum_T[1:]
    cumsum_T[0] = 0

    cumsum_T_squared = np.empty(len(T) + 1)
    np.cumsum(np.square(T), out=cumsum_T_squared[1:])
    cumsum_T_squared[0] = 0

    subseq_sum_T = cumsum_T[m:] - cumsum_T[: n - m + 1]
    subseq_sum_T_squared = cumsum_T_squared[m:] - cumsum_T_squared[: n - m + 1]
    M_T = subseq_sum_T / m
    Σ_T_squared = np.abs(subseq_sum_T_squared / m - np.square(M_T))
    Σ_T = np.sqrt(Σ_T_squared)

    D = np.abs(
        (subseq_sum_T_squared - 2 * subseq_sum_T * M_T + m * np.square(M_T))
        / Σ_T_squared
        - 2 * QT / Σ_T
        + m
    )
    return np.sqrt(D)

def mass(Q, T, M_T=None, Σ_T=None):
    """
    Compute the distance profile using the MASS algorithm

    Parameters
    ----------
    Q : ndarray
        Query array or subsequence

    T : ndarray
        Time series or sequence

    M_T : ndarray (optional)
        Sliding mean of `T`

    Σ_T : ndarray (optional)
        Sliding standard deviation of `T`

    Returns
    -------
    output : ndarray
        Distance profile

    Notes
    -----
    `DOI: 10.1109/ICDM.2016.0179 \
    <https://www.cs.ucr.edu/~eamonn/PID4481997_extend_Matrix%20Profile_I.pdf>`__

    See Table II

    Note that Q, T are not directly required to calculate D

    Note: Unlike the Matrix Profile I paper, here, M_T, Σ_T can be calculated
    once for all subsequences of T and passed in so the redundancy is removed
    """

    QT = sliding_dot_product(Q, T)
    m = Q.shape[0]
    μ_Q, σ_Q = compute_mean_std(Q, m)
    if M_T is None or Σ_T is None:
        M_T, Σ_T = compute_mean_std(T, m)

    return calculate_distance_profile(m, QT, μ_Q, σ_Q, M_T, Σ_T)

convolution = scipy.signal.fftconvolve  # Swap for other convolution function