#################################################################################
# The Institute for the Design of Advanced Energy Systems Integrated Platform
# Framework (IDAES IP) was produced under the DOE Institute for the
# Design of Advanced Energy Systems (IDAES), and is copyright (c) 2018-2021
# by the software owners: The Regents of the University of California, through
# Lawrence Berkeley National Laboratory,  National Technology & Engineering
# Solutions of Sandia, LLC, Carnegie Mellon University, West Virginia University
# Research Corporation, et al.  All rights reserved.
#
# Please see the files COPYRIGHT.md and LICENSE.md for full copyright and
# license information.
#################################################################################
from enum import Enum
import subprocess
from io import StringIO
import sys
import os
import numpy as np
import pandas as pd
import json

from pyomo.environ import Constraint, value, sin, cos, log, exp, Set, Reals
from pyomo.common.config import ConfigValue, In, Path, ListOf, Bool
from pyomo.common.tee import TeeStream
from pyomo.common.fileutils import Executable
from pyomo.common.tempfiles import TempfileManager

from idaes.core.surrogate.base.surrogate_base import SurrogateTrainer, SurrogateBase
from idaes.core.util.exceptions import ConfigurationError
import idaes.logger as idaeslog

# Set up logger
_log = idaeslog.getLogger(__name__)

# TODO: Adaptive sampling

alamo = Executable("alamo")

# Define mapping of Pyomo function names for expression evaluation
GLOBAL_FUNCS = {"sin": sin, "cos": cos, "log": log, "exp": exp}

# The values associated with these must match those expected in the .alm file
class Modelers(Enum):
    BIC = 1
    MallowsCp = 2
    AICc = 3
    HQC = 4
    MSE = 5
    SSEP = 6
    RIC = 7
    MADp = 8

class Screener(Enum):
    none = 0
    lasso = 1
    SIS = 2

supported_options = [
    "xfactor",
    "xscaling",
    "scalez",
    "monomialpower",
    "multi2power",
    "multi3power",
    "ratiopower",
    "expfcns",
    "linfcns",
    "logfcns",
    "sinfcns",
    "cosfcns",
    "constant",
    "grbfcns",
    "rbfparam",
    "modeler",
    "builder",
    "backstepper",
    "convpen",
    "screener",
    "ncvf",
    "sismult",
    "maxtime",
    "maxiter",
    "datalimitterms",
    "maxterms",
    "minterms",
    "numlimitbasis",
    "exclude",
    "ignore",
    "xisint",
    "zisint",
    "tolrelmetric",
    "tolabsmetric",
    "tolmeanerror",
    "tolsse",
    "mipoptca",
    "mipoptcr",
    "linearerror",
    "GAMS",
    "GAMSSOLVER",
    "solvemip",
    "print_to_screen",
]

"""
Excluded Options:
NSAMPLE, NVALSAMPLE, INITIALIZER, PRINT TO FILE, FUNFORM, NTRANS

Excluded Simulator options:
MAXSIM, MINPOINTS, MAXPOINTS, SAMPLER, SIMULATOR, PRESET, TOLMAXERROR, SIMIN,
SIMOUT
"""

# Headers from ALAMO trace file that should be common for all outputs
common_trace = [
    "filename",
    "NINPUTS",
    "NOUTPUTS",
    "INITIALPOINTS",
    "SET",
    "INITIALIZER",
    "SAMPLER",
    "MODELER",
    "BUILDER",
    "GREEDYBUILD",
    "BACKSTEPPER",
    "GREEDYBACK",
    "REGULARIZER",
    "SOLVEMIP",
]

class AlamoTrainer(SurrogateTrainer):
    """
    Standard SurrogateTrainer for ALAMO.

    This defines a set of configuration options for ALAMO along with
    methods to read and write the ALAMO input and output files and to call
    the ALAMO executable.

    Generally, options default to None to indicate that no entry will be
    written in the ALAMO input file. In this case, the default ALAMO settings
    will be used.
    """

    # The following ALAMO options are not (yet) supported
    # Returning model prediction, as we can do that in IDAES
    # Iniital sampling of data, as SurrogateModelTrainer should do that
    # Similarly, adaptive sampling is better handled in Python
    # Single validation set, due to limitations of current API
    # Custom basis functions are not yet implemented

    CONFIG = SurrogateTrainer.CONFIG()

    CONFIG.declare(
        "xfactor",
        ConfigValue(
            default=None,
            domain=ListOf(float),
            description="List of scaling factors for input variables.",
        ),
    )
    CONFIG.declare(
        "xscaling",
        ConfigValue(
            default=None,
            domain=Bool,
            description="Option to scale input variables.",
            doc="Option to scale input variables. If True and xfactors are not "
            "provided, ALAMO sets XFACTORS equal to the range of each input "
            "variable.",
        ),
    )
    CONFIG.declare(
        "scalez",
        ConfigValue(
            default=None, domain=Bool, description="Option to scale output variables."
        ),
    )

    # Basis function options
    CONFIG.declare(
        "monomialpower",
        ConfigValue(
            default=None,
            domain=ListOf(
                int, In(Reals - {0, 1})
            ),  # allow any float except for 0 and 1
            description="Vector of monomial powers considered in basis "
            "functions - cannot include 0 or 1.",
        ),
    )
    CONFIG.declare(
        "multi2power",
        ConfigValue(
            default=None,
            domain=ListOf(float),
            description="Vector of powers to be considered for pairwise "
            "combinations in basis functions.",
        ),
    )
    CONFIG.declare(
        "multi3power",
        ConfigValue(
            default=None,
            domain=ListOf(float),
            description="Vector of three variable combinations of powers to be "
            "considered as basis functions.",
        ),
    )
    CONFIG.declare(
        "ratiopower",
        ConfigValue(
            default=None,
            domain=ListOf(float),
            description="Vector of ratio combinations of powers to be considered "
            "in the basis functions.",
        ),
    )
    CONFIG.declare(
        "constant",
        ConfigValue(
            default=True,
            domain=Bool,
            description="Include constant basis function if True. Default = True",
        ),
    )
    CONFIG.declare(
        "linfcns",
        ConfigValue(
            default=True,
            domain=Bool,
            description="Include linear basis functions if True. Default = True",
        ),
    )
    CONFIG.declare(
        "expfcns",
        ConfigValue(
            default=None,
            domain=Bool,
            description="Include exponential basis functions if True.",
        ),
    )
    CONFIG.declare(
        "logfcns",
        ConfigValue(
            default=None,
            domain=Bool,
            description="Include logarithmic basis functions if True.",
        ),
    )
    CONFIG.declare(
        "sinfcns",
        ConfigValue(
            default=None,
            domain=Bool,
            description="Include sine basis functions if True.",
        ),
    )
    CONFIG.declare(
        "cosfcns",
        ConfigValue(
            default=None,
            domain=Bool,
            description="Include cosine basis functions if True.",
        ),
    )
    CONFIG.declare(
        "grbfcns",
        ConfigValue(
            default=None,
            domain=Bool,
            description="Include Gaussian radial basis functions if True.",
        ),
    )
    CONFIG.declare(
        "rbfparam",
        ConfigValue(
            default=None,
            domain=float,
            description="Multiplicative constant used in the Gaussian radial basis"
            " functions.",
        ),
    )
    CONFIG.declare(
        "custom_basis_functions",
        ConfigValue(
            default=None,
            domain=ListOf(str),
            description="List of custom basis functions to include in surrogate "
            "fitting.",
            doc="List of custom basis functions to include in surrogate model "
            "fitting. These should be in a form that can be rendered as a string "
            "that meets ALAMO's requirements.",
        ),
    )

    # Other fitting options
    CONFIG.declare(
        "modeler",
        ConfigValue(
            default=None,
            domain=In(Modelers),
            description="Fitness metric to be used for model building. Must be an "
            "instance of Modelers Enum.",
        ),
    )
    CONFIG.declare(
        "builder",
        ConfigValue(
            default=None,
            domain=Bool,
            description="If True, a greedy heuristic builds up a model "
            "by adding one variable at a time.",
            doc="If True, a greedy heuristic builds up a model by adding one "
            "variable at a time. This model is used as a starting point for "
            "solving an integer programming formulation according to the choice "
            "of modeler. If an optimizer is not available, the heuristic model "
            "will be the final model to be returned.",
        ),
    )
    CONFIG.declare(
        "backstepper",
        ConfigValue(
            default=None,
            domain=Bool,
            description="If set to 1, a greedy heuristic builds down a model by "
            "starting from the least squares model and removing one variable at "
            "a time.",
        ),
    )
    CONFIG.declare(
        "convpen",
        ConfigValue(
            default=None,
            domain=float,
            description="Convex penalty term to use if Modeler == SSEP or MADp.",
            doc="When MODELER is set to 6 or 8, a penalty consisting of the sum "
            "of square errors (SSEP) or the maximum absolute error (MADp) and a "
            "term penalizing model size is used for model building. The size of "
            "the model is weighted by convpen. If convpen=0, this metric reduces "
            "to the classical sum of square errors (SSEP) or the maximum absolute "
            "deviation (MADp).",
        ),
    )
    CONFIG.declare(
        "screener",
        ConfigValue(
            default=None,
            domain=In(Screener),
            description="Regularization method used to reduce the number of "
            "potential basis functions before optimization. Must be instance of "
            "Screener Enum.",
        ),
    )
    CONFIG.declare(
        "ncvf",
        ConfigValue(
            default=None,
            domain=int,
            description="Number of folds to be used for cross validation by the "
            "lasso screener. ALAMO will use a two-fold validation if fewer than "
            "10 data points are available. NCVF must be a nonnegative integer.",
        ),
    )
    CONFIG.declare(
        "sismult",
        ConfigValue(
            default=None,
            domain=int,
            description="This parameter must be non-negative and is used to "
            "determine the number of basis functions retained by the SIS "
            "screener. The number of basis functions retained equals the floor "
            "of SSISmult n ln(n), where n is the number of measurements "
            "available at the current ALAMO iteration.",
        ),
    )

    CONFIG.declare(
        "maxiter",
        ConfigValue(
            default=None,
            domain=int,
            description="Maximum number of ALAMO iterations. 1 = no adaptive "
            "sampling, 0 = no limit.",
            doc="Maximum number of ALAMO iterations. Each iteration begins with "
            "a model-building step. An adaptive sampling step follows if maxiter "
            "does not equal 1. If maxiter is set to a number less than or equal "
            "to 0, ALAMO will enforce no limit on the number of iterations.",
        ),
    )
    CONFIG.declare(
        "maxtime",
        ConfigValue(
            default=1000,
            domain=float,
            description="Maximum total execution time allowed in seconds. "
            "Default = 1000.",
            doc="Maximum total execution time allowed in seconds. This time "
            "includes all steps of the algorithm, including time to read problem, "
            "preprocess data, solve optimization subproblems, and print results.",
        ),
    )
    CONFIG.declare(
        "datalimitterms",
        ConfigValue(
            default=None,
            domain=Bool,
            description="Limit model terms to number of measurements.",
            doc="If True, ALAMO will limit the number of terms in the model to be "
            "no more than the number of data measurements; otherwise, no limit "
            "based on the number of data measurements will be placed. The user "
            "may provide an additional limit on the number of terms in the model "
            "through the maxterms and minterms options.",
        ),
    )
    CONFIG.declare(
        "maxterms",
        ConfigValue(
            default=None,
            domain=ListOf(int),
            description="List of maximum number of model terms to per output.",
            doc="Row vector of maximum terms allowed in the modeling of output "
            "variables. One per output variable, space separated. A −1 signals "
            "that no limit is imposed.",
        ),
    )
    CONFIG.declare(
        "minterms",
        ConfigValue(
            default=None,
            domain=ListOf(int),
            description="List of minimum number of model terms to per output.",
            doc="Row vector of minimum terms required in the modeling of output "
            "variables. One per output variable, space separated. A 0 signals "
            "that no limit is imposed.",
        ),
    )
    CONFIG.declare(
        "numlimitbasis",
        ConfigValue(
            default=True,  # default this to true to avoid numerical issues
            domain=Bool,
            description="Eliminate infeasible basis functions. Default = True",
            doc="If True, ALAMO will eliminate basis functions that are not "
            "numerically acceptable (e.g., log(x) will be eliminated if x may be "
            "negative); otherwise, no limit based on the number of data "
            "measurements will be placed. The user may provide additional limits "
            "on the the type and number of selected basis functions through the "
            "options exclude and groupcon.",
        ),
    )
    CONFIG.declare(
        "exclude",
        ConfigValue(
            default=None,
            domain=ListOf(int),
            description="List of inputs to exclude during building,",
            doc="Row vector of 0/1 flags that specify which input variables, if "
            "any, ALAMO should exclude during the model building process. All "
            "input variables must be present in the data but ALAMO will not "
            "include basis functions that involve input variables for which "
            "exclude equals 1. This feature does not apply to custom basis "
            "functions or RBFs.",
        ),
    )
    CONFIG.declare(
        "ignore",
        ConfigValue(
            default=None,
            domain=ListOf(int),
            description="List of outputs to ignore during building.",
            doc="Row vector of 0/1 flags that specify which output variables, "
            "if any, ALAMO should ignore. All output variables must be present in "
            "the data but ALAMO does not model output variables for which ignore "
            "equals 1.",
        ),
    )
    CONFIG.declare(
        "xisint",
        ConfigValue(
            default=None,
            domain=ListOf(int),
            description="List of inputs that should be treated as integers.",
            doc="Row vector of 0/1 flags that specify which input variables, if "
            "any, ALAMO should treat as integers. For integer inputs, ALAMO’s "
            "sampling will be restricted to integer values.",
        ),
    )
    CONFIG.declare(
        "zisint",
        ConfigValue(
            default=None,
            domain=ListOf(int),
            description="List of outputs that should be treated as integers.",
            doc="Row vector of 0/1 flags that specify which output variables, if "
            "any, ALAMO should treat as integers. For integer variables, ALAMO’s "
            "model will include the rounding of a function to the nearest integer "
            "(equivalent to the nint function in Fortran.)",
        ),
    )

    CONFIG.declare(
        "tolrelmetric",
        ConfigValue(
            default=None,
            domain=ListOf(float),
            description="Relative tolerance for outputs.",
            doc="Relative convergence tolerance for the chosen fitness metric for "
            "the modeling of output variables. One per output variable, space "
            "separated. Incremental model building will stop if two consecutive "
            "iterations do not improve the chosen metric by at least this amount.",
        ),
    )
    CONFIG.declare(
        "tolabsmetric",
        ConfigValue(
            default=None,
            domain=ListOf(float),
            description="Absolute tolerance for outputs.",
            doc="Absolute convergence tolerance for the chosen fitness metric for "
            "the modeling of output variables. One per output variable, space "
            "separated. Incremental model building will stop if two consecutive "
            "iterations do not improve the chosen metric by at least this amount.",
        ),
    )
    CONFIG.declare(
        "tolmeanerror",
        ConfigValue(
            default=None,
            domain=ListOf(float),
            description="Convergence tolerance for mean errors in outputs.",
            doc="Row vector of convergence tolerances for mean errors in the "
            "modeling of output variables. One per output variable, space "
            "separated. Incremental model building will stop if tolmeanerror, "
            "tolrelmetric, or tolabsmetric is satisfied.",
        ),
    )
    CONFIG.declare(
        "tolsse",
        ConfigValue(
            default=None,
            domain=float,
            description="Absolute tolerance on SSE",
            doc="Absolute tolerance on sum of square errors (SSE). ALAMO will "
            "terminate if it finds a solution whose SSE is within tolsse from "
            "the SSE of the full least squares problem.",
        ),
    )

    CONFIG.declare(
        "mipoptca",
        ConfigValue(
            default=None, domain=float, description="Absolute tolerance for MIP."
        ),
    )
    CONFIG.declare(
        "mipoptcr",
        ConfigValue(
            default=None, domain=float, description="Relative tolerance for MIP."
        ),
    )
    CONFIG.declare(
        "linearerror",
        ConfigValue(
            default=None,
            domain=Bool,
            description="If True, a linear objective is used when solving "
            "mixed-integer optimization problems; otherwise, a squared error will "
            "be employed.",
        ),
    )
    CONFIG.declare(
        "GAMS",
        ConfigValue(
            default=None,
            domain=str,
            description="Complete path of GAMS executable (or name if GAMS is in "
            "the user path).",
        ),
    )
    CONFIG.declare(
        "GAMSSOLVER",
        ConfigValue(
            default=None,
            domain=str,
            description="Name of preferred GAMS solver for solving ALAMO’s "
            "mixed-integer quadratic subproblems. Special facilities have been "
            "implemented in ALAMO and BARON that make BARON the preferred "
            "selection for this option. However, any mixed-integer quadratic "
            "programming solver available under GAMS can be used.",
        ),
    )
    CONFIG.declare(
        "solvemip",
        ConfigValue(
            default=None,
            domain=Bool,
            description="Whether to use an optimizer to solve MIP.",
        ),
    )
    CONFIG.declare(
        "print_to_screen",
        ConfigValue(
            default=None,
            domain=Bool,
            description="Send ALAMO output to stdout. Output is returned by the "
            "call_alamo method.",
        ),
    )

    # I/O file options
    CONFIG.declare(
        "alamo_path",
        ConfigValue(
            default=None, domain=Path, doc="Path to ALAMO executable (if not in path)."
        ),
    )
    CONFIG.declare(
        "filename",
        ConfigValue(
            default=None,
            domain=str,
            description="File name to use for ALAMO files - must be full path of a"
            " .alm file. Other files will be defined from this pattern. If this "
            "option is not None, then working files will not be deleted.",
        ),
    )
    CONFIG.declare(
        "working_directory",
        ConfigValue(
            default=None,
            domain=str,
            description="Full path to working directory for ALAMO to use. "
            "If this option is not None, then working files will not be deleted.",
        ),
    )
    CONFIG.declare(
        "overwrite_files",
        ConfigValue(
            default=False,
            domain=Bool,
            description="Flag indicating whether existing files can be " "overwritten.",
        ),
    )

    # TODO: We need to do some processing of the labels since ALAMO is
    # restrictive about the labels
    # TODO: We need to think more carefully about "input_bounds".
    # Alamo uses bounds during training, but we also want to consider
    # "valid" bounds for the surrogate
    def __init__(self, **settings):
        super().__init__(**settings)

        self._temp_context = None
        self._almfile = None
        self._trcfile = None

        if self.config.alamo_path is not None:
            alamo.executable = self.config.alamo_path

        self._results = None

    def train_surrogate(self):
        """
        General workflow method for training an ALAMO surrogate.

        Takes the existing data set and executes the ALAMO workflow to create
        an AlamoSurrogate based on the current configuration arguments.

        Args:
            None

        Returns:
            tuple : (success, AlamoSurrogate, message) where success indicates
            whether ALAMO was usccessfully executed, an instance of an
            AlamoSurrogate representing the trained surrogate, and message is
            the final status line from the ALAMO output log.
        """
        # Get paths for temp files
        self._get_files()

        return_code = None
        alamo_log = None
        alamo_object = None

        try:
            # Write .alm file
            self._write_alm_file()

            # Call ALAMO executable
            return_code, alamo_log = self._call_alamo()

            # Read back results
            trace_dict = self._read_trace_file(self._trcfile)

            # Populate results and SurrogateModel object
            self._populate_results(trace_dict)
            alamo_object = self._build_surrogate_object()

        finally:
            # Clean up temporary files if required
            self._remove_temp_files()

        success = False
        if return_code == 0:
            # Non-zero return code implies an error
            # specifics returned in the msg
            success = True
        alamo_msg = alamo_log.split("\n")[-3]

        return success, alamo_object, alamo_msg

    # TODO: let's generalize this under the metrics?
    def get_alamo_results(self):
        return self._results

    def _get_files(self):
        """
        Method to get/set paths for .alm and .trc files based on filename
        configuration argument.

        If filename is None, temporary files will be created.

        Args:
            None

        Returns:
            None
        """
        if self._temp_context is None:
            self._temp_context = TempfileManager.new_context()

        if self.config.filename is None:
            # Get a temporary file from the manager
            almfile = self._temp_context.create_tempfile(suffix=".alm")
        else:
            almfile = self.config.filename

            if not self.config.overwrite_files:
                # It is OK if the trace file exists, as ALAMO will append to it
                # The trace file reader handles this case
                if os.path.isfile(almfile):
                    raise FileExistsError(
                        f"A file with the name {almfile} already exists. "
                        f"Either choose a new file name or set "
                        f"overwrite_files = True."
                    )

            self._temp_context.add_tempfile(almfile, exists=False)

        trcfile = os.path.splitext(almfile)[0] + ".trc"
        self._temp_context.add_tempfile(trcfile, exists=False)

        if self.config.working_directory is None:
            wrkdir = self._temp_context.create_tempdir()
        else:
            wrkdir = self.config.working_directory

        # Set attributes to track file names
        self._almfile = almfile
        self._trcfile = trcfile
        self._wrkdir = wrkdir

    def _write_alm_to_stream(
        self, stream, trace_fname=None, training_data=None, validation_data=None
    ):
        """
        Method to write an ALAMO input file (.alm) to a stream.
        Users may provide specific data sets for training and validation.
        If no data sets are provided, the data sets contained in the
        AlamoModelTrainer are used.

        Args:
            stream: stream that data should be writen to
            trace_fname: name for trace file (.trc) to be included in .alm file
            training_data: Pandas dataframe to use for training surrogate
            validation_data: Pandas dataframe to use for validating surrogate

        Returns:
            None
        """
        if training_data is None:
            training_data = self._training_dataframe
        if validation_data is None:
            validation_data = self._validation_dataframe

        # Check bounds on inputs to avoid potential ALAMO failures
        input_max = list()
        input_min = list()
        for k, b in self._input_bounds.items():
            if b is None or b[0] is None or b[1] is None:
                raise ConfigurationError(
                    f"ALAMO configuration error: invalid bounds on input {k} " f"({b})."
                )
            elif b[0] == b[1]:
                raise ConfigurationError(
                    f"ALAMO configuration error: upper and lower bounds on "
                    f"input {k} are equal."
                )
            elif b[1] < b[0]:
                raise ConfigurationError(
                    f"ALAMO configuration error: upper bound is less than "
                    f"lower bound for input {k}."
                )
            input_min.append(b[0])
            input_max.append(b[1])

        # Get number of data points to build alm file
        n_rdata, n_inputs = training_data.shape

        if validation_data is not None:
            n_vdata, n_inputs = validation_data.shape
        else:
            n_vdata = 0

        stream.write("# IDAES Alamopy input file\n")
        stream.write(f"NINPUTS {len(self._input_labels)}\n")
        stream.write(f"NOUTPUTS {len(self._output_labels)}\n")
        stream.write(f"XLABELS {' '.join(map(str, self._input_labels))}\n")
        stream.write(f"ZLABELS {' '.join(map(str, self._output_labels))}\n")
        stream.write(f"XMIN {' '.join(map(str, input_min))}\n")
        stream.write(f"XMAX {' '.join(map(str, input_max))}\n")
        stream.write(f"NDATA {n_rdata}\n")
        if validation_data is not None:
            stream.write(f"NVALDATA {n_vdata}\n")
        stream.write("\n")

        # Other options for config
        # Can be bool, list of floats, None, Enum, float
        # Special cases
        if self.config.monomialpower is not None:
            stream.write(f"MONO {len(self.config.monomialpower)}\n")
        if self.config.multi2power is not None:
            stream.write(f"MULTI2 {len(self.config.multi2power)}\n")
        if self.config.multi3power is not None:
            stream.write(f"MULTI3 {len(self.config.multi3power)}\n")
        if self.config.ratiopower is not None:
            stream.write(f"RATIOS {len(self.config.ratiopower)}\n")
        if self.config.custom_basis_functions is not None:
            stream.write(f"NCUSTOMBAS {len(self.config.custom_basis_functions)}\n")

        for o in supported_options:
            if self.config[o] is None:
                # Write nothing to alm file
                continue
            elif isinstance(self.config[o], Enum):
                # Need to write value of Enum
                stream.write(f"{o} {self.config[o].value}\n")
            elif isinstance(self.config[o], bool):
                # Cast bool to int
                stream.write(f"{o} {int(self.config[o])}\n")
            elif isinstance(self.config[o], (str, float, int)):
                # Write value to file
                stream.write(f"{o} {self.config[o]}\n")
            else:
                # Assume the argument is a list
                stream.write(f"{o} {' '.join(map(str, self.config[o]))}\n")

        stream.write("\nTRACE 1\n")
        if trace_fname is not None:
            stream.write(f"TRACEFNAME {trace_fname}\n")

        def _df_to_data_fragment(df, **kwargs):
            text = df.to_string(
                header=False,
                index=False,
                float_format=lambda x: str(x).format(":g"),
                **kwargs,
            )
            return text

        stream.write("\nBEGIN_DATA\n")
        # Columns will be writen in order in input and output lists
        training_data_str = _df_to_data_fragment(
            training_data,
            columns=self._input_labels + self._output_labels,
        )
        stream.write(training_data_str)
        stream.write("\nEND_DATA\n")

        if validation_data is not None:
            # Add validation data defintion
            stream.write("\nBEGIN_VALDATA\n")
            val_data_str = _df_to_data_fragment(
                validation_data,
                columns=self._input_labels + self._output_labels,
            )
            stream.write(val_data_str)
            stream.write("\nEND_VALDATA\n")

        if self.config.custom_basis_functions is not None:
            stream.write("\nBEGIN_CUSTOMBAS\n")
            for i in self.config.custom_basis_functions:
                stream.write(f"{str(i)}\n")
            stream.write("END_CUSTOMBAS\n")

    def _write_alm_file(self, training_data=None, validation_data=None):
        """
        Method to write an ALAMO input file (.alm) using the current settings.
        Users may provide specific data sets for training and validation.
        If no data sets are provided, the data sets contained in the
        AlamoModelTrainer are used.

        Args:
            training_data: Pandas dataframe to use for training surrogate
            validation_data: Pandas dataframe to use for validating surrogate

        Returns:
            None
        """
        f = open(self._almfile, "w")
        self._write_alm_to_stream(f, self._trcfile, training_data, validation_data)
        f.close()

    def _call_alamo(self):
        """
        Method to call ALAMO executable from Python, passing the current .alm
        file as an argument.

        Args:
            None

        Returns:
            ALAMO: return code
            log: string of the text output from ALAMO
        """
        if alamo.executable is None:
            raise FileNotFoundError(
                "Could not find ALAMO executable. Please ensure that ALAMO "
                "is installed and in the system path, or provide a path to "
                "the executable."
            )

        ostreams = [StringIO(), sys.stdout]

        if self._temp_context is None:
            self._get_files()

        # Set working directory
        cwd = os.getcwd()
        os.chdir(self._wrkdir)

        try:
            # Add lst file to temp file manager
            lstfname = os.path.splitext(os.path.basename(self._almfile))[0] + ".lst"
            lstpath = os.path.join(cwd, lstfname)
            self._temp_context.add_tempfile(lstpath, exists=False)

            with TeeStream(*ostreams) as t:
                results = subprocess.run(
                    [alamo.executable, str(self._almfile)],
                    stdout=t.STDOUT,
                    stderr=t.STDERR,
                    universal_newlines=True,
                )

                t.STDOUT.flush()
                t.STDERR.flush()

            return_code = results.returncode
            alamo_log = ostreams[0].getvalue()

        except OSError:
            _log.error(
                f"Could not execute the command: alamo {str(self._almfile)}. ",
                f"Error message: {sys.exc_info()[1]}.",
            )
            raise

        finally:
            # Revert cwd to where it started
            os.chdir(cwd)

        if "ALAMO terminated with termination code " in alamo_log:
            self._remove_temp_files()
            _log.warn(
                "ALAMO executable returned non-zero return code. Check "
                "the ALAMO output for more information."
            )

        return return_code, alamo_log

    def _read_trace_file(self, trcfile, has_validation_data=False):
        """
        Method to read the results of an ALAMO run from a trace (.trc) file.
        The name location of the trace file is tored on the AlamoModelTrainer
        object and is generally set automatically.

        Args:
            trcfile : str
               Path to the trcfile to read
            output_labels : list of str
               List of strings of the output_labels (in order)
            has_validation_data : bool
                Bool indicating whether valdiation data was included in ALAMO run

        Returns:
            trace_dict: contents of trace file as a dict
        """
        with open(trcfile, "r") as f:
            lines = f.readlines()
        f.close()

        output_labels = self.output_labels()

        trace_read = {}
        # Get headers from first line in trace file
        headers = lines[0].split(", ")
        for i in range(len(headers)):
            header = headers[i].strip("#\n")
            if header in common_trace:
                trace_read[header] = None
            else:
                trace_read[header] = {}

        # Get trace output from final line(s) of file
        # ALAMO will append new lines to existing trace files
        # For multiple outputs, each output has its own line in trace file
        if self._validation_dataframe is not None or has_validation_data:
            omult = 2
        else:
            omult = 1

        for j in range(len(output_labels)):
            trace = lines[(-len(output_labels) + j) * omult].split(", ")

            for i in range(len(headers)):
                header = headers[i].strip("#\n")
                trace_val = trace[i].strip("\n")

                # Replace Fortran powers (^) with Python powers (**)
                trace_val = trace_val.replace("^", "**")
                # Replace = with ==
                trace_val = trace_val.replace("=", "==")

                if header in common_trace:
                    # These should be common across all output
                    if trace_read[header] is None:
                        # No value yet, so set value
                        trace_read[header] = trace_val
                    else:
                        # Check that current value matches the existng value
                        if trace_read[header] != trace_val:
                            raise RuntimeError(
                                f"Mismatch in values when reading ALAMO trace "
                                f"file. Values for {header}: "
                                f"{trace_read[header]}, {header}"
                            )
                else:
                    trace_read[header][output_labels[j]] = trace_val

                # Do some final sanity checks
                if header == "OUTPUT":
                    # OUTPUT should be equal to the current index of outputs
                    if trace_val != str(j + 1):
                        raise RuntimeError(
                            f"Mismatch when reading ALAMO trace file. "
                            f"Expected OUTPUT = {j+1}, found {trace_val}."
                        )
                elif header == "SET":
                    # SET should always be 0 - higher numbers are for
                    # validation data sets
                    if trace_val != "0":
                        raise RuntimeError(
                            f"Mismatch when reading ALAMO trace file. "
                            f"Expected SET = 0, found {trace_val}. "
                            f"This likely indicates the presence of "
                            f"unexpected validation data sets."
                        )
                elif header == "Model":
                    # Var label on LHS should match output label
                    vlabel = trace_val.split("==")[0].strip()
                    if vlabel != output_labels[j]:
                        raise RuntimeError(
                            f"Mismatch when reading ALAMO trace file. "
                            f"Label of output variable in expression "
                            f"({vlabel}) does not match expected label "
                            f"({output_labels[j]})."
                        )

        return trace_read

    def _populate_results(self, trace_dict):
        """
        Method to populate the results object with data from a trace file.

        Args:
            trace_dict: trace file data in form of a dict

        Returns:
            None
        """
        self._results = trace_dict

    def _build_surrogate_object(self):
        """
        Method to construct an AlmaoObject from the current results
        object.

        Args:
            None

        Returns:
            AlamoSurrogate
        """
        return AlamoSurrogate(
            surrogate_expressions=self._results["Model"],
            input_labels=self._input_labels,
            output_labels=self._output_labels,
            input_bounds=self._input_bounds,
        )

    def _remove_temp_files(self):
        """
        Method to remove temporary files created during the ALAMO workflow,
        i.e. the .alm and .trc files.

        Args:
            None

        Returns:
            None
        """
        remove = True
        if self.config.filename is not None:
            remove = False

        self._temp_context.release(remove=remove)
        # Release tempfile context
        self._temp_context = None

class AlamoSurrogate(SurrogateBase):
    """
    Standard SurrogateObject for surrogates trained using ALAMO.

    Contains methods to both populate a Pyomo Block with constraints
    representing the surrogate and to evaluate the surrogate a set of user
    provided points.
    """

    def __init__(
        self, surrogate_expressions, input_labels, output_labels, input_bounds=None
    ):
        super().__init__(input_labels, output_labels, input_bounds)
        self._surrogate_expressions = surrogate_expressions
        self._fcn = None

    def evaluate_surrogate(self, inputs):
        """
        Method to method to evaluate the ALAMO surrogate model at a set of user
        provided values.

        Args:
           dataframe: pandas DataFrame
              The dataframe of input values to be used in the evaluation. The dataframe
              needs to contain a column corresponding to each of the input labels. Additional
              columns are fine, but are not used.

        Returns:
            output: pandas Dataframe
              Returns a dataframe of the the output values evaluated at the provided inputs.
              The index of the output dataframe should match the index of the provided inputs.
        """
        # Create a set of lambda functions for evaluating the surrogate.
        if self._fcn is None:
            fcn = dict()
            for o in self._output_labels:
                fcn[o] = eval(
                    f"lambda {', '.join(self._input_labels)}: "
                    f"{self._surrogate_expressions[o].split('==')[1]}",
                    GLOBAL_FUNCS,
                )
            self._fcn = fcn

        # Use numpy to do the calculations as it is faster
        inputdata = inputs[self._input_labels].to_numpy()
        outputs = np.zeros(shape=(inputs.shape[0], len(self._output_labels)))

        for i in range(inputdata.shape[0]):
            for o in range(len(self._output_labels)):
                o_name = self._output_labels[o]
                outputs[i, o] = value(self._fcn[o_name](*inputdata[i, :]))

        return pd.DataFrame(
            data=outputs, index=inputs.index, columns=self._output_labels
        )

    def populate_block(self, block, additional_options=None):
        """
        Method to populate a Pyomo Block with surrogate model constraints.

        Args:
            block: Pyomo Block component to be populated with constraints.
            additional_options: None
               No additional options are required for this surrogate object
        Returns:
            None
        """

        # TODO: do we need to add the index_set stuff back in?
        output_set = Set(initialize=self._output_labels, ordered=True)

        def alamo_rule(b, o):
            lvars = block.input_vars_as_dict()
            lvars.update(block.output_vars_as_dict())
            return eval(self._surrogate_expressions[o], GLOBAL_FUNCS, lvars)

        block.alamo_constraint = Constraint(output_set, rule=alamo_rule)

    def save(self, strm):
        """
        Save an instance of this surrogate to the strm so the model can be used later.

        Args:
           strm: IO.TextIO
              This is the python stream like a file object or StringIO that will be used
              to serialize the surrogate object. This method writes a string
              of json data to the stream.
        """
        json.dump(
            {
                "surrogate": self._surrogate_expressions,
                "input_labels": self._input_labels,
                "output_labels": self._output_labels,
                "input_bounds": self._input_bounds,
            },
            strm,
        )

    @classmethod
    def load(cls, strm):
        """
        Create an instance of a surrogate from a stream.

        Args:
           strm: stream
              This is the python stream containing the data required to load the surrogate.
              This is often, but does not need to be a string of json data.

        Returns: an instance of the derived class or None if it failed to load
        """
        d = json.load(strm)

        surrogate_expressions = d["surrogate"]
        input_labels = d["input_labels"]
        output_labels = d["output_labels"]

        # Need to convert list of bounds to tuples
        input_bounds = {}
        for k, v in d["input_bounds"].items():
            input_bounds[k] = tuple(v)

        return AlamoSurrogate(
            surrogate_expressions=surrogate_expressions,
            input_labels=input_labels,
            output_labels=output_labels,
            input_bounds=input_bounds,
        )


#!/usr/bin/env python
"""

Test module for 2D Quadrilateral Meshes

"""
from __future__ import division
from builtins import range
from builtins import object
from past.utils import old_div
from proteus.iproteus import *
from proteus.test_utils import TestTools
TestTools.addSubFolders( inspect.currentframe() )

import stokes_2d_p
import stokes_2d_n
import pytest
import inspect
import numpy as np

@pytest.fixture(scope="module")
def simple_mesh():
    nnx = 4
    nny = 4
    x = (-1.0,-1.0)
    L = (2.0,2.0)
    refinementLevels = 2
    nLayersOfOverlap = 1
    parallelPartitioningType = 0
    skipInit = False
    mlMesh = MeshTools.MultilevelQuadrilateralMesh(nnx,nny,1,
                                                   x[0],x[1],0.0,
                                                   L[0],L[1],1.0,
                                                   refinementLevels,
                                                   skipInit,
                                                   nLayersOfOverlap,
                                                   parallelPartitioningType,
                                                   useC=False)
    yield mlMesh, nnx, nny

@pytest.fixture(scope="module")
def simple_mesh_with_c():
    nnx = 4
    nny = 4
    x = (-1.0,-1.0)
    L = (2.0,2.0)
    refinementLevels = 1
    nLayersOfOverlap = 1
    parallelPartitioningType = 0
    skipInit = False
    mlMesh = MeshTools.MultilevelQuadrilateralMesh(nnx,nny,1,
                                                   x[0],x[1],0.0,
                                                   L[0],L[1],1,
                                                   refinementLevels,
                                                   skipInit,
                                                   nLayersOfOverlap,
                                                   parallelPartitioningType,
                                                   useC=True)
    yield mlMesh, nnx, nny

@pytest.mark.MeshTools
def test_mesh_build(simple_mesh):
    """  Test mesh generation and refinment """
    mlMesh,nnx,nny = simple_mesh
    assert mlMesh.meshList[0].nElements_global == (nnx-1)*(nny-1), 'Mesh generator has built incorrect number of quads'
    assert mlMesh.meshList[1].nElements_global == 4*(nnx-1)*(nny-1), 'Mesh generator has built incorrect number of quads'

@pytest.mark.MeshTools
def test_mesh_build_1(simple_mesh_with_c):
    """ Test mesh generation with c """
    mlMesh,nnx,nny = simple_mesh_with_c
    assert mlMesh.meshList[0].nElements_global == (nnx-1)*(nny-1), 'Mesh generator has built incorrect number of quads'
    # TODO (ARB) - add an additional test for refinement when its ready

@pytest.mark.MeshTools
def test_calc_quad_area(simple_mesh):
    mlMesh, nnx, nny = simple_mesh
    for i in range(9):
        assert mlMesh.meshList[0]._calc_quad_area(i) == 4. / 9.

@pytest.mark.MeshTools
def test_calc_quad_area(simple_mesh_with_c):
    mlMesh, nnx, nny = simple_mesh_with_c
    for i in range(9):
        assert mlMesh.meshList[0]._calc_quad_area(i) == 4. / 9.

@pytest.mark.MeshTools
def test_calc_hmax(simple_mesh):
    mlMesh, nnx, nny = simple_mesh
    quad_mesh = mlMesh.meshList[0]
    for i in range(9):
        hmax_i = quad_mesh._calc_hmax(i)
        h = quad_mesh._calc_pt_distance((-1.0,old_div(-1.,3.)),(old_div(-1.,3.),old_div(1.,3.)))
        assert abs(h-hmax_i) < 1e-12

@pytest.mark.MeshTools
def test_buildNodeDiameterArray_1(simple_mesh):
    mlMesh, nnx, nny = simple_mesh
    quad_mesh = mlMesh.meshList[0]
    quad_mesh.buildNodeDiameterArray()
    expected = np.full((16),0.94280904)
    assert np.allclose(quad_mesh.nodeDiametersArray,expected)
    assert abs(quad_mesh.volume-4.0) < 1e-12

@pytest.mark.modelTest
@pytest.mark.moderateTest
@pytest.mark.MeshTools
class Test2DStokesOnQuads(object):
    """ Runs a 2D Poiseulle Stokes problem on Quads with TH elements """

    @classmethod
    def setup_class(cls):
        pass

    @classmethod
    def teardown_class(cls):
        pass

    def setup_method(self,method):
        from importlib import reload
        reload(stokes_2d_p)
        reload(stokes_2d_n)
        pList = [stokes_2d_p]
        nList = [stokes_2d_n]
        so = default_so
        so.tnList = [0.,1.]
        so.name = pList[0].name
        so.sList=[default_s]
        opts.verbose=True
        opts.profile=True
        opts.gatherArchive=True
        self._scriptdir = os.path.dirname(__file__)
        self.ns = NumericalSolution.NS_base(so,pList,nList,so.sList,opts)

    def teardown_method(self,method):
        Filelist = ["rdomain.edge",
                    "rdomain.ele",
                    "rdomain.neig",
                    "rdomain.node",
                    "rdomain.poly",
                    "reference_triangle.poly",
                    "reference_simplex.poly",
                    "proteus.log",
#                    "poiseulleFlow.xmf",
                    "poiseulleFlow.h5",
                    "poiseulleFlow0.h5"]
        TestTools.removeFiles(Filelist)

    # ARB TODO (10/24/18) something has become mixed up with this test and
    # needs to be fixed.  Notably, the *.h5 file is not working correctly
    # in paraview, making it difficult to correct this test. I simply do
    # not have the time to look into this at the moment and it is not a
    # critical piece of code for current projects.  See comments below.

    @pytest.mark.skip
    def test_01_FullRun(self):
        import filecmp
        self.ns.calculateSolution('test1')
        if self.ns.ar[0].global_sync:
            relpath = "comparison_files/poiseulle_global_xmf.output"
        else:
            relpath = "comparison_files/poiseulle_xmf.output"
        # The produced output has diverged from the old comparison
        # output. It needs to be confirmed that the new ouput is
        # in fact correct and drivenCavityNSE_LSC_expected.h5 should
        # be updated accordingly.
        xmf_file = filecmp.cmp('poiseulleFlow.xmf',os.path.join(self._scriptdir,relpath))
        import pdb ; pdb.set_trace()
        assert xmf_file == True, '******** xmf_file compare failed **********'

if __name__ == '__main__':
    pass


import pygrank as pg
import pytest
from .test_core import supported_backends

def test_split():
    data = {"community1": ["A", "B", "C", "D"], "community2": ["B", "E", "F", "G", "H", "I"]}
    training, test = pg.split(data, 1)
    assert training == test
    training, test = pg.split(data, 0.5)
    assert len(training["community2"]) == 3
    assert len(training["community1"]) == 2
    assert len(test["community2"]) == 3
    assert len(set(training["community1"])-set(test["community1"])) == len(training["community1"])
    assert len(set(training["community2"])-set(test["community2"])) == len(training["community2"])
    training, test = pg.split(data, 2)
    assert len(training["community2"]) == 2
    assert len(test["community1"]) == 2
    training, test = pg.split(data["community1"], 0.75)
    assert len(training) == 3
    assert len(test) == 1
    training, test = pg.split(set(data["community1"]), 0.75)
    assert len(training) == 3
    assert len(test) == 1

def test_auc_ndcg_compliance():
    _, graph, group = next(pg.load_datasets_one_community(["bigraph"]))
    training, test = pg.split(group, 0.5)
    for _ in supported_backends():
        scores1 = pg.PageRank()(graph, training)
        scores2 = pg.HeatKernel()(graph, training)
        AUC1 = pg.AUC(test, exclude=training)(scores1)
        AUC2 = pg.AUC(test, exclude=training)(scores2)
        NDCG1 = float(pg.NDCG(test, exclude=training)(scores1))
        NDCG2 = float(pg.NDCG(test, exclude=training)(scores2))
        assert (AUC1 < AUC2) == (NDCG1 < NDCG2)
        with pytest.raises(Exception):
            pg.AUC(test, exclude=test, k=len(graph)+1)(scores2)
        with pytest.raises(Exception):
            pg.NDCG(test, exclude=training, k=len(graph)+1)(scores2)

def test_edge_cases():
    assert pg.pRule([0])([0]) == 0
    assert pg.Cos([0])([0]) == 0
    with pytest.raises(Exception):
        pg.Measure()([0, 1, 0])
    with pytest.raises(Exception):
        pg.AUC([0, 0, 0])([0, 1, 0])
    with pytest.raises(Exception):
        pg.AUC([1, 1, 1])([0, 1, 0])
    with pytest.raises(Exception):
        pg.KLDivergence([0], exclude={"A": 1})([1])
    with pytest.raises(Exception):
        pg.Conductance(next(pg.load_datasets_graph(["graph5"])), max_rank=0.5)([1, 1, 1, 1, 1])
    import networkx as nx
    for _ in supported_backends():
        assert pg.Density(nx.Graph())([]) == 0
        assert pg.Modularity(nx.Graph())([]) == 0
        assert pg.KLDivergence([0,1,0])([0,1,0]) == 0
        assert pg.MKLDivergence([0,1,0])([0,1,0]) == 0
        assert pg.KLDivergence([0])([-1]) == 0

def test_strange_input_types():
    _, graph, group = next(pg.load_datasets_one_community(["bigraph"]))
    training, test = pg.split(group)
    for _ in supported_backends():
        scores = pg.PageRank()(graph, {v: 1 for v in training})
        ndcg = pg.NDCG(pg.to_signal(scores, {v: 1 for v in test}), k=3)({v: scores[v] for v in scores})
        ndcg_biased = pg.NDCG(pg.to_signal(scores, {v: 1 for v in test}), k=3)({v: scores[v] for v in test})
        assert ndcg < ndcg_biased

def test_correlation_compliance():
    graph = next(pg.load_datasets_graph(["graph5"]))
    # TODO: Make spearman and pearson correlation support tensorflow
    alg1 = pg.PageRank(alpha=0.5)
    alg2 = pg.PageRank(alpha=0.99)
    pearson_ordinals = pg.PearsonCorrelation(pg.Ordinals(alg1)(graph))(pg.Ordinals(alg2)(graph))
    spearman = pg.SpearmanCorrelation(alg1(graph))(alg2(graph))
    assert pearson_ordinals == spearman

def test_computations():
    for _ in supported_backends():
        assert pg.Accuracy([1, 2, 3])([1, 2, 3]) == 1
        assert pg.Mabs([3, 1, 1])([2, 0, 2]) == 1
        assert pg.CrossEntropy([1, 1, 1])([1, 1, 1]) < 1.E-12
        assert float(pg.Cos([2, 0, 1])([2, 0, 1])) == 1
        assert float(pg.Cos([2, 0, 1])([-2, 0, -1])) == -1
        assert float(pg.Dot([1, 1, 1])([1, 1, 1])) == 3

def test_aggregated():
    y1 = [1, 1, 0]
    y2 = [1, 0, 0]
    y3 = [1, 1, 0]
    for _ in supported_backends():
        # TODO: investiage why not exactly the same always (numerical precision should be lower)
        assert abs(float(pg.GM().add(pg.AUC(y1), max_val=0.5).add(pg.AUC(y2), min_val=0.9).evaluate(y3)) - 0.45**0.5) < 1.E-6
        assert abs(float(pg.AM().add(pg.AUC(y1), max_val=0.5).add(pg.AUC(y2), min_val=0.9).evaluate(y3)) - 0.7) < 1.E-6

def test_remove_edges():
    graph = next(pg.load_datasets_graph(["graph5"]))
    assert graph.has_edge("A", "B")
    assert graph.has_edge("C", "D")
    pg.remove_intra_edges(graph, {"community1": ["A", "B"], "community2": ["D", "C"]})
    assert graph.has_edge("B", "C")
    assert not graph.has_edge("A", "B")
    assert not graph.has_edge("C", "D")


#
#  File:
#    TRANS_read_ASCII_way2.py
#
#  Synopsis:
#    Illustrates how to read an ASCII file
#
#  Categories:
#    I/O
#
#  Author:
#    Karin Meier-Fleischer, based on NCL example
#
#  Date of initial publication:
#    September 2018
#
#  Description:
#    This example shows how to read an ASCII file.
#
#  Effects illustrated:
#    o  Read ASCII data
#
#  Output:
#    -
#
#  Notes: The data for this example can be downloaded from
#    http://www.ncl.ucar.edu/Document/Manuals/NCL_User_Guide/Data/
#
"""
  Transition Guide Python Example:   TRANS_read_ASCII_way2.py

   - read ASCII file
   - retrieve variable informations

	Test_6h.csv

	2.00;3.50;5.10;8.20
	2.40;3.10;4.80;8.90
	2.60;3.70;5.30;10.10
	2.75;3.90;5.55;10.25
	3.00;4.10;6.05;10.50

  2018-08-28  kmf
"""
from __future__ import print_function
import numpy as np

print("")

#--  data file name

diri = "/Users/k204045/local/miniconda2/envs/pyn_env/lib/ncarg/data/nug/"
fili = "Test_6h.csv"

#-- delimiter

delim = ';'

#-- read the data

f      = open(diri+fili,'r')
data  = f.readlines()

#-- assign list to append elements

vals = []

for i in data:
	line = i.strip()
	cols = line.split(delim)
	vals.append(cols[:])

#-- convert string to float

vals = np.array(vals).astype(float)

nlines = vals[:,0]
ncols  = vals[0,:]

print("vals: " + str(vals))

#-- rows by column

print("--> columns count:    " + str(len(ncols)))
print("--> lines count:      " + str(len(nlines)))
print("--> rank of vals:     " + str(len(vals.shape)))
print("--> shape vals:       " + str(vals.shape))

exit()


import numpy as np

from transonic.util import timeit
from transonic.config import backend_default

def check(functions, arr, columns):
    res0 = functions[0](arr, columns)
    for func in functions[1:]:
        assert np.allclose(res0, func(arr, columns))
    print("Checks passed: results are consistent")

def bench(functions, arr, columns):
    print(backend_default.capitalize())
    for func in functions:
        result = timeit("func(arr, columns)", globals=locals())
        print(f"{func.__name__:20s} {result:.3e} s")
    print()


# ---
# jupyter:
#   jupytext:
#     formats: ipynb,py
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.5.2
#   kernelspec:
#     display_name: Python 3.8.3 64-bit
#     language: python
#     name: python38364bit425a2724ae224223ab43e7c3d3663fd5
# ---

# General references:
# - [Learn Astropy](http://learn.astropy.org/)
#     - [FITS File Handling](https://docs.astropy.org/en/stable/io/fits/index.html)
#     - [Astropy Table](https://docs.astropy.org/en/stable/api/astropy.table.Table.html#astropy.table.Table)
# - [FITS Standard Page](https://fits.gsfc.nasa.gov/fits_standard.html)

# +
# Import Astropy Modules for file operation and converting to table
from astropy.io import fits
from astropy.table import Table

# Import custom module for pre-defined data path
from grb.config import path
# -

dir(fits)

dir(Table)

# One Way to Load the FITS Data
# =============================

# Open the file
# -------------

filename = path.FITS / "160625/L200408111538F357373F92_PH00.fits"
f = fits.open(filename)

# View File Information
# ---------------------

# + tags=[]
f.info()
# -

# Select Needed Data
# ------------------
# `PRIMARY` contains general information about the GRB.

header = f[0].header
header

# See [Time in *Fermi* Data Analysis](https://fermi.gsfc.nasa.gov/ssc/data/analysis/documentation/Cicerone/Cicerone_Data/Time_in_ScienceTools.html#:~:text=The%20Fermitools%20use%20mission%20elapsed%20time%20%28MET%29%2C%20the,the%20UTC%20system.%20Time%20Systems%20in%20a%20Nutshell) for help in `MET`.

METREF = header['MJDREFI'] + header['MJDREFF']
METREF

# `EVENTS` contains the (photons or spacecraft) events information

data = f[1].data
data

# Load Data to Astropy Table
# --------------------------

t = Table(data)
t

# Close The Opened File
# ---------------------
# Always remember to close the opened file.

f.close()

# An Alternative Way to Open A File
# =================================
# [`with` statement](https://docs.python.org/3/reference/compound_stmts.html#with) automatically handles open and close of files.

# + tags=[]
with fits.open(filename) as f:
    f.info()
    data2 = f[1].data
# -

# Check if the two ways get the same data

import numpy as np

np.all(data==data2)

# An Alternative Way to Load FITS data
# ====================================

t2 = Table.read(filename, hdu=1)

t2

# Check if the two ways get the same data.

np.all(t==t2)

# Convert Astropy Table to Pandas DataFrame
# =========================================
# Sometimes it would be better to deal with Pandas DataFrame instead of Astropy Table.
#
# **Note**: Some information, for example, units, will get lost, after the conversion.

t3 = t[['ENERGY', 'TIME']]
df = t3.to_pandas()
df

# + tags=[]
print(f"{type(df) = },\n{type(t3) = }")
# -

# Application to GBM TTE FITS Files
# =================================

# + tags=[]
ttefile = path.FITS / "160625/TTE/glg_tte_b0_bn160625945_v00.fit"
with fits.open(ttefile) as f:
    f.info()
    header = f[0].header
header
# -

# Trigger time relative to MJDREF (in second)

TRIGTIME = header['TRIGTIME']
TRIGTIME

# +
from grb.lat.timeutils import UTCMET

trigtime = UTCMET.met2utc(TRIGTIME)
trigtime
# -

trigtime.value

tobs = header['DATE-OBS']
tobs_met = UTCMET.utc2met(tobs)
tobs_met

tobs_met.value

dt = trigtime - tobs_met
dt

dt.value


"""
=========================================
Creating a colormap from a list of colors
=========================================

For more detail on creating and manipulating colormaps see
:doc:`/tutorials/colors/colormap-manipulation`.

Creating a :doc:`colormap </tutorials/colors/colormaps>`
from a list of colors can be done with the
:meth:`~.colors.LinearSegmentedColormap.from_list` method of
`.LinearSegmentedColormap`. You must pass a list of RGB tuples that define the
mixture of colors from 0 to 1.

Creating custom colormaps
-------------------------
It is also possible to create a custom mapping for a colormap. This is
accomplished by creating dictionary that specifies how the RGB channels
change from one end of the cmap to the other.

Example: suppose you want red to increase from 0 to 1 over the bottom
half, green to do the same over the middle half, and blue over the top
half.  Then you would use::

  cdict = {'red':   ((0.0,  0.0, 0.0),
                     (0.5,  1.0, 1.0),
                     (1.0,  1.0, 1.0)),

           'green': ((0.0,  0.0, 0.0),
                     (0.25, 0.0, 0.0),
                     (0.75, 1.0, 1.0),
                     (1.0,  1.0, 1.0)),

           'blue':  ((0.0,  0.0, 0.0),
                     (0.5,  0.0, 0.0),
                     (1.0,  1.0, 1.0))}

If, as in this example, there are no discontinuities in the r, g, and b
components, then it is quite simple: the second and third element of
each tuple, above, is the same--call it "y".  The first element ("x")
defines interpolation intervals over the full range of 0 to 1, and it
must span that whole range.  In other words, the values of x divide the
0-to-1 range into a set of segments, and y gives the end-point color
values for each segment.

Now consider the green. cdict['green'] is saying that for
0 <= x <= 0.25, y is zero; no green.
0.25 < x <= 0.75, y varies linearly from 0 to 1.
x > 0.75, y remains at 1, full green.

If there are discontinuities, then it is a little more complicated.
Label the 3 elements in each row in the cdict entry for a given color as
(x, y0, y1).  Then for values of x between x[i] and x[i+1] the color
value is interpolated between y1[i] and y0[i+1].

Going back to the cookbook example, look at cdict['red']; because y0 !=
y1, it is saying that for x from 0 to 0.5, red increases from 0 to 1,
but then it jumps down, so that for x from 0.5 to 1, red increases from
0.7 to 1.  Green ramps from 0 to 1 as x goes from 0 to 0.5, then jumps
back to 0, and ramps back to 1 as x goes from 0.5 to 1.::

  row i:   x  y0  y1
                  /
                 /
  row i+1: x  y0  y1

Above is an attempt to show that for x in the range x[i] to x[i+1], the
interpolation is between y1[i] and y0[i+1].  So, y0[0] and y1[-1] are
never used.

"""
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap

# Make some illustrative fake data:

x = np.arange(0, np.pi, 0.1)
y = np.arange(0, 2 * np.pi, 0.1)
X, Y = np.meshgrid(x, y)
Z = np.cos(X) * np.sin(Y) * 10

###############################################################################
# --- Colormaps from a list ---

colors = [(1, 0, 0), (0, 1, 0), (0, 0, 1)]  # R -> G -> B
n_bins = [3, 6, 10, 100]  # Discretizes the interpolation into bins
cmap_name = 'my_list'
fig, axs = plt.subplots(2, 2, figsize=(6, 9))
fig.subplots_adjust(left=0.02, bottom=0.06, right=0.95, top=0.94, wspace=0.05)
for n_bin, ax in zip(n_bins, axs.ravel()):
    # Create the colormap
    cm = LinearSegmentedColormap.from_list(
        cmap_name, colors, N=n_bin)
    # Fewer bins will result in "coarser" colomap interpolation
    im = ax.imshow(Z, interpolation='nearest', origin='lower', cmap=cm)
    ax.set_title("N bins: %s" % n_bin)
    fig.colorbar(im, ax=ax)

###############################################################################
# --- Custom colormaps ---

cdict1 = {'red':   ((0.0, 0.0, 0.0),
                    (0.5, 0.0, 0.1),
                    (1.0, 1.0, 1.0)),

          'green': ((0.0, 0.0, 0.0),
                    (1.0, 0.0, 0.0)),

          'blue':  ((0.0, 0.0, 1.0),
                    (0.5, 0.1, 0.0),
                    (1.0, 0.0, 0.0))
          }

cdict2 = {'red':   ((0.0, 0.0, 0.0),
                    (0.5, 0.0, 1.0),
                    (1.0, 0.1, 1.0)),

          'green': ((0.0, 0.0, 0.0),
                    (1.0, 0.0, 0.0)),

          'blue':  ((0.0, 0.0, 0.1),
                    (0.5, 1.0, 0.0),
                    (1.0, 0.0, 0.0))
          }

cdict3 = {'red':  ((0.0, 0.0, 0.0),
                   (0.25, 0.0, 0.0),
                   (0.5, 0.8, 1.0),
                   (0.75, 1.0, 1.0),
                   (1.0, 0.4, 1.0)),

          'green': ((0.0, 0.0, 0.0),
                    (0.25, 0.0, 0.0),
                    (0.5, 0.9, 0.9),
                    (0.75, 0.0, 0.0),
                    (1.0, 0.0, 0.0)),

          'blue':  ((0.0, 0.0, 0.4),
                    (0.25, 1.0, 1.0),
                    (0.5, 1.0, 0.8),
                    (0.75, 0.0, 0.0),
                    (1.0, 0.0, 0.0))
          }

# Make a modified version of cdict3 with some transparency
# in the middle of the range.
cdict4 = {**cdict3,
          'alpha': ((0.0, 1.0, 1.0),
                #   (0.25,1.0, 1.0),
                    (0.5, 0.3, 0.3),
                #   (0.75,1.0, 1.0),
                    (1.0, 1.0, 1.0)),
          }

###############################################################################
# Now we will use this example to illustrate 3 ways of
# handling custom colormaps.
# First, the most direct and explicit:

blue_red1 = LinearSegmentedColormap('BlueRed1', cdict1)

###############################################################################
# Second, create the map explicitly and register it.
# Like the first method, this method works with any kind
# of Colormap, not just
# a LinearSegmentedColormap:

blue_red2 = LinearSegmentedColormap('BlueRed2', cdict2)
plt.register_cmap(cmap=blue_red2)

###############################################################################
# Third, for LinearSegmentedColormap only,
# leave everything to register_cmap:

plt.register_cmap(name='BlueRed3', data=cdict3)  # optional lut kwarg
plt.register_cmap(name='BlueRedAlpha', data=cdict4)

###############################################################################
# Make the figure:

fig, axs = plt.subplots(2, 2, figsize=(6, 9))
fig.subplots_adjust(left=0.02, bottom=0.06, right=0.95, top=0.94, wspace=0.05)

# Make 4 subplots:

im1 = axs[0, 0].imshow(Z, interpolation='nearest', cmap=blue_red1)
fig.colorbar(im1, ax=axs[0, 0])

cmap = plt.get_cmap('BlueRed2')
im2 = axs[1, 0].imshow(Z, interpolation='nearest', cmap=cmap)
fig.colorbar(im2, ax=axs[1, 0])

# Now we will set the third cmap as the default.  One would
# not normally do this in the middle of a script like this;
# it is done here just to illustrate the method.

plt.rcParams['image.cmap'] = 'BlueRed3'

im3 = axs[0, 1].imshow(Z, interpolation='nearest')
fig.colorbar(im3, ax=axs[0, 1])
axs[0, 1].set_title("Alpha = 1")

# Or as yet another variation, we can replace the rcParams
# specification *before* the imshow with the following *after*
# imshow.
# This sets the new default *and* sets the colormap of the last
# image-like item plotted via pyplot, if any.
#

# Draw a line with low zorder so it will be behind the image.
axs[1, 1].plot([0, 10 * np.pi], [0, 20 * np.pi], color='c', lw=20, zorder=-1)

im4 = axs[1, 1].imshow(Z, interpolation='nearest')
fig.colorbar(im4, ax=axs[1, 1])

# Here it is: changing the colormap for the current image and its
# colorbar after they have been plotted.
im4.set_cmap('BlueRedAlpha')
axs[1, 1].set_title("Varying alpha")
#

fig.suptitle('Custom Blue-Red colormaps', fontsize=16)
fig.subplots_adjust(top=0.9)

plt.show()

#############################################################################
#
# ------------
#
# References
# """"""""""
#
# The use of the following functions, methods, classes and modules is shown
# in this example:

import matplotlib
matplotlib.axes.Axes.imshow
matplotlib.pyplot.imshow
matplotlib.figure.Figure.colorbar
matplotlib.pyplot.colorbar
matplotlib.colors
matplotlib.colors.LinearSegmentedColormap
matplotlib.colors.LinearSegmentedColormap.from_list
matplotlib.cm
matplotlib.cm.ScalarMappable.set_cmap
matplotlib.pyplot.register_cmap
matplotlib.cm.register_cmap


from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import time
import logging

import numpy as np
from six.moves import xrange  # pylint: disable=redefined-builtin
import tensorflow as tf
from tensorflow.python.ops import variable_scope as vs

from evaluate import exact_match_score, f1_score

logging.basicConfig(level=logging.INFO)

def get_optimizer(opt):
    if opt == "adam":
        optfn = tf.train.AdamOptimizer
    elif opt == "sgd":
        optfn = tf.train.GradientDescentOptimizer
    else:
        assert (False)
    return optfn

class Encoder(object):
    def __init__(self, size, vocab_dim):
        self.size = size
        self.vocab_dim = vocab_dim

    def encode(self, inputs, masks, encoder_state_input):
        """
        In a generalized encode function, you pass in your inputs,
        masks, and an initial
        hidden state input into this function.

        :param inputs: Symbolic representations of your input
        :param masks: this is to make sure tf.nn.dynamic_rnn doesn't iterate
                      through masked steps
        :param encoder_state_input: (Optional) pass this as initial hidden state
                                    to tf.nn.dynamic_rnn to build conditional representations
        :return: an encoded representation of your input.
                 It can be context-level representation, word-level representation,
                 or both.
        """

        return

class Decoder(object):
    def __init__(self, output_size):
        self.output_size = output_size

    def decode(self, knowledge_rep):
        """
        takes in a knowledge representation
        and output a probability estimation over
        all paragraph tokens on which token should be
        the start of the answer span, and which should be
        the end of the answer span.

        :param knowledge_rep: it is a representation of the paragraph and question,
                              decided by how you choose to implement the encoder
        :return:
        """

        return

class QASystem(object):
    def __init__(self, encoder, decoder, *args):
        """
        Initializes your System

        :param encoder: an encoder that you constructed in train.py
        :param decoder: a decoder that you constructed in train.py
        :param args: pass in more arguments as needed
        """

        # ==== set up placeholder tokens ========

        # ==== assemble pieces ====
        with tf.variable_scope("qa", initializer=tf.uniform_unit_scaling_initializer(1.0)):
            self.setup_embeddings()
            self.setup_system()
            self.setup_loss()

        # ==== set up training/updating procedure ====
        pass

    def setup_system(self):
        """
        After your modularized implementation of encoder and decoder
        you should call various functions inside encoder, decoder here
        to assemble your reading comprehension system!
        :return:
        """
        raise NotImplementedError("Connect all parts of your system here!")

    def setup_loss(self):
        """
        Set up your loss computation here
        :return:
        """
        with vs.variable_scope("loss"):
            pass

    def setup_embeddings(self):
        """
        Loads distributed word representations based on placeholder tokens
        :return:
        """
        with vs.variable_scope("embeddings"):
            pass

    def optimize(self, session, train_x, train_y):
        """
        Takes in actual data to optimize your model
        This method is equivalent to a step() function
        :return:
        """
        input_feed = {}

        # fill in this feed_dictionary like:
        # input_feed['train_x'] = train_x

        output_feed = []

        outputs = session.run(output_feed, input_feed)

        return outputs

    def test(self, session, valid_x, valid_y):
        """
        in here you should compute a cost for your validation set
        and tune your hyperparameters according to the validation set performance
        :return:
        """
        input_feed = {}

        # fill in this feed_dictionary like:
        # input_feed['valid_x'] = valid_x

        output_feed = []

        outputs = session.run(output_feed, input_feed)

        return outputs

    def decode(self, session, test_x):
        """
        Returns the probability distribution over different positions in the paragraph
        so that other methods like self.answer() will be able to work properly
        :return:
        """
        input_feed = {}

        # fill in this feed_dictionary like:
        # input_feed['test_x'] = test_x

        output_feed = []

        outputs = session.run(output_feed, input_feed)

        return outputs

    def answer(self, session, test_x):

        yp, yp2 = self.decode(session, test_x)

        a_s = np.argmax(yp, axis=1)
        a_e = np.argmax(yp2, axis=1)

        return (a_s, a_e)

    def validate(self, sess, valid_dataset):
        """
        Iterate through the validation dataset and determine what
        the validation cost is.

        This method calls self.test() which explicitly calculates validation cost.

        How you implement this function is dependent on how you design
        your data iteration function

        :return:
        """
        valid_cost = 0

        for valid_x, valid_y in valid_dataset:
          valid_cost = self.test(sess, valid_x, valid_y)

        return valid_cost

    def evaluate_answer(self, session, dataset, sample=100, log=False):
        """
        Evaluate the model's performance using the harmonic mean of F1 and Exact Match (EM)
        with the set of true answer labels

        This step actually takes quite some time. So we can only sample 100 examples
        from either training or testing set.

        :param session: session should always be centrally managed in train.py
        :param dataset: a representation of our data, in some implementations, you can
                        pass in multiple components (arguments) of one dataset to this function
        :param sample: how many examples in dataset we look at
        :param log: whether we print to std out stream
        :return:
        """

        f1 = 0.
        em = 0.

        if log:
            logging.info("F1: {}, EM: {}, for {} samples".format(f1, em, sample))

        return f1, em

    def train(self, session, dataset, train_dir):
        """
        Implement main training loop

        TIPS:
        You should also implement learning rate annealing (look into tf.train.exponential_decay)
        Considering the long time to train, you should save your model per epoch.

        More ambitious appoarch can include implement early stopping, or reload
        previous models if they have higher performance than the current one

        As suggested in the document, you should evaluate your training progress by
        printing out information every fixed number of iterations.

        We recommend you evaluate your model performance on F1 and EM instead of just
        looking at the cost.

        :param session: it should be passed in from train.py
        :param dataset: a representation of our data, in some implementations, you can
                        pass in multiple components (arguments) of one dataset to this function
        :param train_dir: path to the directory where you should save the model checkpoint
        :return:
        """

        # some free code to print out number of parameters in your model
        # it's always good to check!
        # you will also want to save your model parameters in train_dir
        # so that you can use your trained model to make predictions, or
        # even continue training

        tic = time.time()
        params = tf.trainable_variables()
        num_params = sum(map(lambda t: np.prod(tf.shape(t.value()).eval()), params))
        toc = time.time()
        logging.info("Number of params: %d (retreival took %f secs)" % (num_params, toc - tic))


#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""Provide a wrapper around quadratic programming solvers.

References:
    [1] https://github.com/stephane-caron/qpsolvers
"""

import numpy as np

# CVXOPT
# import cvxopt
# CVXPY: nice wrapper around cvxopt
# import cvxpy
# Quadprog
# import quadprog

# QPsolvers optimizers: unified Python interface for multiple QP solvers (cvxopt, cvxpy, quadprog,...)
try:
    import qpsolvers
except ImportError as e:
    raise ImportError(e.__str__() + "\n HINT: you can install qpsolvers directly via 'pip install qpsolvers'.")

from pyrobolearn.optimizers import Optimizer

__author__ = "Brian Delhaisse"
__copyright__ = "Copyright 2018, PyRoboLearn"
__credits__ = ["Brian Delhaisse"]
__license__ = "GNU GPLv3"
__version__ = "1.0.0"
__maintainer__ = "Brian Delhaisse"
__email__ = "briandelhaisse@gmail.com"
__status__ = "Development"

# class CVXOPT(Optimizer):
#     r"""Convex Optimizer
#
#     Note: cvxpy module is a nice wrapper around cvxopt that follows paradigm of a disciplined convex programming.
#
#     References:
#         [1] Python Software for Convex Optimization: https://cvxopt.org/
#         [2] Github repo: https://github.com/cvxopt/cvxopt
#     """
#     pass
#
#
# class CVXPY(Optimizer):
#     r"""Convex Optimizer
#
#     References:
#         [1] CVXPY: http://www.cvxpy.org/
#         [2] Github repo: https://github.com/cvxgrp/cvxpy
#     """
#     pass
#
#
# class QuadProg(object):
#     r"""Quadprog
#
#     References:
#         [1] Github repo: https://github.com/rmcgibbo/quadprog
#     """
#     pass

class QP(Optimizer):
    r"""Quadratic Programming solvers

    This class uses the `qpsolvers` which is a unified Python interface for multiple QP solvers [1,2].

    .. math::

        \min_{x \in R^n} \frac{1}{2} x^T P x + q^T x

    subject to

    .. math::

        Gx \leq h
        Ax = b

    where :math:`x` is the vector of optimization variables, the matrix :math:`P` and vector :math:`q` are used to
    define any quadratic objective function on these variables, while the matrix-vector couples :math:`(G,h)` and
    :math:`(A,b)` respectively define inequality and equality constraints. Vector inequalities apply coordinate by
    coordinate [1].

    - Dense solvers:
        - CVXOPT
        - CVXPY
        - qpOASES
        - quadprog
    - Sparse solvers:
        - ECOS as wrapped by CVXPY
        - Gurobi
        - MOSEK
        - OSQP

    Check the available solvers by calling `print(qpsolvers.available_solvers)`.

    Notes: Many solvers (including CVXOPT, OSQP and quadprog) assume that `P` is a symmetric matrix, and may return
    erroneous results when that is not the case. You can set ``sym_proj=True`` to project `P` on its symmetric part,
    at the cost of some computation time.

    References:
        [1] QP in Python: https://scaron.info/blog/quadratic-programming-in-python.html
        [2] Github repo: https://github.com/stephane-caron/qpsolvers
    """

    def __init__(self, method='quadprog', *args, **kwargs):
        """
        Initialize the QP solver.

        Args:
            method (str): QP method/library to use. Select between ['cvxopt', 'cvxpy', 'ecos', 'gurobi', 'mosek',
                'osqp', 'qpoases', 'quadprog']
        """
        super(QP, self).__init__(*args, **kwargs)

        solvers = set(qpsolvers.available_solvers)
        if len(solvers) == 0:
            raise ValueError("No QP solvers have been found on this computer. Please install one of the QP modules")
        if method not in solvers:
            method = 'quadprog'
        self.method = method

        # check methods that require a symmetric matrix for P
        methods = ['cvxopt', 'osqp', 'quadprog']
        self.sym_proj = True if self.method in set(methods) else False

    ##################
    # Static Methods #
    ##################

    @staticmethod
    def is_symmetric(X, tol=1e-8):
        """Check if the given matrix is symmetric."""
        return np.allclose(X, X.T, atol=tol)

    ###########
    # Methods #
    ###########

    def optimize(self, Q, p, x0=None, G=None, h=None, A=None, b=None):
        r"""
        Optimize the given quadratic problem.

        .. math::

            \min_{x \in \mathbb{R}^N} \frac{1}{2} x^T Q x + p^T x

        subject to

        .. math::

            Gx \leq h
            Ax = b

        Args:
            Q (np.array[N,N]): matrix used in the QP objective function where `N` is the size of the vector `x` being
                optimized.
            p (np.array[N]): vector used in the QP objective function where `N` is the size of the vector `x` being
                optimized.
            G (np.array[M,N]): matrix used in the inequality constraint, where `M` is the number of inequalities, and
                `N` is the size of the vector `x` being optimized. Note that if you have lower and upper bounds for
                the vector `x`, you can set :attr:`G` to be the concatenation of :math:`[-I, I]^\top`, where :math:`I`
                is the identity matrix.
            h (np.array[M]): vector used in the inequality constraint, where `M` is the number of inequalities. Note
                that if you have lower and upper bounds for the vector `x`, you can set :attr:`h` to be the
                concatenation of :math:`[-b_l^\top, b_u^\top]`, where :math:`b_l` and :math:`b_u` are the lower and
                upper bounds respectively.
            A (np.array[K,N]): matrix used in the equality constraint.
            b (np.array[K,N]): vector used in the equality constraint.

        Returns:
            np.array: QP solution
        """
        return qpsolvers.solve_qp(Q, p, G, h, A, b, solver=self.method, initvals=x0, sym_proj=self.sym_proj)


import matplotlib.pyplot as plt

def visualize(x_s, y_s, labels, title, x_label, y_label, savefig=False, out_dir="."):
    """ This function generates multiple curves within a plot.

    Args:
    - x_s (list of lists): A list of x's for each curve.
    - y_s (list of lists): A list of y's for each curve.
    - labels (str): The labels for every curve.
    - title (str): The title of the plot.
    """
    assert len(x_s) == len(y_s) == len(labels)
    plt.clf()

    for (x, y, label) in zip(x_s, y_s, labels):
        plt.plot(x, y, label=label)

    plt.title(title)
    plt.xlabel(x_label)
    plt.ylabel(y_label)
    plt.legend()

    if savefig:
        fig = plt.gcf()
        fig.savefig(fname=f"{out_dir}/{title}.png", dpi=150, format="png")
    else:
        plt.show()

if __name__ == "__main__":
    import numpy as np
    x_s = [np.arange(10), np.arange(10)]
    y_s = [np.arange(10), np.arange(10) + 1]
    labels = ["1", "2"]
    title = "test"
    x_label = "x"
    y_label = "y"

    visualize(x_s, y_s, labels, title, x_label, y_label)


# -*- coding: utf-8 -*-
"""
Created on Sun May 10 12:01:46 2020

@author: 66IN
"""

import numpy as np
import sys

'''
a = np.array([(1,2,3),(4,5,6)])

print((a))

s = range(1000)
print(sys.getsizeof(5)*len(s))

d = np.arange(1000)
print(d.size*d.itemsize)

'''


""" Doctests for Nipy / NumPy-specific nose/doctest modifications
"""
# try the #random directive on the output line
def check_random_directive():
    '''
    >>> 2+2
    <BadExample object at 0x084D05AC>  #random: may vary on your system
    '''

# check the implicit "import numpy as np"
def check_implicit_np():
    '''
    >>> np.array([1,2,3])
    array([1, 2, 3])
    '''

# there's some extraneous whitespace around the correct responses
def check_whitespace_enabled():
    '''
    # whitespace after the 3
    >>> 1+2
    3

    # whitespace before the 7
    >>> 3+4
     7
    '''

def check_empty_output():
    """ Check that no output does not cause an error.

    This is related to nose bug 445; the numpy plugin changed the
    doctest-result-variable default and therefore hit this bug:
    http://code.google.com/p/python-nose/issues/detail?id=445

    >>> a = 10
    """

def check_skip():
    """ Check skip directive

    The test below should not run

    >>> 1/0 #doctest: +SKIP
    """

def func():
    return 1

def check_have_module_context():
    """ Check that, unlike numpy, we do have the module namespace

    >>> func()
    1
    """

def check_fails():
    """ Check inversion directive

    The directive is mainly for tests

    >>> 'black' #doctest: +NOT_EQUAL
    'white'
    >>> 'white' #doctest: +NOT_EQUAL
    'black'
    """

def check_ignore_output():
    """ Check IGNORE_OUTPUT option works

    >>> 'The answer' #doctest: +IGNORE_OUTPUT
    42
    >>> 'The answer' #doctest: +IGNORE_OUTPUT
    'The answer'
    """

def check_sympy_equal():
    """ Check SYMPY_EQUAL option

    >>> from sympy import symbols
    >>> a, b, c = symbols('a, b, c')
    >>> a + b #doctest: +SYMPY_EQUAL
    b + a
    >>> a + b #doctest: +SYMPY_EQUAL
    a + b
    >>> a + b #doctest: +SYMPY_EQUAL +NOT_EQUAL
    a + c
    >>> a + b #doctest: +SYMPY_EQUAL +NOT_EQUAL
    a - b
    """

def check_fp_equal():
    """ Check floating point equal

    >>> 0.12345678 #doctest: +FP_6DP
    0.1234569
    >>> 0.12345678 #doctest: +FP_6DP +NOT_EQUAL
    0.1234564
    >>> 0.12345678 #doctest: +FP_4DP
    0.1235
    >>> 0.12345678 #doctest: +FP_6DP +NOT_EQUAL
    0.1235
    """

def check_array_repr():
    """ Stripping of array repr

    >>> arr = np.arange(5, dtype='i2')

    The test should match with and without the array repr

    >>> arr #doctest: +STRIP_ARRAY_REPR
    [0, 1, 2, 3, 4]
    >>> arr #doctest: +STRIP_ARRAY_REPR
    array([0, 1, 2, 3, 4], dtype=int16)
    """

def check_combinations():
    """ Check the processing combines as expected

    >>> 0.33333 #doctest: +SYMPY_EQUAL +NOT_EQUAL
    0.3333
    >>> 0.33333 #doctest: +SYMPY_EQUAL +FP_4DP
    0.3333
    >>> arr = np.arange(5, dtype='i2')

    This next will not sympify unless the array repr is removed

    >>> arr #doctest: +STRIP_ARRAY_REPR +SYMPY_EQUAL
    array([0, 1, 2, 3, 4], dtype=int16)
    """

if __name__ == '__main__':
    # Run tests outside nipy test rig
    import sys
    import nose
    from nipy.testing.doctester import NipyDoctest
    argv = [sys.argv[0], __file__, '--with-nipydoctest'] + sys.argv[1:]
    nose.core.TestProgram(argv=argv, addplugins=[NipyDoctest()])


import atom
import unittest
import numpy as np

class AtomUnittest(unittest.TestCase):
    """Unit tests for the Atom Class"""
    def test_creation(self):
        """Testing the default and copy constructors"""
        target_string = "H1 1 2 3"
        first = atom.Atom("H1",np.array([1,2,3]))
        second = first.copy()
        self.assertTrue(target_string == str(first))
        self.assertTrue(target_string == str(second))

    def test_move(self):
        """Testing that Atom.move() functions properly """
        test_atom = atom.Atom("TEST", np.array([0,0,0]))
        test_atom.move([1,2,3])

        self.assertListEqual(list(test_atom.coords),[1,2,3])

        test_atom.move([-3,-3,-3])
        self.assertListEqual(list(test_atom.coords),[-2,-1,0])

    def test_repr_str_equivalent(self):
        """Testing that the __str__() and __repr__() methods produce equivalent, correct output"""
        target_string = "TEST -1000 0 1000"
        test_atom = atom.Atom("TEST",np.array([-1000,0,1000]))

        self.assertEqual(target_string,str(test_atom))
        self.assertEqual(str(test_atom),test_atom.__repr__())

if __name__ == '__main__':
    unittest.main()