import argparse

parser = argparse.ArgumentParser(description='Train a neural network to decode a code.',
                                 formatter_class=argparse.RawDescriptionHelpFormatter,
                                 epilog='''\
''')
parser.add_argument('dist', type=int,
                    help='the distance of the code')
parser.add_argument('out', type=str,
                    help='the name of the output file (used as a prefix for the log file as well)')
parser.add_argument('--trainset', type=str,
                    help='the name of the training set file (generated by `generate_training_data.py`); if not specified --onthefly is assumed')
parser.add_argument('--onthefly', type=int, nargs=2, default=[2000000, 50000],
                    help='generate the training set on the fly, specify training and validation size (default: %(default)s)')
parser.add_argument('--prob', type=float, default=0.9,
                    help='the probability of no error on the physical qubit when generating training data (considered only if --onthefly is present) (default: %(default)s)')
parser.add_argument('--load', type=str, default='',
                    help='the file from which to load a pretrained model weights (optional, requires correct hyperparameters)')
parser.add_argument('--eval', action='store_true',
                    help='if present, calculate the fraction of successful corrections based on sampling the NN using the validation set')
parser.add_argument('--giveup', type=int, default=1000000,
                    help='after how many samples to give up decoding a given test error vector (considered only if --eval is present) (default: %(default)s)')
parser.add_argument('--batch', type=int, default=512,
                    help='the batch size (default: %(default)s)')
parser.add_argument('--epochs', type=int, default=20,
                    help='the number of epochs (default: %(default)s)')
parser.add_argument('--learningrate', type=float, default=0.002,
                    help='the learning rate (default: %(default)s)')
parser.add_argument('--hact', type=str, default='tanh',
                    help='the activation for hidden layers (default: %(default)s)')
parser.add_argument('--act', type=str, default='sigmoid',
                    help='the activation for the output layer (default: %(default)s)')
parser.add_argument('--loss', type=str, default='binary_crossentropy',
                    help='the loss to be optimized (default: %(default)s)')
parser.add_argument('--layers', type=float, default=[4, 4, 4], nargs='+',
                    help='the list of sizes of the hidden layers (as a factor of the output layer) (default: %(default)s)')
parser.add_argument('--normcentererr', action='store_true',
                    help='if present, zero-center and normalize the error training and evaluation data (should be done in advance for pregenerated data)')
parser.add_argument('--normcenterstab', action='store_true',
                    help='if present, zero-center and normalize the stabilizer training and evaluation data (should be done in advance for pregenerated data)')
parser.add_argument('--batchnorm', type=float, default=0,
                    help='if nonzero, batchnormalize each layer with the given momentum (default: %(default)s)')
parser.add_argument('--Zstab', action='store_true',
                    help='if present, include the Z stabilizer in the neural network')
parser.add_argument('--Xstab', action='store_true',
                    help='if present, include the X stabilizer in the neural network')

args = parser.parse_args()
print(args)

from neural import create_model, data_generator, do_normcenterstab, undo_normcentererr, smart_sample
from codes import ToricCode
import numpy as np
import tqdm


if args.trainset:
    f = np.load(args.trainset)
    x_test = []
    y_test = []
    if args.Zstab:
        x_test.append(f['arr_4'])
        y_test.append(f['arr_5'])
    if args.Xstab:
        x_test.append(f['arr_6'])
        y_test.append(f['arr_7'])
    x_test = np.hstack(x_test)
    y_test = np.hstack(y_test)

model = create_model(L=args.dist,
                     hidden_sizes=args.layers,
                     hidden_act=args.hact,
                     act=args.act,
                     loss=args.loss,
                     Z=args.Zstab, X=args.Xstab,
                     learning_rate=args.learningrate,
                     normcentererr_p=args.prob if args.normcentererr else None,
                     batchnorm=args.batchnorm
                    )
L = args.dist
code = ToricCode(L)
out_dimZ = 2*L**2 * args.Zstab
out_dimX = 2*L**2 * args.Xstab
in_dim = L**2 * (args.Xstab+args.Zstab)
H = code.H(args.Zstab,args.Xstab)


if args.load:
    model.load_weights(args.load)
if args.epochs:
    if args.trainset:
        raise NotImplementedError("This is still using the OLD keras API. Update it!")
        x_train = []
        y_train = []
        if args.Zstab:
            x_train.append(f['arr_0'])
            y_train.append(f['arr_1'])
        if args.Xstab:
            x_train.append(f['arr_2'])
            y_train.append(f['arr_3'])
        x_train = np.hstack(x_train)
        y_train = np.hstack(y_train)
        hist = model.fit(x_train, y_train,
                         nb_epoch=args.epochs,
                         batch_size=args.batch,
                         validation_data=(x_test, y_test)
                        )
    else:
        dat = data_generator(H, out_dimZ, out_dimX, in_dim, args.prob, args.batch,
                             normcenterstab=args.normcenterstab, normcentererr=args.normcentererr)
        val = data_generator(H, out_dimZ, out_dimX, in_dim, args.prob, args.batch,
                             normcenterstab=args.normcenterstab, normcentererr=args.normcentererr)
        hist = model.fit_generator(dat, args.onthefly[0]//args.batch, args.epochs,
                                   validation_data=val, validation_steps=args.onthefly[1]//args.batch)
    model.save_weights(args.out)
    with open(args.out+'.log', 'w') as f:
        f.write(str((hist.params, hist.history)))
if args.eval:
    E = ToricCode(L).E(args.Zstab,args.Xstab)
    both = args.Zstab and args.Xstab
    if both:
        Hz = ToricCode(L).H(True, False)
        Hx = ToricCode(L).H(False, True)
    outlen = 2*L**2*(args.Zstab+args.Xstab)
    inlen = L**2*(args.Zstab+args.Xstab)
    c = cz = cx = 0
    if args.trainset:
        stabflipgen = zip(x_test, y_test)
        size = len(y_test)
    else:
        size = args.onthefly[1]
        stabflipgen = data_generator(H, out_dimZ, out_dimX, in_dim, args.prob, batch_size=1, size=size)
    full_log = np.zeros((size, E.shape[0]+args.Zstab+args.Xstab), dtype=int)
    for i, (stab, flips) in tqdm.tqdm(enumerate(stabflipgen), total=size):
        if args.normcenterstab:
            stab_normed = do_normcenterstab(stab, args.prob)
            pred = model.predict(stab_normed).ravel()
        else:
            pred = model.predict(stab).ravel()
        if args.normcentererr:
            pred = undo_normcentererr(pred, args.prob)
        stab = stab.ravel()
        flips = flips.ravel()
        sample = pred>np.random.uniform(size=outlen)
        if both:
            attemptsZ = smart_sample(Hz, stab[:inlen//2], pred[:outlen//2], sample[:outlen//2], args.giveup)
            attemptsX = smart_sample(Hx, stab[inlen//2:], pred[outlen//2:], sample[outlen//2:], args.giveup)
        else:
            attempts = smart_sample(H, stab, pred, sample, args.giveup)
        errors = E.dot((sample+flips)%2)%2
        if np.any(errors) or np.any(stab!=H.dot(sample)%2):
            c += 1
            if both:
                cz += np.any(errors[:len(errors)//2])
                cx += np.any(errors[len(errors)//2:])
        if both:
            full_log[i,:-2] = errors
            full_log[i,-2] = attemptsZ
            full_log[i,-1] = attemptsX
        else:
            full_log[i,:-1] = errors
            full_log[i,-1] = attempts
    with open(args.out+'.eval', 'w') as f:
        if both:
            f.write(str(((1-c/size),(1-cz/size),(1-cx/size))))
        else:
            f.write(str(((1-c/size),)))
    np.savetxt(args.out+'.eval.log', full_log, fmt='%d')
