#!/usr/bin/env python3
# -*- coding: utf-8 -*-
'''
sequential mnist with a drawing task. Given an input sequence of MNIST pixels,
predict (x_t, y_t) coordinates for drawing. Same as train-model.py but records model activities/gradients
'''

import sys
import os
import argparse
import datetime
import yaml

import numpy as np
import torch
import torch.nn as nn
from ignite.engine import Events
from ignite.metrics import Loss, Accuracy
import glob

src_path = 'src_path' #path to src directory
sys.path.insert(0, src_path)

from dataset.mnist import load_mnist, draw_2dlines, draw_2ddigits

from model.init import init_model
from training.setup import set_seed_and_device, setup_training, setup_logging
from training.engine import run_training

import logging
logging.getLogger("ignite").setLevel(logging.NOTSET)

import torchvision.transforms as trans
from correlation import get_popcorr, get_pearson_corr, get_pearson_corr_batch, get_coding_space, get_sparsity, get_weight_info

def train_model(
    # Model parameters
    model='LSTM', nlayers=1, nhid=50, dropout=0.0,
    # Data paramters
    data_path='../../data/exp_raw/', input_size=28, 
    shuffle=True, val_split=0.2, download=False,
    # Training parameters
    epochs=10, batch_size=50, optim='RMSprop', lr=1e-4,
    l2_norm=0.0, rate_reg=0.0, clip=1.0, early_stopping=False,
    decay_lr=False, lr_scale=0.1, lr_decay_patience=10, keep_hidden=False,
    numbers=None, draw_digits=False, spars_int=None, fixed_mag=True, 
    #truncation/synthesiser params
    bptt=None, synth_nlayers=0, synth_nhid=None, nfibres=None, bias=None, record_grads=False,
    # Replicability and storage
    save='../../data/sims/seq-mnist/test', seed=18092, no_cuda=False,
    verbose=False, log_interval=10, nseeds=3,
):
    # Set training seed and get device
    device = set_seed_and_device(seed, no_cuda)
      
    #define drawing task
    seq_len = int(28**2/input_size)
    if numbers is not None:
        numbers = list(map(int, numbers.strip('[]').split(','))) #convert string to list    

    if draw_digits:
        curves = draw_2ddigits(seq_len, numbers)
    else:
        if numbers is not None:
            curves = draw_2dlines(len(numbers), seq_len, fixed_mag=fixed_mag)
        else:
            curves = draw_2dlines(10, seq_len, fixed_mag=fixed_mag)
    if numbers is not None:
        target_transform = trans.Lambda(lambda x: curves[numbers.index(x)])
    else:
        target_transform = trans.Lambda(lambda x: curves[x])
    #Load training data
    train_data, validation_data, test_data = load_mnist(
        data_path, input_size, batch_size, 
        shuffle=shuffle, download=download, target_transform=target_transform, 
        numbers=numbers
    )
   
    #Initialise model
    model_type = model  
    model = init_model(
        model_type=model,
        n_layers=nlayers, hidden_size=nhid,
        input_size=input_size, output_size=2,
        device=device,
        dropout=dropout,
        synth_nlayers=synth_nlayers,
        synth_nhid=synth_nhid,
        predict_last=False
    )

    #bias for RNN intialisation (not modeled in neurips2021)
    if bias != None:
        with torch.no_grad():
            if model_type in ['DNI_LSTM', 'DNI_TANH', 'cDNI_LSTM', 'cDNI_TANH']:
                model.rnn.rnn.weight_hh_l0[:] = model.rnn.rnn.weight_hh_l0[:] * bias
                model.rnn.rnn.weight_ih_l0[:] = model.rnn.rnn.weight_ih_l0[:] * bias
                           
                model.rnn.rnn.bias_hh_l0[:] = model.rnn.rnn.bias_hh_l0[:]*bias
                model.rnn.rnn.bias_ih_l0[:] = model.rnn.rnn.bias_ih_l0[:]*bias            
            else:
                model.rnn.weight_hh_l0[:] = model.rnn.weight_hh_l0[:] * bias
                model.rnn.weight_ih_l0[:] = model.rnn.weight_ih_l0[:] * bias


    #sparsify input 'mossy fibre' synthesiser weights if desired (not modeled in neurips2021)
    if model_type in ['LSTM', 'TANH']:
        nfibres = None
    mask = None
    if nfibres is not None:
        #synth_sparsity = 0.5
        
        shape = model.rnn.backward_interface.synthesizer.input_trigger.weight.shape
        
        mask = torch.ones(shape)
        
        if nfibres != -1:
            for i in range(shape[0]):
                connections = np.random.choice(shape[1], nfibres, replace=False)
                mask[i, connections] = 0
        mask = mask.type(torch.bool)
        with torch.no_grad():
            model.rnn.backward_interface.synthesizer.input_trigger.weight[mask] = 0
            print("Initialised all but {} weights to zero".format(nfibres))
        
        if nfibres == -1:
            mask = None

    #if bptt not set then set as sequence length
    if bptt is None:
        npixels = 28**2
        assert float.is_integer(npixels/input_size), "Input size ({:d}) doesn't divide" \
                                                     "number of pixels ({:d})".format(input_size, npixels)
        seq_len = int(npixels/input_size)
        bptt = seq_len
                                                        
    # Set up the training regime
    setup = setup_training(
        model, validation_data, optim, ['mse'], lr, l2_norm,
        rate_reg, clip, early_stopping, decay_lr, lr_scale, lr_decay_patience,
        keep_hidden, save, device, True, True,
        bptt, batch_size, False, record_grads, spars_int, mask=mask #JOP
    )

    trainer, validator, checkpoint, metrics = setup[:4]
    training_tracer, validation_tracer, timer = setup[4:]

    if verbose:
        setup_logging(
            trainer, validator, metrics,
            len(train_data), log_interval
        )

    # Run training
    test_metrics = run_training(
        model=model,
        train_data=train_data,
        trainer=trainer,
        epochs=epochs,
        test_data=test_data,
        metrics=metrics,
        model_checkpoint=checkpoint,
        device=device
    )
    
    if record_grads:
        test_metrics, grads = test_metrics
     
    # Testing preformance
    test_mse = test_metrics['mse']

    print('Training ended: test MSE {:5.4f}'.format(
        test_mse))

    print('Saving results....')

    # Save traces
    training_tracer.save(save)
    validation_tracer.save(save)

    # Save experiment metadata
    model_params = {
        'model_type': model_type,
        'hidden_size': nhid,
        'n_layers': nlayers,
        'input_size': input_size,
        'dropout': dropout,
        'synth_nlayers': synth_nlayers,
        'synth_nhid': synth_nhid, 
        'nfibres': nfibres
    }

    learning_params = {
        'optimizer': optim,
        'learning-rate': lr,
        'l2-norm': l2_norm,
        'criterion': 'mse',
        'batch_size': batch_size,
        'bptt': bptt,
        'spars_int': spars_int
    }

    # Save data parameters in a dictionary for testing
    data_params = {
        'val-split': val_split, 'keep-hidden': keep_hidden,
        'numbers': numbers, 'draw_digits': draw_digits, 'fixed_mag': fixed_mag
    }

    meta = {
        'data-params': data_params,
        'model-params': model_params,
        'learning-params': learning_params,
        'info': {
            'test-mse': test_mse,
            'training-time': timer.value(),
            'timestamp': datetime.datetime.now()
        },
        'seed': seed
    }

    with open(save + '/meta.yaml', mode='w') as f:
        yaml.dump(meta, f)
    np.save(save + '/curves.npy', curves)

    print('Done.')
    
    #JOP
    train_mse = training_tracer.trace
    val_mse = [vt[0] for vt in validation_tracer.trace]
    
    if not record_grads:
        return train_mse, val_mse 
    else:
        return val_mse, grads, model

##############################################################################
# PARSE THE INPUT
##############################################################################
if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description='Train an LSTM variant on the Delayed Addition task')

    # Model parameters
    parser.add_argument('--model', type=str, default='LSTM',
                        help='RNN model tu use. One of:'
                        '|TANH|DNI_TANH|LSTM|DNI_LSTM') #JOP
    parser.add_argument('--nlayers', type=int, default=1,
                        help='number of layers')
    parser.add_argument('--nhid', type=int, default=30,
                        help='number of hidden units per layer')
    parser.add_argument('--dropout', type=float, default=0.0,
                        help='the drop rate for each layer of the network')

    # Data parameters
    parser.add_argument('--data-path', type=str, default='../../data/exp_raw/MNIST',
        help='location of the data set')
    parser.add_argument('--input-size', type=int, default=28,
        help='the default dimensionality of each input timestep.'
        'defaults to 1, meaning instances are treated like one large 1D sequence')
    parser.add_argument('--val-split', type=float, default=0.2,
        help='proportion of trainig data used for validation')
    parser.add_argument('--shuffle', action='store_true',
        help='shuffle the data at the start of each epoch.')

    # Training parameters
    parser.add_argument('--epochs', type=int, default=40,
                        help='max number of training epochs')
    parser.add_argument('--batch-size', type=int, default=50, metavar='N',
                        help='size of each batch. keep in mind that if the '
                        'training size modulo this quantity is not zero, then'
                        'the it will be increased to create a full batch.')
    parser.add_argument('--optim', type=str, default='adam',
                        help='gradient descent method, supports on of:'
                        'adam|sparseadam|adamax|rmsprop|sgd|adagrad|adadelta')
    parser.add_argument('--lr', type=float, default=1e-4,
                        help='initial learning rate')
    parser.add_argument('--l2-norm', type=float, default=0,
                        help='weight of L2 norm')
    parser.add_argument('--rate-reg', type=float, default=0.0,
                        help='regularization factor for hidden variables')
    parser.add_argument('--clip', type=float, default=1,
                        help='gradient clipping')
    parser.add_argument('--early-stopping', action='store_true',
                        help='use early stopping')
    parser.add_argument('--decay_lr', action='store_true',
                        help='if provided will decay the learning after not'
                        'improving for a certain amount of epochs')
    parser.add_argument('--lr-scale', type=float, default=0.1,
                        help='the factor by which learning rate should be scaled'
                        'after the patience has been exhausted')
    parser.add_argument('--lr_decay_patience', type=int, default=10,
                        help='specifies the number of epochs to wait until'
                        'decay is applied before applying the decay factor')
    parser.add_argument('--keep-hidden', action='store_true',
                        help='keep the hidden state values across an epoch'
                        'of training, detaching them from the computation graph'
                        'after each batch for gradient consistency')

    # Replicability and storage
    parser.add_argument('--save', type=str,
                        default='../../data/sims/seq-mnist/test',
                        help='path to save the final model')
    parser.add_argument('--seed', type=int, default=18092,
                        help='random seed')
    parser.add_argument('--nseeds', type=int, default=3,
			            help='number of seeds')

    # CUDA
    parser.add_argument('--no-cuda', action='store_true',
                        help='flag to disable CUDA')

    # Print options
    parser.add_argument('--verbose', action='store_true',
                        help='print the progress of training to std output.')
    parser.add_argument('--log-interval', type=int, default=10, metavar='N',
                        help='report interval')
    
    #DNI specific
    parser.add_argument('--bptt', type=int, metavar='N',
                        help='truncation size')
    parser.add_argument('--record-grads', action='store_true',
                        help='Record hidden activity and corresponding true/synthetic gradients')
    parser.add_argument('--synth-nlayers', type=int, default=1,
                        help='Number of hidden layers for synthesiser')
    parser.add_argument('--synth-nhid', type=int, metavar='N',
                        help='Number of hidden units in synthesiser')

    parser.add_argument('--numbers', type=str,
                        help='Only train model for these numbers')
    parser.add_argument('--draw-digits', action='store_true',
                        help='Train model to draw digits themselves, as opposed to lines')   
    parser.add_argument('--fixed-mag', action='store_true',
                        help='When drawing straight lines, have fixed magnitude.')
    parser.add_argument('--bias', type=float, help='bias for recurrent weights (manipulates correlations')


    parser.add_argument('--spars-int', type=int, 
                        help='the interval at which to sparsify the targets')
    parser.add_argument('--nfibres', type=int, metavar='N',
                        help='Number of non-zero weights from hidden to synthesiser (like mossy fibres)')  

    args = parser.parse_args()
   
    args.model = 'DNI_LSTM'
    args.bptt = 3
    args.epochs = 1
    args.nseeds = 1
    args.record_grads = True
    
    assert args.record_grads, 'In records_grad script, set it to true!'

    nhid = args.nhid
    nbatch = args.batch_size
    
    task = 'seqmnist-dd' if args.draw_digits else 'seqmnist-ld'

    root = 'path_to_save_results/{}/'.format(task)
        
    seeds = [521, 1243, 34521, 135235, 236236, 984, 68516, 98945, 3223, 69490, 4093202, 95389, 3984, 129, 59348, 53985, 23153, 230481, 34598, 12309][:args.nseeds]  

    #bias for initialisation of RNN weights (not modeled in neurips2021)
    biases = [1]

    nfibres_list = [None]
    corr_methods = [get_popcorr, get_pearson_corr, get_coding_space, get_sparsity]
    
    score_shape = (len(seeds), len(biases), len(nfibres_list), args.epochs) #add 1 for first batch
    scores = np.zeros(score_shape)
    corr_shape = score_shape[:3] + (args.epochs + 1, len(corr_methods)) # x pop/pearson/etc
    
    corrs_hidd = np.zeros(corr_shape) 
    corrs_truegrads = np.zeros(corr_shape)
    corrs_shidds = np.zeros(corr_shape)
    corrs_sgrad = np.zeros(corr_shape)
    corrs_tgrad = np.zeros(corr_shape)

    weights_info = np.zeros(score_shape[:2] + (2, 4,4)) #dni/sparsedni, #weights, #measures(mean, absolute mean, absolute mean without zeros, variance)
    
    for i, seed in enumerate(seeds):
        args.seed = seed
        for j, bias in enumerate(biases):
            args.bias = bias
            for k, nfibres in enumerate(nfibres_list):
                if nfibres==-1:
                    args.model = 'LSTM'       
                else:
                    args.model = 'DNI_LSTM'    
                args.nfibres = nfibres
                
                #run training
                val_mse, grads, model = train_model(**vars(args))
                scores[i,j,k] = val_mse
                
                #statistics for learnt weights (not modeled in neurips2021)
                if k > 0:
                    synthesiser = model.rnn.backward_interface.synthesizer
                    for w, weight in enumerate([synthesiser.input_trigger.weight, synthesiser.input_trigger.bias, synthesiser.layers[0].weight, synthesiser.layers[0].bias]):
                        weights_info[i, j, k-1, w, 0] = get_weight_info(weight)
                        weights_info[i, j, k-1, w, 1] = get_weight_info(weight, absolute=True)
                        weights_info[i, j, k-1, w, 2] = get_weight_info(weight, absolute=True, ignore_zeros=True)
                        weights_info[i, j, k-1, w, 3] = get_weight_info(weight, getvar=True)                       

                hidds, truegrads_output, truegrads_cell, tgrads_output, tgrads_cell, sgrads, shidds = grads               
                hidds = torch.stack(hidds).squeeze(1)

                if nfibres != -1:
                    truegrads_output = torch.stack(truegrads_output).squeeze(1)
                    truegrads_cell = torch.stack(truegrads_cell).squeeze(1)
                    tgrads_output = torch.stack(tgrads_output).squeeze(1)
                    tgrads_cell = torch.stack(tgrads_cell).squeeze(1)
                    sgrads = torch.stack(sgrads).squeeze(1)
                    shidds = torch.stack(shidds).squeeze(1)
                
                savedir = "{}_bias-{}_seed-{}".format(args.model, args.bias, args.seed)
                if nfibres is not None:
                    savedir += "_nfibres-{}".format(nfibres)

                savedir = os.path.join(root, savedir) + "/"
                if not os.path.exists(savedir): 
                    os.makedirs(savedir)
                    
                """
                np.save(savedir + 'hidds.npy', hidds.detach().numpy())
                np.save(savedir + 'sgrads.npy', sgrads.detach().numpy())
                np.save(savedir + 'truegrads_output.npy', truegrads_output.detach().numpy())
                np.save(savedir + 'truegrads_cell.npy', truegrads_cell.detach().numpy())
                np.save(savedir + 'tgrads_output.npy', tgrads_output.detach().numpy())
                np.save(savedir + 'tgrads_cell.npy', tgrads_cell.detach().numpy())
                np.save(savedir + 'shidds.npy', shidds.detach().numpy())
                
                np.save(savedir + 'scores.npy', val_mse)

                np.save(savedir + 'biases.npy', np.array(biases))
                np.save(savedir + 'seeds.npy', np.array(seeds))
                """


    print("All finished and results saved in", root)
