#!/usr/bin/env python

#
# This file is part of TransportMaps.
#
# TransportMaps is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# TransportMaps is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with TransportMaps.  If not, see <http://www.gnu.org/licenses/>.
#
# Transport Maps Library
# Copyright (C) 2015-2017 Massachusetts Institute of Technology
# Uncertainty Quantification group
# Department of Aeronautics and Astronautics
#
# Author: Transport Map Team
# Website: transport-maps.mit.edu
# Support: transport-maps.mit.edu/qa/
#

from __future__ import print_function

import sys
import getopt
import os
import os.path
import shutil
import time
import datetime
import logging
import dill as pickle
import numpy as np
import scipy.stats as stats
import TransportMaps as TM
import TransportMaps.CLI as TMCLI
import TransportMaps.Maps as MAPS
import TransportMaps.Distributions as DIST
import TransportMaps.Distributions.Decomposable as DECDIST
import TransportMaps.Algorithms as ALG
import TransportMaps.XML as TMXML

sys.path.append(os.getcwd())

# Data storage object
stg = type('', (), {})()

def usage():
    usage_str = """
Usage: tmap-sequential-tm [-h -f -I] 
  --dist=DIST --output=OUTPUT 
  (--mtype=MTYPE --span=SPAN --btype=BTYPE --order=ORDER)
    / (--map-0-descr=MAP_DESCR --map-descr=MAP_DESCR)
  --qtype=QTYPE --qnum=QNUM
  [--tol=TOL --with-reg=REG --ders=DERS]
  [(--hyper-mtype=MTYPE --hyper-span=SPAN --hyper-btype=BTYPE --hyper-order=ORDER)
    / (--hyper-map-descr=MAP_DESCR)
   --hyper-qtype=QTYPE --hyper-qnum=QNUM
   --hyper-tol=TOL --hyper-with-reg=REG]
  [--safe-mode --reload --log=LOG --nprocs=NPROCS --batch=BATCH]
"""
    print(usage_str)

def description():
    docs_monotone_str = \
        '  --mtype=MTYPE           monotone format for the transport\n' + \
        TMCLI.print_avail_options(TMCLI.AVAIL_MONOTONE,'                          ')
    docs_span_str = \
        '  --span=SPAN             span type for all the components of the map\n' + \
        TMCLI.print_avail_options(TMCLI.AVAIL_SPAN,'                          ')
    docs_btype_str = \
        '  --btype=BTYPE           basis types for all the components of the map\n' + \
        TMCLI.print_avail_options(TMCLI.AVAIL_BTYPE,'                          ')
    docs_qtype_str = \
        '  --qtype=QTYPE           quadrature type for the discretization of ' + \
        'the KL-divergence\n' + \
        TMCLI.print_avail_options(TMCLI.AVAIL_QTYPE,'                          ')
    docs_ders_str = \
        '  --ders=DERS             derivatives to be used in the optimization\n' + \
        TMCLI.print_avail_options(TMCLI.AVAIL_DERS,'                          ')
    docs_hyper_monotone_str = \
        '  --hyper-mtype=MTYPE     monotone format for the transport\n' + \
        TMCLI.print_avail_options(TMCLI.AVAIL_MONOTONE,'                          ')
    docs_hyper_span_str = \
        '  --hyper-span=SPAN       span type for all the components of the map\n' + \
        TMCLI.print_avail_options(TMCLI.AVAIL_SPAN,'                          ')
    docs_hyper_btype_str = \
        '  --hyper-btype=BTYPE     basis types for all the components of the map\n' + \
        TMCLI.print_avail_options(TMCLI.AVAIL_BTYPE,'                          ')
    docs_hyper_qtype_str = \
        '  --hyper-qtype=QTYPE     quadrature type for the discretization of ' + \
        'the KL-divergence\n' + \
        TMCLI.print_avail_options(TMCLI.AVAIL_QTYPE,'                          ')
    docs_log_str = \
        '  --log=LOG               log level (default=30). Uses package logging.\n' + \
        TMCLI.print_avail_options(TMCLI.AVAIL_LOGGING,'                          ')

    docs_str = """DESCRIPTION
Given a file (--dist) storing the target distribution, produce the transport map that
pushes forward the base distribution (default: standard normal) to the target distribution,
using the algorithm for sequential low-dimensional couplings.
The input distribution must be SequentialHiddenMarkcovChainDistribution with
hdim hyperparameters and sdim state dimension.
All files involved are stored and loaded using the python package dill.

OPTIONS - input/output:
  --dist=DIST             path to the file containing the target distribution 
  --output=OUTPUT         path to the output file containing the transport map,  
                          the base distribution, the target distribution and all 
                          the additional parameters used for the construction 
OPTIONS - map description (using default maps):
""" + docs_monotone_str + docs_span_str + docs_btype_str + \
"""  --order=ORDER           order of the transport map
OPTIONS - map description (manual):
  --map-0-descr=MAP_DESCR XML file with the skeleton of the zero-th transport map (hdim+sdim)
  --map-descr=MAP_DESCR   XML file with the skeleton of the transport map (hdim+2*sdim)
OPTIONS - KL-minimization solver:
""" + docs_qtype_str + \
"""  --qnum=QNUM             quadrature level
  --tol=TOL               optimization tolerance (default: 1e-4)
  --reg=REG               a float L2 regularization parameter
                          (default: no regularization)
""" + docs_ders_str + \
"""OPTIONS - hyper-parameters map description (using default maps):
""" + docs_hyper_monotone_str + docs_hyper_span_str + docs_hyper_btype_str + \
"""  --hyper-order=ORDER     order of the transport map
OPTIONS - hyper-parameters map description (manual):
  --hyper-map-descr=MAP_DESCR   XML file containing the skeleton of the transport map
OPTIONS - regression solver:
""" + docs_hyper_qtype_str + \
"""  --hyper-qnum=QNUM       quadrature level
  --hyper-tol=TOL         optimization tolerance (default: 1e-4)
  --hyper-reg=REG         a float L2 regularization parameter
                          (default: no regularization)
OPTIONS - other:
  --safe-mode             always store intermediate maps (allows for restarting)
  --reload                automatically resume from stored data (conflicts with -f)
""" + docs_log_str + \
"""  --nprocs=NPROCS         number of processors to be used (default=1)
  --batch=BATCH           list of batch sizes for function evaluation, gradient
                          evaluation and Hessian evaluation
  -v                      verbose output (not affecting --log)
  -f                      force overwrite of OUTPUT file
  -I                      enter interactive mode after finishing
  -h                      print this help
"""
    print(docs_str)

def full_usage():
    usage()

def full_doc():
    full_usage()
    description()

##################### INPUT PARSING #####################
argv = sys.argv[1:]
INTERACTIVE = False
# I/O
DIST_FNAME = None
OUT_FNAME = None
BASE_DIST_FNAME = None
FORCE = False
# Map type
MONOTONE = None
SPAN = None
BTYPE = None
ORDER = None
MAP_0_DESCR = None
MAP_DESCR = None
# Quadrature type
stg.QTYPE = None
stg.QNUM = None
# KL-minimization Solver options
stg.TOL = 1e-4
stg.REG = None
stg.DERS = 2
# Hyper Map type
HYPER_MONOTONE = None
HYPER_SPAN = None
HYPER_BTYPE = None
HYPER_ORDER = None
HYPER_MAP_DESCR = None
# Hyper Quadrature type
stg.HYPER_QTYPE = None
stg.HYPER_QNUM = None
# Regression Solver options
stg.HYPER_TOL = 1e-4
stg.HYPER_REG = None
# Safe mode
SAFE_MODE = False
RELOAD = False
# Logging
VERBOSE = False
LOGGING_LEVEL = 30 # Warnings
# Parallelization
NPROCS = 1
BATCH_SIZE = [None, None, None]
try:
    opts, args = getopt.getopt(argv,"hfvI",
                               [
                                   # I/O
                                   "dist=", "output=", "base-dist=",
                                   # Map type
                                   "mtype=", "span=", "btype=", "order=",
                                   "map-0-descr=", "map-descr=",
                                   # Quadrature type
                                   "qtype=", "qnum=",
                                   # KL-minimization solver options
                                   "tol=", "reg=", "ders=",
                                   # Hyper map type
                                   "hyper-mtype=", "hyper-span=", "hyper-btype=",
                                   "hyper-order=",
                                   "hyper-map-descr=",
                                   # Hyper Quadrature type
                                   "hyper-qtype=", "hyper-qnum=",
                                   # Regression solver options
                                   "hyper-tol=", "hyper-reg=",
                                   # Safe mode
                                   "safe-mode", "reload",
                                   # Logging
                                   "log=",
                                   # Parallelization and batching option
                                   "nprocs=", "batch="
                               ])
except getopt.GetoptError:
    full_usage()
    raise
for opt, arg in opts:
    if opt == '-h':
        full_doc()
        sys.exit()

    # Force overwrite
    elif opt == '-f':
        FORCE = True

    # Verbose
    elif opt == '-v':
        VERBOSE = True
        
    # Interactive
    elif opt == '-I':
        INTERACTIVE = True
        
    # I/O
    elif opt in ['--dist']:
        DIST_FNAME = arg
    elif opt in ['--output']:
        OUT_FNAME = arg
    elif opt in ['--base-dist']:
        BASE_DIST_FNAME = arg
        
    # Map type
    elif opt in ['--mtype']:
        MONOTONE = arg
    elif opt in ['--span']:
        SPAN = arg
    elif opt in ['--btype']:
        BTYPE = arg
    elif opt in ['--order']:
        ORDER = int(arg)
    elif opt in ['--map-0-descr']:
        MAP_0_DESCR = arg
    elif opt in ['--map-descr']:
        MAP_DESCR = arg

    # Quadrature type
    elif opt in ['--qtype']:
        stg.QTYPE = int(arg)
    elif opt in ['--qnum']:
        stg.QNUM = [int(q) for q in arg.split(',')]
        
    # KL-minimization solver options
    elif opt in ['--tol']:
        stg.TOL = float(arg)
    elif opt in ['--reg']:
        stg.REG = {'type': 'L2',
                   'alpha': float(arg)}
    elif opt in ['--ders']:
        stg.DERS = int(arg)

    # Hyper-parameters Map type
    elif opt in ['--hyper-mtype']:
        HYPER_MONOTONE = arg
    elif opt in ['--hyper-span']:
        HYPER_SPAN = arg
    elif opt in ['--hyper-btype']:
        HYPER_BTYPE = arg
    elif opt in ['--hyper-order']:
        HYPER_ORDER = int(arg)
    elif opt in ['--hyper-map-descr']:
        HYPER_MAP_DESCR = arg

    # Hyper-parameters Quadrature type
    elif opt in ['--hyper-qtype']:
        stg.HYPER_QTYPE = int(arg)
    elif opt in ['--hyper-qnum']:
        stg.HYPER_QNUM = [int(q) for q in arg.split(',')]
        
    # Regression solver options
    elif opt in ['--hyper-tol']:
        stg.HYPER_TOL = float(arg)
    elif opt in ['--hyper-reg']:
        stg.HYPER_REG = {'type': 'L2',
                   'alpha': float(arg)}

    # Safe mode
    elif opt in ['--safe-mode']:
        SAFE_MODE = True
    elif opt in ['--reload']:
        RELOAD = True
        
    # Logging
    elif opt in ['--log']:
        LOGGING_LEVEL = int(arg)

    # Parallelization and batching
    elif opt in ['--nprocs']:
        NPROCS = int(arg)
    elif opt in ['--batch']:
        BATCH_SIZE = [int(b) for b in arg.split(',')]
        
    else:
        raise ValueError("Option %s not recognized" % opt)

logging.basicConfig(level=LOGGING_LEVEL)

def tstamp_print(msg, *args, **kwargs):
    tstamp = datetime.datetime.fromtimestamp(
        time.time()
    ).strftime('%Y-%m-%d %H:%M:%S')
    print(tstamp + " " + msg, *args, **kwargs)

def filter_tstamp_print(msg, *args, **kwargs):
    if VERBOSE:
        tstamp_print(msg, *args, **kwargs)

def filter_print(*args, **kwargs):
    if VERBOSE:
        print(*args, **kwargs)

def safe_store(data, fname):
    # Backup copy
    if os.path.exists(fname):
        shutil.copyfile(fname, fname + '.bak')
    # Store data
    with open(fname, 'wb') as out_stream:
        pickle.dump(data, out_stream)
    # Remove backup
    try:
        os.remove(fname + '.bak')
    except FileNotFoundError:
        pass
        
# Check for required arguments
if None in [DIST_FNAME, OUT_FNAME]:
    usage()
    tstamp_print("ERROR: Options --dist and --output must be specified")
    sys.exit(3)
if FORCE and RELOAD:
    usage()
    tstamp_print("ERROR: options clash -f --reload")
    sys.exit(3)
if not RELOAD and None in [stg.QTYPE, stg.QNUM]:
    usage()
    tstamp_print("ERROR: Options --qtype and --qnum must be specified")
    sys.exit(3)
if not RELOAD and stg.QTYPE < 3:
    stg.QNUM = stg.QNUM[0]
map_descr_list = [MONOTONE, SPAN, BTYPE, ORDER]
if (MAP_DESCR is None and MAP_0_DESCR is not None) or \
   (MAP_DESCR is not None and MAP_0_DESCR is None):
    usage()
    tstamp_print("ERROR: --map-0-descr and --map-descr must be " + \
                 "specified simultaneously.")
    sys.exit(3)
if MAP_DESCR is None and None in map_descr_list:
    usage()
    tstamp_print("ERROR: Either options --mtype, --span, --btype, " + \
                 "--order are specified or option --map-descr is specified")
    sys.exit(3)
elif MAP_DESCR is not None and not all([s is None for s in map_descr_list]):
    usage()
    tstamp_print("ERROR: Either options --mtype, --span, --btype, " + \
                 "--order are specified or option --map-descr is specified")
    sys.exit(3)
BATCH_SIZE = [None] * (stg.DERS + 1 - len(BATCH_SIZE)) + BATCH_SIZE

if os.path.exists(OUT_FNAME):
    if not FORCE:
        if not RELOAD:
            sel = ''
            while sel not in ['y', 'Y', 'n', 'N', 'q']:
                if sys.version_info[0] == 3:
                    sel = input("The file %s already exists. " % OUT_FNAME + \
                                "Do you want to reload it? [Y/n/q] ") or 'y'
                else:
                    sel = raw_input("The file %s already exists. " % OUT_FNAME + \
                                    "Do you want to reload it? [Y/n/q] ") or 'y'
        if RELOAD or sel == 'y' or sel == 'Y':
            tstamp_print("Reloading data...")
            with open(OUT_FNAME, 'rb') as in_stream:
                stg = pickle.load(in_stream)
            RELOAD = True
        elif sel == 'q':
            tstamp_print("Terminating.")
            sys.exit(0)
        else:
            sel = ''
            while sel not in ['y', 'Y', 'n', 'N', 'q']:
                if sys.version_info[0] == 3:
                    sel = input("Do you want to overwrite? [y/N/q] ") or 'n'
                else:
                    sel = raw_input("Do you want to overwrite? [y/N/q] ") or 'n'
            if sel == 'n' or sel == 'N' or sel == 'q':
                tstamp_print("Terminating.")
                sys.exit(0)            

try:
    ##################### DATA LOADING #####################
    if not RELOAD:
        # Load target distribution
        with open(DIST_FNAME,'rb') as istr:
            stg.target_distribution = pickle.load(istr)
        dim = stg.target_distribution.dim

        if not isinstance(stg.target_distribution,
                          DECDIST.SequentialHiddenMarkovChainDistribution):
            raise ValueError("The input distribution must be a " + \
                             "SequentialHiddenMarkovChainDistribution")
        if len(stg.target_distribution.pi_list) < 1:
            raise ValueError("The input hidden Markov chain distribution " + \
                             "must perform at least " + \
                             "one transition (Z0,Z1). " + \
                             "If no transition is performed (Z0) use tmap-tm.")

        # Load base distribution
        stg.base_distribution = DIST.StandardNormalDistribution(dim)

        # Build integrator
        stg.integrator = ALG.SequentialHiddenMarkovChainIntegrator(
            stg.target_distribution.pi_hyper)

    hdim = stg.target_distribution.hyper_dim
    sdim = stg.target_distribution.state_dim

    # Start pool if necessary
    mpi_pool = None
    if NPROCS > 1:
        mpi_pool = TM.get_mpi_pool()
        mpi_pool.start(NPROCS)

    # Start integration
    nstart = stg.integrator.nsteps + 1
    try:
        for n, pi, ll in zip(
                range(nstart, stg.target_distribution.get_n_steps()),
                stg.target_distribution.pi_list[nstart:],
                stg.target_distribution.ll_list[nstart:] ):
            filter_tstamp_print("Step %5d" % n, end="")
            if n == 0:
                # Instantiate Transport Map
                if MAP_0_DESCR is not None:
                    tm_approx = TMXML.load_xml(MAP_0_DESCR)
                else:
                    if MONOTONE == 'linspan':
                        tm_approx = TM.Default_IsotropicMonotonicLinearSpanTriangularTransportMap(
                            hdim+sdim, ORDER, span=SPAN, btype=BTYPE, common_basis_flag=True)
                    elif MONOTONE == 'intexp':
                        tm_approx = TM.Default_IsotropicIntegratedExponentialTriangularTransportMap(
                            hdim+sdim, ORDER, span=SPAN, btype=BTYPE, common_basis_flag=True)
                    elif MONOTONE == 'intsq':
                        tm_approx = TM.Default_IsotropicIntegratedSquaredTriangularTransportMap(
                            hdim+sdim, ORDER, span=SPAN, btype=BTYPE)
                    else:
                        raise ValueError("Monotone type not recognized (linspan|intexp|intsq)")
                # Prepare solution parameters
                if isinstance(stg.QNUM, list):
                    qnum = stg.QNUM[:hdim+sdim]
                solve_params = {'qtype': stg.QTYPE, 'qparams': qnum,
                                'tol': stg.TOL, 'regularization': stg.REG,
                                'ders': stg.DERS}
                # Assimilate
                stg.integrator.assimilate(pi, ll, tm_approx, solve_params)
            else:
                # Instantiate Transport Map
                if MAP_DESCR is not None:
                    tm_approx = TMXML.load_xml(MAP_DESCR)
                else:
                    if MONOTONE == 'linspan':
                        tm_approx = TM.Default_IsotropicMonotonicLinearSpanTriangularTransportMap(
                            hdim+2*sdim, ORDER, span=SPAN, btype=BTYPE, common_basis_flag=True)
                    elif MONOTONE == 'intexp':
                        tm_approx = TM.Default_IsotropicIntegratedExponentialTriangularTransportMap(
                            hdim+2*sdim, ORDER, span=SPAN, btype=BTYPE, common_basis_flag=True)
                    elif MONOTONE == 'intsq':
                        tm_approx = TM.Default_IsotropicIntegratedSquaredTriangularTransportMap(
                            hdim+2*sdim, ORDER, span=SPAN, btype=BTYPE)
                    else:
                        raise ValueError("Monotone type not recognized (linspan|intexp)")
                # Prepare solution parameters
                solve_params = {'qtype': stg.QTYPE, 'qparams': stg.QNUM,
                                'tol': stg.TOL, 'regularization': stg.REG,
                                'ders': stg.DERS}
                # Instantiate Hyper-parameters transport map if necessary
                hyper_tm_approx = None
                if MAP_DESCR is not None:
                    hyper_tm_approx = TMXML.load_xml(HYPER_MAP_DESCR)
                else:
                    if HYPER_MONOTONE == 'linspan':
                        hyper_tm_approx = \
                            TM.Default_IsotropicMonotonicLinearSpanTriangularTransportMap(
                                hdim, HYPER_ORDER, span=HYPER_SPAN, btype=HYPER_BTYPE,
                                common_basis_flag=True)
                    elif HYPER_MONOTONE == 'intexp':
                        hyper_tm_approx = \
                            TM.Default_IsotropicIntegratedExponentialTriangularTransportMap(
                                hdim, HYPER_ORDER, span=HYPER_SPAN, btype=HYPER_BTYPE,
                                common_basis_flag=True)
                    elif HYPER_MONOTONE == 'intsq':
                        hyper_tm_approx = \
                            TM.Default_IsotropicIntegratedSquaredTriangularTransportMap(
                                hdim, HYPER_ORDER, span=HYPER_SPAN, btype=HYPER_BTYPE)
                    else:
                        raise ValueError("Monotone type not recognized (linspan|intexp)")
                # Prepare regression parameters if necessary
                reg_params = None
                if hyper_tm_approx is not None:
                    rho_hyper = DIST.StandardNormalDistribution(hdim)
                    reg_params = {'d': rho_hyper,
                                  'qtype': stg.HYPER_QTYPE, 'qparams': stg.HYPER_QNUM,
                                  'tol': stg.HYPER_TOL, 'regularization': stg.HYPER_REG}
                # Assimilate
                stg.integrator.assimilate(pi, ll, tm_approx, solve_params,
                                          hyper_tm_approx, reg_params)
            if SAFE_MODE: # Store partial integration
                safe_store(stg, OUT_FNAME)
            filter_print("    [DONE]")
    finally:
        if mpi_pool is not None:
            mpi_pool.stop()

    # Collect the smoothing map
    stg.tmap = stg.integrator.get_smoothing_map()

    stg.approx_base_distribution = DIST.PullBackTransportMapDistribution(stg.tmap, stg.target_distribution)
    stg.approx_target_distribution = DIST.PushForwardTransportMapDistribution(stg.tmap, stg.base_distribution)

    # STORE
    safe_store(stg, OUT_FNAME)

finally:
    if INTERACTIVE:
        from IPython import embed
        embed()
