#  ___________________________________________________________________________
#
#  Pyomo: Python Optimization Modeling Objects
#  Copyright 2017 National Technology and Engineering Solutions of Sandia, LLC
#  Under the terms of Contract DE-NA0003525 with National Technology and 
#  Engineering Solutions of Sandia, LLC, the U.S. Government retains certain 
#  rights in this software.
#  This software is distributed under the 3-clause BSD License.
#  ___________________________________________________________________________

import logging

# This is an auto geometry generator for quadratic ROM
import numpy as np
from six import StringIO
from pyomo.contrib.trustregion.cache import GeometryCache

logger = logging.getLogger('pyomo.contrib.trustregion')

def _pset_to_mat(pset, lx):
    mat = []
    for p in pset:
        basisValue = [1]
        for i1 in range(0, lx):
            basisValue.append(p[i1])
        for i1 in range(0, lx):
            for i2 in range(i1, lx):
                basisValue.append(p[i1]*p[i2])
        mat.append(basisValue)
    return np.matrix(mat)

def generate_quadratic_rom_geometry(lx, NUM_SEEDS=None):
    if lx in GeometryCache:
        condOpt, txt = GeometryCache[lx]
        psetOpt = np.loadtxt(StringIO(txt))
        if psetOpt.ndim < 2:
            psetOpt = psetOpt.reshape(psetOpt.size, 1)
        matOpt = _pset_to_mat(psetOpt, lx)
        if NUM_SEEDS is None:
            return condOpt, psetOpt, matOpt
        logger.info(
            "Loading cached geometry with condition number %f" % (condOpt,))
    else:
        condOpt = np.inf
        psetOpt = None
        matOpt = None

    # 500 seems to be a reasonable number of iterations if we are
    # starting from scratch
    if NUM_SEEDS is None:
        NUM_SEEDS = 5000

    logger.info("Generating %d random geometries for LX=%s" % (NUM_SEEDS,lx))
    dim = int((lx*lx+lx*3)/2 + 1)
    x1 = np.zeros(lx)
    for i in range(0,NUM_SEEDS):
        #TODO: the following line returns error
        # ValueError: cannot reshape array of size 0 into shape (0)
        # for np.random.multivariate_normal(np.zeros(0),np.eye(0),0)
        pset = np.random.multivariate_normal(x1,np.eye(lx),dim-1)
        for j in range(dim-1):
            pset[j] = pset[j]/np.linalg.norm(pset[j])
        pset = np.append(pset,[x1],axis=0)
        mat = _pset_to_mat(pset, lx)
        cond = np.linalg.cond(mat)
        if(cond<condOpt):
            logger.info("new: %6d : %10.4f : %10.4f" % ( i, condOpt, cond))
            condOpt = cond
            psetOpt = pset
            matOpt = mat
    if(psetOpt is None):
        logger.error("lx = %d failed in initialization "
                     "(no non-singular geometries found)!\n" % lx)
    return condOpt, psetOpt, matOpt


if __name__ == '__main__':
    logger.setLevel(logging.INFO)
    import sys, os
    if len(sys.argv) < 2:
        print("Usage: %s [LX] [COUNT]")
        print("   LX:    single number or range in the form nn:nn")
        print("   COUNT: number of random matricies to generate")

    from os.path import abspath, dirname, join
    from inspect import getfile, currentframe
    CACHE_FILE = join(dirname(abspath(getfile(currentframe()))), 'cache.py')
    if not os.path.isfile (CACHE_FILE) or not os.access(CACHE_FILE, os.W_OK):
        logger.error(
            "Cannot write to the geometry cache (%s).  "
            "This utility is only expected to be run as a script in "
            "editable (development) source trees." % (CACHE_FILE,))
        sys.exit()

    if len(sys.argv) > 1:
        LX_SET = range(*tuple(int(i) for i in sys.argv[1].split(':')))
    else:
        LX_SET = range(1,24)
    if len(sys.argv) > 2:
        COUNT = int(sys.argv[2])
    else:
        COUNT = None
    for lx in LX_SET:
        local_count = COUNT if COUNT else 10000*lx
        cond, pset, mat = generate_quadratic_rom_geometry(lx, local_count)
        if pset is not None:
            txt = StringIO()
            np.savetxt(txt, pset)
            GeometryCache[lx] = (cond, txt.getvalue())
            with open(CACHE_FILE,'w') as FILE:
                FILE.write("""
# ___________________________________________________________________________
#
# Pyomo: Python Optimization Modeling Objects
# Copyright 2017 National Technology and Engineering Solutions of Sandia, LLC
# Under the terms of Contract DE-NA0003525 with National Technology and
# Engineering Solutions of Sandia, LLC, the U.S. Government retains certain
# rights in this software.
# This software is distributed under the 3-clause BSD License.
# ___________________________________________________________________________
#
# Cache of autogenerated quadratic ROM geometries
#
# THIS FILE IS AUTOGENERATED BY pyomo.contrib.trustregion.GeometryGenerator
#
# DO NOT EDIT BY HAND
#
GeometryCache = {
""")
                for i in sorted(GeometryCache):
                    FILE.write('    %d: ( %f, """%s""" ),\n'
                               % (i, GeometryCache[i][0], GeometryCache[i][1]))
                FILE.write("}\n")
            logger.info("Condition number: lx = %d is %f\n" % (lx,cond))