import sys
import os
import argparse
from datetime import datetime
from numpy import float32
from collections import OrderedDict

from blocks.extensions.training import SharedVariableModifier
from blocks.algorithms import GradientDescent, Adam, CompositeRule, StepClipping, RemoveNotFinite
from blocks.model import Model
from blocks.extensions import Timing, Printing, FinishAfter, ProgressBar
from blocks.extensions.monitoring import TrainingDataMonitoring, DataStreamMonitoring
from blocks.extensions.plot import Plot
from picklable_itertools.extras import equizip
import theano
from theano import config
from blocks.graph import ComputationGraph
from theano.sandbox.rng_mrg import MRG_RandomStreams as RandomStreams
import numpy as np
from fuel.streams import DataStream
from fuel.schemes import ShuffledScheme
from fuel.transformers import Padding, Mapping
from theano import tensor as T
from blocks.main_loop import MainLoop
from blocks.dump import load_parameter_values, save_parameter_values

from utils import MismulitclassificationRate, MismulitmistakeRate, NegativeLogLikelihood
from models.rnnrbm import Rnnrbm
from utils import test_value
from midi import MidiSequence2, MidiSequence
from miditools.utils import midiwrite


floatX = config.floatX
rng = RandomStreams(seed=np.random.randint(1 << 30))
sys.setrecursionlimit(10000)
np.set_printoptions(threshold='nan')

x = T.tensor3('features')
x_mask = T.matrix('features_mask')
y = T.tensor3('targets')
y_mask = T.matrix('targets_mask')

x = test_value(x, np.ones((10, 15, 88), dtype=floatX))
y = test_value(y, np.ones((10, 15, 88), dtype=floatX))
x_mask = test_value(x_mask, np.ones((10, 15), dtype=floatX))
y_mask = test_value(y_mask, np.ones((10, 15), dtype=floatX))

pre_training_params = {
    '/rbm.buh': '/rnnrbm/uh.buh',
    '/rbm.Wrbm': '/rnnrbm/rbm.Wrbm',
    '/rbm.buv': '/rnnrbm/uv.buv',
    # '/rnn/out_layer/linear_0.b',
    # '/rnn/out_layer/linear_0.W',
    # '/rnnrbm/uh.Wuh',
    # '/rnnrbm/uv.Wuv',
    '/rnn/lstm_rnn_layer.W_cell_to_out': '/rnnrbm/lstm_rnn_layer.W_cell_to_out',
    '/rnn/lstm_rnn_layer.W_cell_to_in': '/rnnrbm/lstm_rnn_layer.W_cell_to_in',
    '/rnn/lstm_rnn_layer.W_cell_to_forget': '/rnnrbm/lstm_rnn_layer.W_cell_to_forget',
    '/rnn/h2h_transform.W': '/rnnrbm/h2h_transform.W',
    '/rnn/gru_rnn_layer.W': '/rnnrbm/gru_rnn_layer.W',
    '/rnn/input_transfrom.W': '/rnnrbm/input_transfrom.W',
    '/rnn/lstm_rnn_layer.W_state': '/rnnrbm/lstm_rnn_layer.W_state',
}

rbm_cdk = {
    3000: 3,
    6000: 7,
    9000: 12,
    9500: 20,
    9700: 25,
    9800: 30,
    9900: 40,
    9950: 50,
    9980: 70,
    9990: 80
}

rnnrbm_cdk = {
    5000: 20,
    10000: 30,
    10500: 40,
    11000: 50,
    11500: 60,
    12000: 80,
    12500: 90

}




def pretrain_rnn(train, rnnrbm, test=None, epochs=1000, bokeh=True):
    lr = theano.shared(float32(0.1))

    probs, _, _, _ = rnnrbm.rnn_pretrain_pred(x, x_mask)
    cost = NegativeLogLikelihood().apply(y, probs, y_mask)

    error_rate = MismulitclassificationRate().apply(y, probs, y_mask)
    error_rate.name = "error on note as a whole"
    mistake_rate = MismulitmistakeRate().apply(y, probs, y_mask)
    mistake_rate.name = "single error within note"
    cost.name = 'final_cost'

    model = Model(cost)
    cg = ComputationGraph([cost])
    step_rule = CompositeRule(
        [RemoveNotFinite(), StepClipping(30.0), Adam(learning_rate=lr), StepClipping(6.0),
         RemoveNotFinite()])
    algorithm = GradientDescent(step_rule=step_rule, cost=cost, params=cg.parameters)
    extensions = [SharedVariableModifier(parameter=lr,
                                         function=lambda n, v: float32(0.7 * v) if n % 700 == 0 else v),
                  FinishAfter(after_n_epochs=epochs),
                  TrainingDataMonitoring(
                      [cost, error_rate, mistake_rate, ],  # hidden_states, debug_val, param_nans,
                      # aggregation.mean(algorithm.total_gradient_norm)],  #+ params,
                      prefix="train",
                      after_epoch=False, every_n_batches=40),
                  Timing(),
                  Printing(),
                  ProgressBar()]
    if test is not None:
        extensions.append(DataStreamMonitoring(
            [cost, error_rate, mistake_rate],
            data_stream=test,
            updates=cg.updates,
            prefix="test", after_epoch=False, every_n_batches=40))

    if bokeh:
        extensions.append(Plot(
            'Pretrain RNN',
            channels=[
                ['train_error on note as a whole', 'train_single error within note',
                 'test_error on note as a whole',
                 'test_single error within note'],
                ['train_rbm_cost'],
                # ['train_total_gradient_norm'],
            ]))

    main_loop = MainLoop(algorithm=algorithm,
                         data_stream=train,
                         model=model,
                         extensions=extensions
                         )
    return main_loop


def pretrain_rbm(train, rnnrbm, test=None, epochs=900, bokeh=True, load_path=None):
    cdk = theano.shared(1)
    lr = theano.shared(float32(0.1))

    cost, v_sample = rnnrbm.rbm.cost(visible=x, k=cdk, batch_size=x.shape[0], mask=x_mask)
    error_rate = MismulitclassificationRate().apply(x, v_sample[-1], x_mask)
    error_rate.name = "error on note as a whole"
    mistake_rate = MismulitmistakeRate().apply(x, v_sample[-1], x_mask)
    mistake_rate.name = "single error within note"

    cost.name = 'rbm_cost'
    model = Model(cost)
    cg = ComputationGraph([cost])
    step_rule = CompositeRule(
        [RemoveNotFinite(), StepClipping(30.0), Adam(learning_rate=lr), StepClipping(6.0),
         RemoveNotFinite()])
    gradients = dict(equizip(cg.parameters, T.grad(cost, cg.parameters, consider_constant=[v_sample])))
    algorithm = GradientDescent(step_rule=step_rule, gradients=gradients, cost=cost,
                                params=cg.parameters)
    algorithm.add_updates(cg.updates)
    extensions = [SharedVariableModifier(parameter=cdk,
                                         function=lambda n, v: rbm_cdk[n] if rbm_cdk.get(n) else v),
                  SharedVariableModifier(parameter=lr,
                                         function=lambda n, v: float32(0.7 * v) if n % 1500 == 0 else v),
                  FinishAfter(after_n_epochs=epochs),
                  TrainingDataMonitoring(
                      [cost, error_rate, mistake_rate, ],  # hidden_states, debug_val, param_nans,
                      # aggregation.mean(algorithm.total_gradient_norm)],  #+ params,
                      prefix="train",