import sys
from models import VGG16Base,VGG16Attention
import utils
import json

def train(model_name,parameters_file):

    # Read the parameters
    print('-----Parameter Reading Started-----\n')
    f = open(parameters_file)
    parameters = json.load(f)
    print('-----Parameter Reading Complete-----\n')


    # Instantiate the model
    print('-----Model Instantiation Started-----\n')
    if model_name == 'VGG16':
        model = VGG16Base(num_classes=parameters['num_classes'])
    else:
        model = VGG16Attention(num_classes=parameters['num_classes'])
    
    parameters['model_descriptor'] = model_name
    
    print('-----Model Instantiation Complete-----\n')
    
    # Perform training
    trained_model, TRAIN_LOSS,VAL_LOSS,VAL_CLASS_AUC = utils.perform_training(  model,
                                                                                parameters)
    return trained_model, TRAIN_LOSS,VAL_LOSS,VAL_CLASS_AUC

if __name__ == '__main__':
    
    parameters_file = './train_params.json'
    model_name = None
    
    # Use the model name coming from input
    if len(sys.argv) == 2:
        model_name = sys.argv[1]
    
    if len(sys.argv) == 3:
        model_name = sys.argv[1]
        parameters_file = sys.argv[2]

    # Check if correct model name is specified
    if model_name in ['VGG16','VGG16-ATTN']:
        # Initiate Training
        train(model_name,parameters_file)
    else:
        print('Invalid model name provided. Correct name VGG16 or VGG16-ATTN')
    
    


