import HONHelpers as hon
import caffe
import google.protobuf.text_format
import os
from datetime import datetime
import argparse

def check_base_lr(solver_in):
    solver = caffe.io.caffe_pb2.SolverParameter()
    with open(solver_in) as f:
        google.protobuf.text_format.Merge(f.read(), solver)
    return solver.base_lr

def set_image_source_train_val(train_val_in, train_val_out, source_train, source_val):
    net = caffe.io.caffe_pb2.NetParameter()
    with open(train_val_in) as f:
        google.protobuf.text_format.Merge(f.read(), net)
    # phases correspond to 0 train, 1 test
    # layers 0 and 1 should be about data
    assert net.layer[0].include[0].phase == 0 and net.layer[0].name == u'data'
    assert net.layer[1].include[0].phase == 1 and net.layer[0].name == u'data'
    net.layer[0].image_data_param.source = source_train
    net.layer[1].image_data_param.source = source_val

    with open(train_val_out, 'w') as f:
        f.write(str(net))


def set_num_output_in_deploy(deploy_in, deploy_out, num_outputs):
    net = caffe.io.caffe_pb2.NetParameter()
    with open(deploy_in) as f:
        google.protobuf.text_format.Merge(f.read(), net)
    assert net.layer[-2].type == 'InnerProduct'
    assert net.layer[-1].type == 'Softmax'
    net.layer[-2].inner_product_param.num_output = num_outputs
    net.layer[-2].name = 'fc8_' + str(num_outputs)
    with open(deploy_out, 'w') as f:
        f.write(str(net))

def set_num_output_in_train_val(train_val_in, train_val_out, num_outputs):
    net = caffe.io.caffe_pb2.NetParameter()
    with open(train_val_in) as f:
        google.protobuf.text_format.Merge(f.read(), net)
    assert net.layer[-3].type == 'InnerProduct'
    net.layer[-3].inner_product_param.num_output = num_outputs
    net.layer[-3].name = 'fc8_' + str(num_outputs)
    with open(train_val_out, 'w') as f:
        f.write(str(net))


def set_net_solver(solver_in, solver_out, net_spec, snapshot_prefix, base_lr):
    solver = caffe.io.caffe_pb2.SolverParameter()
    with open(solver_in) as f:
        google.protobuf.text_format.Merge(f.read(), solver)
    solver.snapshot_prefix = snapshot_prefix
    solver.base_lr = base_lr
    solver.net = net_spec
    with open(solver_out, 'w') as f:
        f.write(str(solver))

if __name__ == '__main__':   
    parser = argparse.ArgumentParser()
    parser.add_argument("experiment_name", 
        help="Directory relative to the experiments/ directory \
        containing a hon_vgg_train_val.prototxt and solver.protoxt, \
        within which individual webcam experiments should be created", type=str)
    args = parser.parse_args()
    EXPERIMENT_ROOT = os.path.join(hon.experiment_root, args.experiment_name)

    for webcam in hon.webcams + ['all']:
        # make directory for each webcam
        webcam_root = os.path.join(EXPERIMENT_ROOT, webcam)
        if not os.path.exists(webcam_root):
            os.makedirs(webcam_root)
        # change source of images for the train_val prototxt
        # change solver.prototxt to look at correct net specification and write to correct location
        train_val_in = os.path.join(EXPERIMENT_ROOT, 'hon_vgg_train_val.prototxt')
        train_val_out = os.path.join(webcam_root, 'hon_vgg_train_val.prototxt')
        solver_in = os.path.join(EXPERIMENT_ROOT, 'solver.prototxt')
        solver_out = os.path.join(webcam_root, 'solver.prototxt')
        source_train = os.path.join(webcam_root, 'train.txt')    
        source_val = os.path.join(webcam_root, 'val.txt')
        base_lr = 1e-9 #  
        set_image_source_train_val(train_val_in, train_val_out, source_train, source_val)
        snapshot_prefix = os.path.join(webcam_root, 'model')
        set_net_solver(solver_in, solver_out, train_val_out, snapshot_prefix, base_lr)

    readme = os.path.join(EXPERIMENT_ROOT, 'README')
    with open(readme, 'w') as f:
        f.write("The solver.prototxt and hon_vgg_train_val.prototxt are generated by individual_webcam_experiments, {}".format(str(datetime.now())))