#!/usr/bin/env python

# utility for computing RMS difference of two models

import argparse
import sys

import numpy as np

def _parse_arguments():
    parser = argparse.ArgumentParser(
        description='Computes RMS difference of two models')

    # defines keyword arguments
    parser.add_argument('--format', default='sem')
    parser.add_argument('--parameters', default='vp,vs')
    parser.add_argument('--max-rms', default='0.01')
    parser.add_argument('--nproc', default=1, type=int)

    # defines positional arguments
    parser.add_argument('path1')
    parser.add_argument('path2')

    return parser.parse_args()

def load(path, format, parameters, nproc):
    if format in ['SEM', 'sem']:
        from seisflows.plugins.io import sem as io
    else:
        raise ValueError

    model = {}
    for key in parameters:
        model[key] = []
        for iproc in range(nproc):
            model[key] += [io.read(path, key, iproc)]
    return model


def stack(model):
    nn = 0
    for islice in model:
        nn += len(islice)

    vv = np.zeros(nn)
    jj = 0
    for islice in model:
        ii = jj
        jj += len(islice)
        vv[ii:jj] = islice[:]

    return vv


if __name__ == '__main__':
    # parse arguments
    args = _parse_arguments()

    if 'parameters' in args:
        parameters = args.parameters.split(',')
    

    # parse max-rms argument
    max_rms = []
    if 'max_rms' in args:
        for string in args.max_rms.split(','):
            max_rms += [float(string)]
    nn = len(parameters)
    if len(parameters) > len(max_rms) == 1:
       max_rms *= len(parameters)
    if len(parameters) != len(max_rms):
        raise Exception


    # read models
    model1 = load(args.path1, args.format, parameters, args.nproc)
    model2 = load(args.path2, args.format, parameters, args.nproc)

    _status = []
    for ii in range(nn):
        v1 = stack(model1[parameters[ii]])
        v2 = stack(model2[parameters[ii]])

        # compute rms
        rms = np.linalg.norm(v1-v2)
        if np.linalg.norm(v1) > 0:
            rms /= np.linalg.norm(v2)

        if rms < max_rms[ii]:
            print '\nSUCCESS: %s, %e, %e\n' % (parameters[ii], rms, max_rms[ii])
            _status += [0]
        else:
            print '\nFAILED:  %s, %e, %e\n' % (parameters[ii], rms, max_rms[ii])
            _status += [1]


    if any(_status):
        sys.exit(-1)

