import copy
import pickle
import json
import csv
import argparse
import sys

import torch
import numpy as np
from torch.autograd import Variable

sys.path.append('src/utils/')
from utils import config_parser
from utils.filter_helper import compute_similarity, sample_poles, return_filter_coeffs, extract_eigen_values
from utils.data_provider import create_signal_dataset
from utils.plotting import plot_loss_history, plot_polar_animation, plot_polar_plot

sys.path.append('src/models/')
from models.linear_rnn import LinearRNN
from models.oob_lstm import OOBLSTM

torch.manual_seed(0)

# Set up Boolean Flags
flag_optimal = True
flag_bptt = True
flag_bptt_biased = False
flag_bptt_l2 = False
flag_bptt_lstm = False

# Parse Config Parameters
input_dim, hidden_dim, output_dim, num_layers = config_parser.get_config(config_parser.Sections.RNN)
nepoch, nepoch_biased, spectral_norm_enabled, spectral_threshold, sgd = config_parser.get_config(config_parser.Sections.TRANING)
N, seq_len, wash_out, batch_size, split_ratio = config_parser.get_config(config_parser.Sections.LINEAR)

# Parse Arguments
parser = argparse.ArgumentParser()
parser.add_argument("--mode", default='single')
parser.add_argument("--Nmax", default='3')
parser.add_argument("--trials", default='3')
parser.add_argument("--threshold", default='0.2')
parser.add_argument("--plotting", default='False')
args = parser.parse_args()

def run_network(N, hidden_dim):

    z, p = sample_poles(N)
    real_poles, imag_poles = [x.real for x in p], [x.imag for x in p]

    b, a = return_filter_coeffs(z, p)

    realW_bptt_biased, imagW_bptt_biased = [], []
    realW_bptt_l2, imagW_bptt_l2 = [], []

    # Data sets for train & test
    X, Y = create_signal_dataset(batch_size, seq_len, 0, 1, z, p)
    split_index = int(split_ratio*batch_size)

    train_X, train_Y = X[:split_index], Y[:split_index]
    test_X, test_Y = X[split_index:], Y[split_index:]

    train_inp = Variable(torch.Tensor(train_X.reshape((train_X.shape[0], -1, 1))), requires_grad=True)
    train_out = Variable(torch.Tensor(train_Y.reshape((train_Y.shape[0], -1, 1))))

    test_inp = Variable(torch.Tensor(test_X.reshape((test_X.shape[0], -1, 1))), requires_grad=True)
    test_out = Variable(torch.Tensor(test_Y.reshape((test_Y.shape[0], -1, 1))))

    if flag_optimal:
        #region Approximate via RNN with optimal weights
        network_optimal = LinearRNN(input_dim, N, output_dim, num_layers, train_inp.shape[0])

        #Shifted diagonal matrix and IIR Filter coefficients:
        a_new = -a
        a_new[0] = 1
        network_optimal.state_dict()['rnn.weight_hh_l0'].data.copy_(torch.diag(torch.ones(N-1),diagonal=-1))
        network_optimal.state_dict()['rnn.weight_hh_l0'][0,:].data.copy_(torch.tensor(a_new[1:]))
        #Vector with a "1" for input x:
        network_optimal.state_dict()['rnn.weight_ih_l0'].data.copy_(torch.zeros((N,1)))
        network_optimal.state_dict()['rnn.weight_ih_l0'][0,0].data.copy_(torch.tensor(1.0))
        #Weight for linear output layer to select h[0]:
        network_optimal.state_dict()['linear.weight'].data.copy_(torch.zeros((1,N)))
        network_optimal.state_dict()['linear.weight'][0,0].data.copy_(torch.tensor(1.0))

        pred_network_optimal = network_optimal(test_inp)

        error_network_optimal = network_optimal.loss_func(pred_network_optimal[:,wash_out:,:], test_out[:,wash_out:,:]).detach().numpy()
        print(error_network_optimal)

        realW_optimal, imagW_optimal = extract_eigen_values(network_optimal.rnn.weight_hh_l0.detach().numpy())
        #endregion

    if flag_bptt:
        #region Approximate via RNN with BPTT
        network_bptt = LinearRNN(input_dim, hidden_dim, output_dim, num_layers, train_inp.shape[0])
        network_bptt_l2 = copy.deepcopy(network_bptt)

        # for param in network_bptt.linear.parameters():
        #     param.requires_grad = False

        error_train_bptt, spectra_rnn_bptt = network_bptt.train(train_inp, train_out, nepoch, wash_out, 0.001, 0, spectral_norm_enabled, spectral_threshold, sgd)
        pred_bptt_linear = network_bptt(test_inp)

        error_pred_bptt = network_bptt.loss_func(pred_bptt_linear[:,wash_out:,:], test_out[:,wash_out:,:]).detach().numpy()
        print(error_pred_bptt)

        realW_bptt, imagW_bptt = extract_eigen_values(network_bptt.rnn.weight_hh_l0.detach().numpy())

        filename = 'bin/filter/error_train_bptt.pkl'
        outfile = open(filename,'wb')
        pickle.dump(error_train_bptt, outfile)
        outfile.close()
        #endregion

    if flag_bptt_biased:
        #region Approximate via RNN with biased weights
        network_bptt_biased = LinearRNN(input_dim, N, output_dim, num_layers, train_inp.shape[0])

        def add_noise(shape):
            return torch.normal(mean=torch.zeros(shape), std=torch.ones(shape))*0.05

        #Shifted diagonal matrix and IIR Filter coefficients:
        a_new = -a
        a_new[0] = 1
        network_bptt_biased.state_dict()['rnn.weight_hh_l0'].data.copy_(torch.diag(torch.ones(N-1),diagonal=-1))
        network_bptt_biased.state_dict()['rnn.weight_hh_l0'][0,:].data.copy_(torch.tensor(a_new[1:]))
        network_bptt_biased.state_dict()['rnn.weight_hh_l0'].data -= add_noise(network_bptt_biased.state_dict()['rnn.weight_hh_l0'].data.shape)
        #Vector with a "1" for input x:
        network_bptt_biased.state_dict()['rnn.weight_ih_l0'].data.copy_(torch.zeros((N,1)))
        network_bptt_biased.state_dict()['rnn.weight_ih_l0'][0,0].data.copy_(torch.tensor(1.0))
        network_bptt_biased.state_dict()['rnn.weight_ih_l0'].data -= add_noise(network_bptt_biased.state_dict()['rnn.weight_ih_l0'].data.shape)
        #Weight for linear output layer to select h[0]:
        network_bptt_biased.state_dict()['linear.weight'].data.copy_(torch.zeros((1,N)))
        network_bptt_biased.state_dict()['linear.weight'][0,0].data.copy_(torch.tensor(1.0))
        network_bptt_biased.state_dict()['linear.weight'].data -= add_noise(network_bptt_biased.state_dict()['linear.weight'].data.shape)

        error_train_bptt_biased, spectra_network_bptt_biased = network_bptt_biased.train(train_inp, train_out, nepoch_biased, wash_out, 0.001, 0, spectral_norm_enabled, spectral_threshold, sgd)
        pred_bptt_biased = network_bptt_biased(test_inp)

        error_pred_bptt_biased = network_bptt_biased.loss_func(pred_bptt_biased[:,wash_out:,:], test_out[:,wash_out:,:]).detach().numpy()
        print(error_pred_bptt_biased)

        real_spectra_network_bptt_biased, imag_spectra_network_bptt_biased = spectra_network_bptt_biased['real'], spectra_network_bptt_biased['imag']
        realW_bptt_biased, imagW_bptt_biased = extract_eigen_values(network_bptt_biased.rnn.weight_hh_l0.detach().numpy())

        plot_polar_animation(real_poles, imag_poles,real_spectra_network_bptt_biased, imag_spectra_network_bptt_biased, nepoch_biased)

        filename = 'bin/filter/error_train_bptt_biased.pkl'
        outfile = open(filename,'wb')
        pickle.dump(error_train_bptt_biased, outfile)
        outfile.close()
        #endregion

    if flag_bptt_l2:
        #region Approximate via RNN with BPTT
        error_train_bptt_l2, _ = network_bptt_l2.train(train_inp, train_out, nepoch, wash_out, 0.001, 0.01, spectral_norm_enabled, spectral_threshold, sgd)
        pred_bptt_l2 = network_bptt_l2(test_inp)

        error_pred_bptt_l2 = network_bptt_l2.loss_func(pred_bptt_l2[:,wash_out:,:], test_out[:,wash_out:,:]).detach().numpy()
        print(error_pred_bptt_l2)

        realW_bptt_l2, imagW_bptt_l2 = extract_eigen_values(network_bptt_l2.rnn.weight_hh_l0.detach().numpy())
        #endregion

    if flag_bptt_lstm:
        #region Approximate via LSTM with BPTT
        network_lstm = OOBLSTM(input_dim, hidden_dim, output_dim, num_layers, batch_size)
        error_train_bptt_lstm = network_lstm.train(X, Y, nepoch, wash_out, 0.001, spectral_norm_enabled, spectral_threshold, sgd)
        pred_bptt_lstm = network_lstm(test_inp)

        error_pred_bptt_lstm = network_lstm.loss_func(pred_bptt_lstm[:,wash_out:,:], test_out[:,wash_out:,:]).detach().numpy()
        print(error_pred_bptt_lstm)
        #endregion

    if bool(args.plotting):
        plot_loss_history(error_train_bptt, error_network_optimal*np.ones(nepoch))

        plot_polar_plot(real_poles, imag_poles, realW_bptt, imagW_bptt, realW_bptt_biased, imagW_bptt_biased, realW_bptt_l2, imagW_bptt_l2)

    # Compute Solution Similarity to Optimal
    overlaps = compute_similarity(float(args.threshold), real_poles, imag_poles, realW_bptt, imagW_bptt)
    print(overlaps)

    return {
        "N": N,
        "hidden_dims": hidden_dim,
        "threshold": float(args.threshold),
        "overlaps": overlaps,
        "poles": str([real_poles, imag_poles]),
        "eigenvalues": str([realW_bptt, imagW_bptt]),
        "error_network_optimal": str(error_network_optimal),
        "error_pred_bptt": str(error_pred_bptt),
    }

def main():
    if args.mode == 'bulk':
        f = open("bin/filter/bulk_run_"+args.trials+".csv", "w")
        fieldnames = ['N', 'trial_number', 'overlaps']
        writer = csv.DictWriter(f, fieldnames=fieldnames)

        for N_current in range(3, int(args.Nmax)+1, 2):
            for i in range(int(args.trials)):
                result = run_network(N_current, N_current)
                writer.writerow({'N':result['N'], 'trial_number':i, 'overlaps':result['overlaps']})
        f.close()

    elif args.mode == 'single':
        results = run_network(N, hidden_dim)
    
        # Write Results to Local File
        with open('bin/filter/run.txt', 'w') as f:
                f.write(json.dumps(results, indent=4))
        f.close()

if __name__ == "__main__":
  main()
import copy
import pickle
import json
import csv
import argparse
import sys

import torch
import numpy as np
from torch.autograd import Variable

sys.path.append('src/utils/')
from utils import config_parser
from utils.filter_helper import compute_similarity, sample_poles, return_filter_coeffs, extract_eigen_values
from utils.data_provider import create_signal_dataset
from utils.plotting import plot_loss_history, plot_polar_animation, plot_polar_plot

sys.path.append('src/models/')
from models.linear_rnn import LinearRNN
from models.oob_lstm import OOBLSTM

torch.manual_seed(0)

# Set up Boolean Flags
flag_optimal = True
flag_bptt = True
flag_bptt_biased = False
flag_bptt_l2 = False
flag_bptt_lstm = False

# Parse Config Parameters
input_dim, hidden_dim, output_dim, num_layers = config_parser.get_config(config_parser.Sections.RNN)
nepoch, nepoch_biased, spectral_norm_enabled, spectral_threshold, sgd = config_parser.get_config(config_parser.Sections.TRANING)
N, seq_len, wash_out, batch_size, split_ratio = config_parser.get_config(config_parser.Sections.LINEAR)

# Parse Arguments
parser = argparse.ArgumentParser()
parser.add_argument("--mode", default='single')
parser.add_argument("--Nmax", default='3')
parser.add_argument("--trials", default='3')
parser.add_argument("--threshold", default='0.2')
parser.add_argument("--plotting", default='False')
args = parser.parse_args()

def run_network(N, hidden_dim):

    z, p = sample_poles(N)
    real_poles, imag_poles = [x.real for x in p], [x.imag for x in p]

    b, a = return_filter_coeffs(z, p)

    realW_bptt_biased, imagW_bptt_biased = [], []
    realW_bptt_l2, imagW_bptt_l2 = [], []

    # Data sets for train & test
    X, Y = create_signal_dataset(batch_size, seq_len, 0, 1, z, p)
    split_index = int(split_ratio*batch_size)

    train_X, train_Y = X[:split_index], Y[:split_index]
    test_X, test_Y = X[split_index:], Y[split_index:]

    train_inp = Variable(torch.Tensor(train_X.reshape((train_X.shape[0], -1, 1))), requires_grad=True)
    train_out = Variable(torch.Tensor(train_Y.reshape((train_Y.shape[0], -1, 1))))

    test_inp = Variable(torch.Tensor(test_X.reshape((test_X.shape[0], -1, 1))), requires_grad=True)
    test_out = Variable(torch.Tensor(test_Y.reshape((test_Y.shape[0], -1, 1))))

    if flag_optimal:
        #region Approximate via RNN with optimal weights
        network_optimal = LinearRNN(input_dim, N, output_dim, num_layers, train_inp.shape[0])

        #Shifted diagonal matrix and IIR Filter coefficients:
        a_new = -a
        a_new[0] = 1
        network_optimal.state_dict()['rnn.weight_hh_l0'].data.copy_(torch.diag(torch.ones(N-1),diagonal=-1))
        network_optimal.state_dict()['rnn.weight_hh_l0'][0,:].data.copy_(torch.tensor(a_new[1:]))
        #Vector with a "1" for input x:
        network_optimal.state_dict()['rnn.weight_ih_l0'].data.copy_(torch.zeros((N,1)))
        network_optimal.state_dict()['rnn.weight_ih_l0'][0,0].data.copy_(torch.tensor(1.0))
        #Weight for linear output layer to select h[0]:
        network_optimal.state_dict()['linear.weight'].data.copy_(torch.zeros((1,N)))
        network_optimal.state_dict()['linear.weight'][0,0].data.copy_(torch.tensor(1.0))

        pred_network_optimal = network_optimal(test_inp)

        error_network_optimal = network_optimal.loss_func(pred_network_optimal[:,wash_out:,:], test_out[:,wash_out:,:]).detach().numpy()
        print(error_network_optimal)

        realW_optimal, imagW_optimal = extract_eigen_values(network_optimal.rnn.weight_hh_l0.detach().numpy())
        #endregion

    if flag_bptt:
        #region Approximate via RNN with BPTT
        network_bptt = LinearRNN(input_dim, hidden_dim, output_dim, num_layers, train_inp.shape[0])
        network_bptt_l2 = copy.deepcopy(network_bptt)

        # for param in network_bptt.linear.parameters():
        #     param.requires_grad = False

        error_train_bptt, spectra_rnn_bptt = network_bptt.train(train_inp, train_out, nepoch, wash_out, 0.001, 0, spectral_norm_enabled, spectral_threshold, sgd)
        pred_bptt_linear = network_bptt(test_inp)

        error_pred_bptt = network_bptt.loss_func(pred_bptt_linear[:,wash_out:,:], test_out[:,wash_out:,:]).detach().numpy()
        print(error_pred_bptt)

        realW_bptt, imagW_bptt = extract_eigen_values(network_bptt.rnn.weight_hh_l0.detach().numpy())

        filename = 'bin/filter/error_train_bptt.pkl'
        outfile = open(filename,'wb')
        pickle.dump(error_train_bptt, outfile)
        outfile.close()
        #endregion

    if flag_bptt_biased:
        #region Approximate via RNN with biased weights
        network_bptt_biased = LinearRNN(input_dim, N, output_dim, num_layers, train_inp.shape[0])

        def add_noise(shape):
            return torch.normal(mean=torch.zeros(shape), std=torch.ones(shape))*0.05

        #Shifted diagonal matrix and IIR Filter coefficients:
        a_new = -a
        a_new[0] = 1
        network_bptt_biased.state_dict()['rnn.weight_hh_l0'].data.copy_(torch.diag(torch.ones(N-1),diagonal=-1))
        network_bptt_biased.state_dict()['rnn.weight_hh_l0'][0,:].data.copy_(torch.tensor(a_new[1:]))
        network_bptt_biased.state_dict()['rnn.weight_hh_l0'].data -= add_noise(network_bptt_biased.state_dict()['rnn.weight_hh_l0'].data.shape)
        #Vector with a "1" for input x:
        network_bptt_biased.state_dict()['rnn.weight_ih_l0'].data.copy_(torch.zeros((N,1)))
        network_bptt_biased.state_dict()['rnn.weight_ih_l0'][0,0].data.copy_(torch.tensor(1.0))
        network_bptt_biased.state_dict()['rnn.weight_ih_l0'].data -= add_noise(network_bptt_biased.state_dict()['rnn.weight_ih_l0'].data.shape)
        #Weight for linear output layer to select h[0]:
        network_bptt_biased.state_dict()['linear.weight'].data.copy_(torch.zeros((1,N)))
        network_bptt_biased.state_dict()['linear.weight'][0,0].data.copy_(torch.tensor(1.0))
        network_bptt_biased.state_dict()['linear.weight'].data -= add_noise(network_bptt_biased.state_dict()['linear.weight'].data.shape)

        error_train_bptt_biased, spectra_network_bptt_biased = network_bptt_biased.train(train_inp, train_out, nepoch_biased, wash_out, 0.001, 0, spectral_norm_enabled, spectral_threshold, sgd)
        pred_bptt_biased = network_bptt_biased(test_inp)

        error_pred_bptt_biased = network_bptt_biased.loss_func(pred_bptt_biased[:,wash_out:,:], test_out[:,wash_out:,:]).detach().numpy()
        print(error_pred_bptt_biased)

        real_spectra_network_bptt_biased, imag_spectra_network_bptt_biased = spectra_network_bptt_biased['real'], spectra_network_bptt_biased['imag']
        realW_bptt_biased, imagW_bptt_biased = extract_eigen_values(network_bptt_biased.rnn.weight_hh_l0.detach().numpy())

        plot_polar_animation(real_poles, imag_poles,real_spectra_network_bptt_biased, imag_spectra_network_bptt_biased, nepoch_biased)

        filename = 'bin/filter/error_train_bptt_biased.pkl'
        outfile = open(filename,'wb')
        pickle.dump(error_train_bptt_biased, outfile)
        outfile.close()
        #endregion

    if flag_bptt_l2:
        #region Approximate via RNN with BPTT
        error_train_bptt_l2, _ = network_bptt_l2.train(train_inp, train_out, nepoch, wash_out, 0.001, 0.01, spectral_norm_enabled, spectral_threshold, sgd)
        pred_bptt_l2 = network_bptt_l2(test_inp)

        error_pred_bptt_l2 = network_bptt_l2.loss_func(pred_bptt_l2[:,wash_out:,:], test_out[:,wash_out:,:]).detach().numpy()
        print(error_pred_bptt_l2)

        realW_bptt_l2, imagW_bptt_l2 = extract_eigen_values(network_bptt_l2.rnn.weight_hh_l0.detach().numpy())
        #endregion

    if flag_bptt_lstm:
        #region Approximate via LSTM with BPTT
        network_lstm = OOBLSTM(input_dim, hidden_dim, output_dim, num_layers, batch_size)
        error_train_bptt_lstm = network_lstm.train(X, Y, nepoch, wash_out, 0.001, spectral_norm_enabled, spectral_threshold, sgd)
        pred_bptt_lstm = network_lstm(test_inp)

        error_pred_bptt_lstm = network_lstm.loss_func(pred_bptt_lstm[:,wash_out:,:], test_out[:,wash_out:,:]).detach().numpy()
        print(error_pred_bptt_lstm)
        #endregion

    if bool(args.plotting):
        plot_loss_history(error_train_bptt, error_network_optimal*np.ones(nepoch))

        plot_polar_plot(real_poles, imag_poles, realW_bptt, imagW_bptt, realW_bptt_biased, imagW_bptt_biased, realW_bptt_l2, imagW_bptt_l2)

    # Compute Solution Similarity to Optimal
    overlaps = compute_similarity(float(args.threshold), real_poles, imag_poles, realW_bptt, imagW_bptt)
    print(overlaps)

    return {
        "N": N,
        "hidden_dims": hidden_dim,
        "threshold": float(args.threshold),
        "overlaps": overlaps,
        "poles": str([real_poles, imag_poles]),
        "eigenvalues": str([realW_bptt, imagW_bptt]),
        "error_network_optimal": str(error_network_optimal),
        "error_pred_bptt": str(error_pred_bptt),
    }

def main():
    if args.mode == 'bulk':
        f = open("bin/filter/bulk_run_"+args.trials+".csv", "w")
        fieldnames = ['N', 'trial_number', 'overlaps']
        writer = csv.DictWriter(f, fieldnames=fieldnames)

        for N_current in range(3, int(args.Nmax)+1, 2):
            for i in range(int(args.trials)):
                result = run_network(N_current, N_current)
                writer.writerow({'N':result['N'], 'trial_number':i, 'overlaps':result['overlaps']})
        f.close()

    elif args.mode == 'single':
        results = run_network(N, hidden_dim)
    
        # Write Results to Local File
        with open('bin/filter/run.txt', 'w') as f:
                f.write(json.dumps(results, indent=4))
        f.close()

if __name__ == "__main__":
  main()
# Set up Boolean Flags # Parse Config Parameters # Parse Arguments # Data sets for train & test #region Approximate via RNN with optimal weights #Shifted diagonal matrix and IIR Filter coefficients: #Vector with a "1" for input x: #Weight for linear output layer to select h[0]: #endregion #region Approximate via RNN with BPTT # for param in network_bptt.linear.parameters(): #     param.requires_grad = False #endregion #region Approximate via RNN with biased weights #Shifted diagonal matrix and IIR Filter coefficients: #Vector with a "1" for input x: #Weight for linear output layer to select h[0]: #endregion #region Approximate via RNN with BPTT #endregion #region Approximate via LSTM with BPTT #endregion # Compute Solution Similarity to Optimal # Write Results to Local File