#################################################################################
# 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
# -*- 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)


#!/usr/bin/env python
# coding: utf-8

#
#

# # PHY321: Classical Mechanics 1
# **Homework 1, due January 21 (midnight)**
#
# Date: **Jan 11, 2022**

# ### Practicalities about  homeworks and projects
#
# 1. You can work in groups (optimal groups are often 2-3 people) or by yourself. If you work as a group you can hand in one answer only if you wish. **Remember to write your name(s)**!
#
# 2. Homeworks (final version) are available approximately ten days before the  deadline.
#
# 3. How do I(we)  hand in?  You can hand in the paper and pencil exercises as a scanned document. For this homework this applies to exercises 1-5. You should upload the scan to D2L. Alternatively, you can hand in everyhting (if you are ok with typing mathematical formulae using say Latex) as a jupyter notebook at D2L. The numerical exercise (exercise 6 here) should always be handed in as a jupyter notebook by the deadline at D2L.

# ### Exercise 1 (12 pt), math reminder, properties of exponential function
#
# The first exercise is meant to remind ourselves about properties of
# the exponential function and imaginary numbers. This is highly
# relevant later in this course when we start analyzing oscillatory
# motion and some wave mechanics. As physicists we should thus feel comfortable with expressions that
# include $\exp{(\imath\omega t)}$. Here $t$ could be interpreted as time and $\omega$ as a frequency and $\imath$ is the imaginary unit number.
#
# * 1a (3pt): Perform Taylor expansions in powers of $\omega t$ of the functions $\cos{(\omega t)}$ and $\sin{(\omega t)}$.
#
# * 1b (3pt): Perform a Taylor expansion of $\exp{(i\omega t)}$.
#
# * 1c (3pt): Using parts (a) and (b) here, show that $\exp{(\imath\omega t)}=\cos{(\omega t)}+\imath\sin{(\omega t)}$.
#
# * 1d (3pt): Show that $\ln{(−1)} = \imath\pi$.

# ### Exercise 2 (12 pt), Vector algebra
#
# * 2a (6pt) One of the many uses of the scalar product is to find the angle between two given vectors. Find the angle between the vectors $\boldsymbol{a}=(1,2,4)$ and $\boldsymbol{b}=(4,2,1)$ by evaluating their scalar product.
#
# * 2b (6pt) For a cube with sides of length 1, one vertex at the origin, and sides along the $x$, $y$, and $z$ axes, the vector of the body diagonal from the origin can be written $\boldsymbol{a}=(1, 1, 1)$ and the vector of the face diagonal in the $xy$ plane from the origin is $\boldsymbol{b}=(1,1,0)$. Find first the lengths of the body diagonal and the face diagonal. Use then part (2a) to find the angle between the body diagonal and the face diagonal.

# ### Exercise 3 (10 pt), More vector mathematics
#
# * 3a (5pt) Show (using the fact that multiplication of reals is distributive) that $\boldsymbol{a}(\boldsymbol{b}+\boldsymbol{c})=\boldsymbol{a}\boldsymbol{b}+\boldsymbol{a}\boldsymbol{c}$.
#
# * 3b (5pt) Show that (using product rule for differentiating reals)  $\frac{d}{dt}(\boldsymbol{a}\boldsymbol{b})=\boldsymbol{a}\frac{d\boldsymbol{b}}{dt}+\boldsymbol{b}\frac{d\boldsymbol{a}}{dt}$

# ### Exercise 4 (10 pt), Algebra of cross products
#
# * 4a (5pt) Show that the cross products are distribuitive $\boldsymbol{a}\times(\boldsymbol{b}+\boldsymbol{c})=\boldsymbol{a}\times\boldsymbol{b}+\boldsymbol{a}\times\boldsymbol{c}$.
#
# * 4b (5pt) Show that $\frac{d}{dt}(\boldsymbol{a}\times\boldsymbol{b})=\boldsymbol{a}\times\frac{d\boldsymbol{b}}{dt}+\frac{d\boldsymbol{a}}{dt}\times \boldsymbol{b}$. Be careful with the order of factors

# ### Exercise 5 (10 pt), Area of triangle and law of sines
#
# Exercise 1.18 in the textbook of Taylor, Classical Mechanics. Part (1.18a) gives 5pt and part (1.18b) gives also 5pt.

# ### Exercise 6 (40pt), Numerical elements, getting started with some simple data
#
# **This exercise should be handed in as a jupyter-notebook** at D2L. Remember to write your name(s).
#
# Our first numerical attempt will involve reading data from file or
# just setting up two vectors, one for position and one for time. Our data are from
# [Usain Bolt's world record 100m during the olympic games in Beijing in
# 2008](https://www.youtube.com/watch?v=93dC0o2aHto). The data show the time used in units of 10m (see below). Before we however
# venture into this, we need to repeat some basic Python syntax with an
# emphasis on
#
# * basic Python syntax for arrays
#
# * define and operate on vectors and matrices in Python
#
# * create plots for motion in 1D space
#
# For more information, see the [introductory slides](https://mhjensen.github.io/Physics321/doc/pub/week2/html/week2.html).
# Here are some of the basic packages we will be using this week

# In[1]:

get_ipython().run_line_magic('matplotlib', 'inline')

import numpy as np
import matplotlib.pyplot as plt
get_ipython().run_line_magic('matplotlib', 'inline')

# The first exercise here deals with simply getting familiar with vectors and matrices.
#
# We will be working with vectors and matrices to get you familiar with them
#
# 1. Initalize two three-dimensional $xyz$ vectors in the below cell using np.array([x,y,z]). Vectors are represented through arrays in python
#
# 2. V1 should have x1=1, y1 =2, and z1=3.
#
# 3. Vector 2 should have x2=4, y2=5,  and z2=6.
#
# 4. Print both vectors to make sure your code is working properly.

# In[2]:

V1 = np.array([1,2,3])
V2 = np.array([4,5,6])
print("V1: ", V1)
print("V2: ", V2)

# If this is not too familiar, here's a useful link for creating vectors in python
# <https://docs.scipy.org/doc/numpy-1.13.0/user/basics.creation.html>. Alternatively, look up the [introductory slides](https://mhjensen.github.io/Physics321/doc/pub/week2/html/week2.html).
#
# Now lets do some basic mathematics with vectors.
#
# Compute and print the following, and double check with hand calculations:
#
# * 6a (2pt)  Calculate $\boldsymbol{V}_1-\boldsymbol{V}_2$.
#
# * 6b (2pt)  Calculate $\boldsymbol{V}_2-\boldsymbol{V}_1$.
#
# * 6c (2pt) Calculate the dot product $\boldsymbol{V}_1\boldsymbol{V}_2$.
#
# * 6d (2pt) Calculate the cross product $\boldsymbol{V}_1\times\boldsymbol{V}_2$.
#
# Here is some useful explanation on numpy array operations if you feel a bit confused by what is happening,
# see <https://www.pluralsight.com/guides/overview-basic-numpy-operations>.
#
# The following code prints the first two exercises

# In[3]:

print(V1-V2)
print(V2-V1)

# For the dot product of V1 and V2 below we can use the **dot** function of **numpy** as follows

# In[4]:

print(V1.dot(V2))

# As a small challenge try to write your own function for the **dot** product of two vectors.
#
# Matrices can be created in a similar fashion in python.  In this
# language we can work with them through the package numpy (which we
# have already imported)

# In[5]:

M1 = np.matrix([[1,2,3],
             [4,5,6],
             [7,8,9]])
M2 = np.matrix([[1,2],
             [3,4],
             [5,6]])
M3 = np.matrix([[9,8,7],
             [4,5,6],
             [7,6,9]])

# Matrices can be added in the same way vectors are added in python as shown here

# In[6]:

print("M1+M3: ", M1+M3)

# What happens if we try to do $M1+M2$?
#
# That's enough vectors and matrices for now. Let's move on to some physics problems! Yes, the actual subject we are studying for.
#
# We can opt for two different ways of handling the data. The data is listed in the table here and represents the total time Usain Bolt used in steps of  10 meters of distance. The label $i$ is just a counter and we start from zero since Python arrays are by default set from zero. The variable $t$ is time in seconds and $x$ is the position in meters.
#
#

#

#
 i  	  0  	  1  	  2  	  3  	  4  	  5  	  6  	  7  	  8  	  9

#

#

#
   x[m]    	    10      	    20      	    30      	    40      	    50      	    60      	    70      	    80      	    90      	    100

#
   t[s]    	    1.85    	    2.87    	    3.78    	    4.65    	    5.50    	    6.32    	    7.14    	    7.96    	    8.79    	    9.69

#

#

#
# * 6e (6pt) You can here make a file with the above data and read them in and set up two vectors, one for time and one for position. Alternatively, you can just set up these two vectors directly and define two vectors in your Python code.
#
# The following example code may help here

# In[7]:

# we just initialize time and position
x = np.array([10.0, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0, 80.0, 90.0, 100.0])
t = np.array([1.85, 2.87, 3.78, 4.65, 5.50, 6.32, 7.14, 7.96, 8.79, 9.69])
plt.plot(t,x, color='black')
plt.xlabel("Time t[s]")
plt.ylabel("Position x[m]")
plt.title("Usain Bolt's world record run")
plt.show()

# * 6f (6pt) Plot the position as function of time
#
# * 6g (10pt) Compute thereafter the mean velocity for every interval $i$ and the total velocity (from $i=0$ to the given interval $i$) for each interval and plot these two quantities as function of time. Comment your results.
#
# * 6h (10pt) Finally, compute and plot the mean acceleration for each interval and the total acceleration. Again, comment your results. Can you see whether he slowed down during the last meters?


"""Tools for setting up interactive sessions. """

from sympy.interactive.printing import init_printing

preexec_source = """\
from __future__ import division
from sympy import *
x, y, z, t = symbols('x y z t')
k, m, n = symbols('k m n', integer=True)
f, g, h = symbols('f g h', cls=Function)
"""

verbose_message = """\
These commands were executed:
%(source)s
Documentation can be found at http://www.sympy.org
"""

no_ipython = """\
Couldn't locate IPython. Having IPython installed is greatly recommended.
See http://ipython.scipy.org for more details. If you use Debian/Ubuntu,
just install the 'ipython' package and start isympy again.
"""

def _make_message(ipython=True, quiet=False, source=None):
    """Create a banner for an interactive session. """
    from sympy import __version__ as sympy_version
    from sympy.polys.domains import GROUND_TYPES
    from sympy.utilities.misc import ARCH
    from sympy import SYMPY_DEBUG

    import sys
    import os

    python_version = "%d.%d.%d" % sys.version_info[:3]

    if ipython:
        shell_name = "IPython"
    else:
        shell_name = "Python"

    info = ['ground types: %s' % GROUND_TYPES]

    cache = os.getenv('SYMPY_USE_CACHE')

    if cache is not None and cache.lower() == 'no':
        info.append('cache: off')

    if SYMPY_DEBUG:
        info.append('debugging: on')

    args = shell_name, sympy_version, python_version, ARCH, ', '.join(info)
    message = "%s console for SymPy %s (Python %s-%s) (%s)\n" % args

    if not quiet:
        if source is None:
            source = preexec_source

        _source = ""

        for line in source.split('\n')[:-1]:
            if not line:
                _source += '\n'
            else:
                _source += '>>> ' + line + '\n'

        message += '\n' + verbose_message % {'source': _source}

    return message

def int_to_Integer(s):
    """
    Wrap integer literals with Integer.

    This is based on the decistmt example from
    http://docs.python.org/library/tokenize.html.

    Only integer literals are converted.  Float literals are left alone.
    Example
    =======

    >>> from sympy.interactive.session import int_to_Integer
    >>> from sympy import Integer
    >>> s = '1.2 + 1/2 - 0x12 + a1'
    >>> int_to_Integer(s)
    '1.2 +Integer (1 )/Integer (2 )-Integer (0x12 )+a1 '
    >>> s = 'print (1/2)'
    >>> int_to_Integer(s)
    'print (Integer (1 )/Integer (2 ))'
    >>> exec(s) #doctest: +SKIP
    0.5
    >>> exec(int_to_Integer(s))
    1/2
    """
    from tokenize import generate_tokens, untokenize, NUMBER, NAME, OP
    from StringIO import StringIO

    def _is_int(num):
        """
        Returns true if string value num (with token NUMBER) represents an integer.
        """
        # XXX: Is there something in the standard library that will do this?
        if '.' in num or 'j' in num.lower() or 'e' in num.lower():
            return False
        return True

    result = []
    g = generate_tokens(StringIO(s).readline)   # tokenize the string
    for toknum, tokval, _, _, _  in g:
        if toknum == NUMBER and _is_int(tokval):  # replace NUMBER tokens
            result.extend([
                (NAME, 'Integer'),
                (OP, '('),
                (NUMBER, tokval),
                (OP, ')')
            ])
        else:
            result.append((toknum, tokval))
    return untokenize(result)

# XXX: Something like this might be used, but it only works on single line
# inputs.  See
# http://mail.scipy.org/pipermail/ipython-user/2012-August/010846.html and
# https://github.com/ipython/ipython/issues/1491.  So instead we are forced to
# just monkey-patch run_cell until IPython builds a better API.
#
# class IntTransformer(object):
#     """
#     IPython command line transformer that recognizes and replaces int
#     literals.
#
#     Based on
#     https://bitbucket.org/birkenfeld/ipython-physics/src/71b2d850da00/physics.py.
#
#     """
#     priority = 99
#     enabled = True
#     def transform(self, line, continue_prompt):
#         import re
#         from tokenize import TokenError
#         leading_space = re.compile(' *')
#         spaces = re.match(leading_space, line).span()[1]
#         try:
#             return ' '*spaces + int_to_Integer(line)
#         except TokenError:
#             return line
#
# int_transformer = IntTransformer()
#
# def enable_automatic_int_sympification(app):
#     """
#     Allow IPython to automatically convert integer literals to Integer.
#
#     This lets things like 1/2 be executed as (essentially) Rational(1, 2).
#     """
#     app.shell.prefilter_manager.register_transformer(int_transformer)

def enable_automatic_int_sympification(app):
    """
    Allow IPython to automatically convert integer literals to Integer.
    """
    hasshell = hasattr(app, 'shell')

    import ast
    if hasshell:
        old_run_cell = app.shell.run_cell
    else:
        old_run_cell = app.run_cell
    def my_run_cell(cell, *args, **kwargs):
        try:
            # Check the cell for syntax errors.  This way, the syntax error
            # will show the original input, not the transformed input.  The
            # downside here is that IPython magic like %timeit will not work
            # with transformed input (but on the other hand, IPython magic
            # that doesn't expect transformed input will continue to work).
            ast.parse(cell)
        except SyntaxError:
            pass
        else:
            cell = int_to_Integer(cell)
        old_run_cell(cell, *args, **kwargs)

    if hasshell:
        app.shell.run_cell = my_run_cell
    else:
        app.run_cell = my_run_cell

def enable_automatic_symbols(app):
    """Allow IPython to automatially create symbols (``isympy -a``). """
    # XXX: This should perhaps use tokenize, like int_to_Integer() above.
    # This would avoid re-executing the code, which can lead to subtle
    # issues.  For example:
    #
    # In [1]: a = 1
    #
    # In [2]: for i in range(10):
    #    ...:     a += 1
    #    ...:
    #
    # In [3]: a
    # Out[3]: 11
    #
    # In [4]: a = 1
    #
    # In [5]: for i in range(10):
    #    ...:     a += 1
    #    ...:     print b
    #    ...:
    # b
    # b
    # b
    # b
    # b
    # b
    # b
    # b
    # b
    # b
    #
    # In [6]: a
    # Out[6]: 12
    #
    # Note how the for loop is executed again because `b` was not defined, but `a`
    # was already incremented once, so the result is that it is incremented
    # multiple times.

    import re
    re_nameerror = re.compile("name '(?P<symbol>[A-Za-z_][A-Za-z0-9_]*)' is not defined")

    def _handler(self, etype, value, tb, tb_offset=None):
        """Handle :exc:`NameError` exception and allow injection of missing symbols. """
        if etype is NameError and tb.tb_next and not tb.tb_next.tb_next:
            match = re_nameerror.match(str(value))

            if match is not None:
                # XXX: Make sure Symbol is in scope. Otherwise you'll get infinite recursion.
                self.run_cell("%(symbol)s = Symbol('%(symbol)s')" %
                    {'symbol': match.group("symbol")}, store_history=False)

                try:
                    code = self.user_ns['In'][-1]
                except (KeyError, IndexError):
                    pass
                else:
                    self.run_cell(code, store_history=False)
                    return None
                finally:
                    self.run_cell("del %s" % match.group("symbol"),
                                  store_history=False)

        stb = self.InteractiveTB.structured_traceback(etype, value, tb, tb_offset=tb_offset)
        self._showtraceback(etype, value, stb)

    if hasattr(app, 'shell'):
        app.shell.set_custom_exc((NameError,), _handler)
    else:
        # This was restructured in IPython 0.13
        app.set_custom_exc((NameError,), _handler)

def init_ipython_session(argv=[], auto_symbols=False, auto_int_to_Integer=False):
    """Construct new IPython session. """
    import IPython

    if IPython.__version__ >= '0.11':
        # use an app to parse the command line, and init config
        from IPython.frontend.terminal import ipapp
        app = ipapp.TerminalIPythonApp()

        # don't draw IPython banner during initialization:
        app.display_banner = False
        app.initialize(argv)

        if auto_symbols:
            enable_automatic_symbols(app)
        if auto_int_to_Integer:
            enable_automatic_int_sympification(app)

        return app.shell
    else:
        from IPython.Shell import make_IPython
        return make_IPython(argv)

def init_python_session():
    """Construct new Python session. """
    from code import InteractiveConsole

    class SymPyConsole(InteractiveConsole):
        """An interactive console with readline support. """

        def __init__(self):
            InteractiveConsole.__init__(self)

            try:
                import readline
            except ImportError:
                pass
            else:
                import os
                import atexit

                readline.parse_and_bind('tab: complete')

                if hasattr(readline, 'read_history_file'):
                    history = os.path.expanduser('~/.sympy-history')

                    try:
                        readline.read_history_file(history)
                    except IOError:
                        pass

                    atexit.register(readline.write_history_file, history)

    return SymPyConsole()

def init_session(ipython=None, pretty_print=True, order=None,
        use_unicode=None, quiet=False, auto_symbols=False, auto_int_to_Integer=False, argv=[]):
    """
    Initialize an embedded IPython or Python session. The IPython session is
    initiated with the --pylab option, without the numpy imports, so that
    matplotlib plotting can be interactive.

    Parameters
    ==========

    pretty_print: boolean
        If True, use pretty_print to stringify;
        if False, use sstrrepr to stringify.
    order: string or None
        There are a few different settings for this parameter:
        lex (default), which is lexographic order;
        grlex, which is graded lexographic order;
        grevlex, which is reversed graded lexographic order;
        old, which is used for compatibility reasons and for long expressions;
        None, which sets it to lex.
    use_unicode: boolean or None
        If True, use unicode characters;
        if False, do not use unicode characters.
    quiet: boolean
        If True, init_session will not print messages regarding its status;
        if False, init_session will print messages regarding its status.
    auto_symbols: boolean
        If True, IPython will automatically create symbols for you.
        If False, it will not.
        The default is False.
    auto_int_to_Integer: boolean
        If True, IPython will automatically wrap int literals with Integer, so
        that things like 1/2 give Rational(1, 2).
        If False, it will not.
        The default is False.
    ipython: boolean or None
        If True, printing will initialize for an IPython console;
        if False, printing will initialize for a normal console;
        The default is None, which does what False does.
    argv: list of arguments for IPython
        See sympy.bin.isympy for options that can be used to initialize IPython.

    See Also
    ========

    sympy.interactive.printing.init_printing: for examples and the rest of the parameters.

    Examples
    ========

    >>> from sympy import init_session, Symbol, sin, sqrt
    >>> sin(x) #doctest: +SKIP
    NameError: name 'x' is not defined
    >>> init_session() #doctest: +SKIP
    >>> sin(x) #doctest: +SKIP
    sin(x)
    >>> sqrt(5) #doctest: +SKIP
      ___
    \/ 5
    >>> init_session(pretty_print=False) #doctest: +SKIP
    >>> sqrt(5) #doctest: +SKIP
    sqrt(5)
    >>> y + x + y**2 + x**2 #doctest: +SKIP
    x**2 + x + y**2 + y
    >>> init_session(order='grlex') #doctest: +SKIP
    >>> y + x + y**2 + x**2 #doctest: +SKIP
    x**2 + y**2 + x + y
    >>> init_session(order='grevlex') #doctest: +SKIP
    >>> y * x**2 + x * y**2 #doctest: +SKIP
    x**2*y + x*y**2
    >>> init_session(order='old') #doctest: +SKIP
    >>> x**2 + y**2 + x + y #doctest: +SKIP
    x + y + x**2 + y**2
    >>> theta = Symbol('theta') #doctest: +SKIP
    >>> theta #doctest: +SKIP
    theta
    >>> init_session(use_unicode=True) #doctest: +SKIP
    >>> theta # doctest: +SKIP
    \u03b8
    """
    import sys

    in_ipython = False

    if ipython is False:
        ip = init_python_session()
        mainloop = ip.interact
    else:
        try:
            import IPython
        except ImportError:
            if ipython is not True:
                if not quiet:
                    print no_ipython
                ip = init_python_session()
                mainloop = ip.interact
            else:
                raise RuntimeError("IPython is not available on this system")
        else:
            ipython = True

            if IPython.__version__ >= '0.11':
                try:
                    ip = get_ipython()
                except NameError:
                    ip = None
            else:
                ip = IPython.ipapi.get()
                if ip:
                    ip = ip.IP

            if ip is not None:
                in_ipython = True
            else:
                ip = init_ipython_session(argv=argv,
                    auto_symbols=auto_symbols, auto_int_to_Integer=auto_int_to_Integer)

            if IPython.__version__ >= '0.11':
                # runsource is gone, use run_cell instead, which doesn't
                # take a symbol arg.  The second arg is `store_history`,
                # and False means don't add the line to IPython's history.
                ip.runsource = lambda src, symbol='exec': ip.run_cell(src, False)

                #Enable interactive plotting using pylab.
                try:
                    ip.enable_pylab(import_all=False)
                except ImportError:
                    #Causes an import error if matplotlib is not installed.
                    pass
            if not in_ipython:
                mainloop = ip.mainloop

    if auto_symbols and (not ipython or IPython.__version__ < '0.11'):
        raise RuntimeError("automatic construction of symbols is possible only in IPython 0.11 or above")
    if auto_int_to_Integer and (not ipython or IPython.__version__ < '0.11'):
        raise RuntimeError("automatic int to Integer transformation is possible only in IPython 0.11 or above")

    _preexec_source = preexec_source

    ip.runsource(_preexec_source, symbol='exec')
    init_printing(pretty_print=pretty_print, order=order, use_unicode=use_unicode, ip=ip)

    message = _make_message(ipython, quiet, _preexec_source)

    if not in_ipython:
        mainloop(message)
        sys.exit('Exiting ...')
    else:
        ip.write(message)
        ip.set_hook('shutdown_hook', lambda ip: ip.write("Exiting ...\n"))


#!/usr/bin/env python
# coding: utf-8

#
#
# #  Gravitational Wave Open Data Workshop #3
#
#
# #### Tutorial 1.2: Introduction to GWpy
#
# This tutorial will briefly describe GWpy, a python package for gravitational astrophysics, and walk-through how you can use this to speed up access to, and processing of, GWOSC data.
#
# [Click this link to view this tutorial in Google Colaboratory](https://colab.research.google.com/github/gw-odw/odw-2019/blob/master/Day_1/Tuto%201.2%20Open%20Data%20access%20with%20GWpy.ipynb)
#
#
This notebook were generated using python 3.7, but should work on python 2.7, 3.6, or 3.7.

# ##  Installation  (execute only if running on a cloud platform or if you haven't done the installation already!)
#
# Note: we use [`pip`](https://docs.python.org/3.6/installing/), but **it is recommended** to use [conda](https://docs.ligo.org/lscsoft/conda/) on your own machine, as explained in the [installation instructions](https://github.com/gw-odw/odw-2019/blob/master/setup.md). This usage might look a little different than normal, simply because we want to do this directly from the notebook.

# In[2]:

# -- Uncomment following line if running in Google Colab
#! pip install -q 'gwpy==1.0.1'

# **Important:** With Google Colab, you may need to restart the runtime after running the cell above.

# ## Initialization

# In[1]:

import gwpy
print(gwpy.__version__)

# ## A note on object-oriented programming
#
# Before we dive too deeply, its worth a quick aside on object-oriented programming (OOP).
# GWpy is heavily object-oriented, meaning almost all of the code you run using GWpy is based around an object of some type, e.g. `TimeSeries`.
# Most of the methods (functions) we will use are attached to an object, rather than standing alone, meaning you should have a pretty good idea of what sort of data you are dealing with (without having to read the documentation!).
#
# For a quick overview of object-oriented programming in Python, see [this blog post by Jeff Knupp](https://jeffknupp.com/blog/2014/06/18/improve-your-python-python-classes-and-object-oriented-programming/).

# ## Handling data in the time domain

# #### Finding open data
#
# We have seen already that the `gwosc` module can be used to query for what data are available on GWOSC.
# The next thing to do is to actually read some open data. Let's try to get some for GW150914, the first direct detection of an astrophysical gravitational-wave signal from a BBH (binary black hole system).
#
# We can use the [`TimeSeries.fetch_open_data`](https://gwpy.github.io/docs/stable/api/gwpy.timeseries.TimeSeries.html#gwpy.timeseries.TimeSeries.fetch_open_data) method to download data directly from https://www.gw-openscience.org, but we need to know the GPS times.
# We can query for the GPS time of an event as follows:

# In[2]:

from gwosc.datasets import event_gps
gps = event_gps('GW150914')
print(gps)

# Now we can build a `[start, end)` GPS segment to 10 seconds around this time, using integers for convenience:

# In[3]:

segment = (int(gps)-5, int(gps)+5)
print(segment)

# and can now query for the full data.
# For this example we choose to retrieve data for the LIGO-Livingston interferometer, using the identifier `'L1'`.
# We could have chosen any of
#
# - `'G1`' - GEO600
# - `'H1'` - LIGO-Hanford
# - `'L1'` - LIGO-Livingston
# - `'V1'` - (Advanced) Virgo
#
# In the future, the Japanese observatory KAGRA will come online, with the identifier `'K1'`.

# In[4]:

from gwpy.timeseries import TimeSeries
ldata = TimeSeries.fetch_open_data('L1', *segment, verbose=True)
print(ldata)

# ##### The `verbose=True` flag lets us see that GWpy has discovered two files that provides the data for the given interval, downloaded them, and loaded the data.
# The files are not stored permanently, so next time you do the same call, it will be downloaded again, however, if you know you might repeat the same call many times, you can use `cache=True` to store the file on your computer.
#
# Notes:
#
# * To read data from a local file instead of from the GWOSC server, we can use [`TimeSeries.read`](https://gwpy.github.io/docs/stable/api/gwpy.timeseries.TimeSeries.html#gwpy.timeseries.TimeSeries.read) method.
#
# We have now downloaded real LIGO data for GW150914! These are the actual data used in the analysis that discovered the first binary black hole merger.

# To sanity check things, we can easily make a plot, using the [`plot()`](https://gwpy.github.io/docs/stable/timeseries/plot.html) method of the `data` `TimeSeries`.

#

# Since this is the first time we are plotting something in this notebook, we need to make configure `matplotlib` (the plotting library) to work within the notebook properly:
#

# Matplotlib documentation can be found [`here`](https://matplotlib.org/contents.html).

# In[5]:

get_ipython().run_line_magic('matplotlib', 'inline')
plot = ldata.plot()

# Notes: There are alternatives ways to access the GWOSC data.
#
# * [`readligo`](https://losc.ligo.org/s/sample_code/readligo.py) is a light-weight Python module that returns the time series into a Numpy array.
# * The [PyCBC](http://github.com/ligo-cbc/pycbc) package has the `pycbc.frame.query_and_read_frame` and `pycbc.frame.read_frame` methods. We use [PyCBC](http://github.com/ligo-cbc/pycbc) in Tutorial 2.1, 2.2 and 2.3.

# ## Handling data in the frequency domain using the Fourier transform
#
# The [Fourier transform](https://en.wikipedia.org/wiki/Fourier_transform) is a widely-used mathematical tool to expose the frequency-domain content of a time-domain signal, meaning we can see which frequencies contian lots of power, and which have less.
#
# We can calculate the Fourier transform of our `TimeSeries` using the [`fft()`](https://gwpy.github.io/docs/stable/api/gwpy.timeseries.TimeSeries.html#gwpy.timeseries.TimeSeries.fft) method:

# In[6]:

fft = ldata.fft()
print(fft)

# The result is a [`FrequencySeries`](https://gwpy.github.io/docs/stable/frequencyseries/), with complex amplitude, representing the amplitude and phase of each frequency in our data.
# We can use `abs()` to extract the amplitude and plot that:

# In[7]:

plot = fft.abs().plot(xscale="log", yscale="log")
plot.show(warn=False)

# This doesn't look correct at all!
# The problem is that the FFT works under the assumption that our data are periodic, which means that the edges of our data look like discontinuities when transformed.
# We need to apply a window function to our time-domain data before transforming, which we can do using the [`scipy.signal`](https://docs.scipy.org/doc/scipy/reference/signal.html) module:

# In[8]:

from scipy.signal import get_window
window = get_window('hann', ldata.size)
lwin = ldata * window

# Let's try our transform again and see what we get

# In[9]:

fftamp = lwin.fft().abs()
plot = fftamp.plot(xscale="log", yscale="log")
plot.show(warn=False)

# This looks a little more like what we expect for the amplitude spectral density of a gravitational-wave detector.

# ## Calculating the power spectral density
#
# In practice, we typically use a large number of FFTs to estimate an averages power spectral density over a long period of data.
# We can do this using the [`asd()`](https://gwpy.github.io/docs/stable/api/gwpy.timeseries.TimeSeries.html#gwpy.timeseries.TimeSeries.asd) method, which uses [Welch's method](https://en.wikipedia.org/wiki/Welch%27s_method) to combine FFTs of overlapping, windowed chunks of data.

# In[10]:

asd = ldata.asd(fftlength=4, method="median")
plot = asd.plot()
plot.show(warn=False)

# In[11]:

ax = plot.gca()
ax.set_xlim(10, 1400)
ax.set_ylim(2e-24, 1e-20)
plot

# The ASD is a standard tool used to study the frequency-domain sensitivity of a gravitational-wave detector.
# For the LIGO-Livingston data we loaded, we can see large spikes at certain frequencies, including
#
# - ~300 Hz
# - ~500 Hz
# - ~1000 Hz
#
# The [O2 spectral lines](https://www.gw-openscience.org/o2speclines/) page on GWOSC describes a number of these spectral features for O2, with some of them being forced upon us, and some being deliberately introduced to help with interferometer control.
#
# Loading more data allows for more FFTs to be averaged during the ASD calculation, meaning random variations get averaged out, and we can see more detail:

# In[12]:

ldata2 = TimeSeries.fetch_open_data('L1', int(gps)-512, int(gps)+512, cache=True)
lasd2 = ldata2.asd(fftlength=4, method="median")
plot = lasd2.plot()
ax = plot.gca()
ax.set_xlim(10, 1400)
ax.set_ylim(5e-24, 1e-20)
plot.show(warn=False)

# Now we can see some more features, including sets of lines around ~30 Hz and ~65 Hz, and some more isolate lines through the more sensitive region.
#
# For comparison, we can load the LIGO-Hanford data and plot that as well:

# In[13]:

# get Hanford data
hdata2 = TimeSeries.fetch_open_data('H1', int(gps)-512, int(gps)+512, cache=True)
hasd2 = hdata2.asd(fftlength=4, method="median")

# and plot using standard colours
ax.plot(hasd2, label='LIGO-Hanford', color='gwpy:ligo-hanford')

# update the Livingston line to use standard colour, and have a label
lline = ax.lines[0]
lline.set_color('gwpy:ligo-livingston')  # change colour of Livingston data
lline.set_label('LIGO-Livingston')

ax.set_ylabel(r'Strain noise [$1/\sqrt{\mathrm{Hz}}$]')
ax.legend()
plot

# Now we can see clearly the relative sensitivity of each LIGO instrument, the common features between both, and those unique to each observatory.

# # Challenges:

# ##### Quiz Question 1:
#
# The peak amplitude in the LIGO-Livingston data occurs at approximately 5 seconds into the plot above and is undetectable above the background noise by the eye. Plot the data for the LIGO-Hanford detector around GW150914. Looking at your new LIGO-Handford plot, can your eye identify a signal peak?

# In[14]:

from gwosc.datasets import event_gps
gps1 = event_gps('GW150914')
print(gps1)
segment1 = (int(gps)-7, int(gps)+7)
print(segment1)

from gwpy.timeseries import TimeSeries
ldata1 = TimeSeries.fetch_open_data('L1', *segment1, verbose=True)
print(ldata1)

get_ipython().run_line_magic('matplotlib', 'inline')
plot1 = ldata1.plot()

# In[24]:

import matplotlib.pyplot as plt

ldata2 = TimeSeries.fetch_open_data('L1', int(gps)-512, int(gps)+512, cache=True)
lasd2 = ldata2.asd(fftlength=4, method="median")

# get Hanford data
hdata2 = TimeSeries.fetch_open_data('H1', int(gps)-512, int(gps)+512, cache=True)
hasd2 = hdata2.asd(fftlength=4, method="median")

fig, ax = plt.subplots(figsize=(13,7))
# and plot using standard colours
ax.plot(hasd2, label='LIGO-Hanford', color='gwpy:ligo-hanford')
ax.plot(lasd2, label='LIGO-Livingston', color='gwpy:ligo-livingston')

ax.set_xlim(2, 1400)
ax.set_ylim(5e-24, 2e-19)
# update the Livingston line to use standard colour, and have a label

ax.set_yscale('log')
ax.set_xscale('log')
ax.set_ylabel(r'Strain noise [$1/\sqrt{\mathrm{Hz}}$]')
ax.set_xlabel(r'Frequency [Hz]')

ax.legend()
plt.show()

# # Quiz Question 2 :
#
# Make an ASD around the time of an O3 event, GW190412 for L1 detector .  Compare this with the ASDs around GW150914 for L1 detector.  Which data have lower noise - and so are more sensitive - around 100 Hz?
#

# In[25]:

from gwosc.datasets import event_gps
gps_GW190412 = event_gps('GW190412')
print(gps_GW190412)
segment_GW190412 = (int(gps_GW190412)-7, int(gps_GW190412)+7)
print(segment_GW190412)

from gwpy.timeseries import TimeSeries
ldata_GW190412 = TimeSeries.fetch_open_data('L1', *segment_GW190412, verbose=True)
print(ldata_GW190412)

get_ipython().run_line_magic('matplotlib', 'inline')
plot_GW190412 = ldata_GW190412.plot()

# In[32]:

import matplotlib.pyplot as plt

ldata_GW190412 = TimeSeries.fetch_open_data('L1', int(gps_GW190412)-512, int(gps_GW190412)+512, cache=True)
lasd_GW190412 = ldata_GW190412.asd(fftlength=4, method="median")

# get Hanford data
hdata_GW190412 = TimeSeries.fetch_open_data('H1', int(gps_GW190412)-512, int(gps_GW190412)+512, cache=True)
hasd_GW190412 = hdata_GW190412.asd(fftlength=4, method="median")

fig, ax = plt.subplots(figsize=(13,7))
# and plot using standard colours
ax.plot(hasd_GW190412, label='LIGO-Hanford', color='gwpy:ligo-hanford')
ax.plot(lasd_GW190412, label='LIGO-Livingston', color='gwpy:ligo-livingston')

ax.set_xlim(2, 1400)
ax.set_ylim(4e-24, 7e-19)
# update the Livingston line to use standard colour, and have a label

ax.set_yscale('log')
ax.set_xscale('log')
ax.set_ylabel(r'Strain noise [$1/\sqrt{\mathrm{Hz}}$]')
ax.set_xlabel(r'Frequency [Hz]')

ax.legend()
plt.show()

# In[33]:

import matplotlib.pyplot as plt

ldata_GW190412 = TimeSeries.fetch_open_data('L1', int(gps_GW190412)-512, int(gps_GW190412)+512, cache=True)
lasd_GW190412 = ldata_GW190412.asd(fftlength=4, method="median")

ldata2 = TimeSeries.fetch_open_data('L1', int(gps)-512, int(gps)+512, cache=True)
lasd2 = ldata2.asd(fftlength=4, method="median")

fig, ax = plt.subplots(figsize=(13,7))
# and plot using standard colours
ax.plot(lasd2, label='LIGO-L1-GW150914', color='blue')
ax.plot(lasd_GW190412, label='LIGO-L1-GW190412', color='green')

ax.set_xlim(2, 1400)
ax.set_ylim(4e-24, 7e-19)
# update the Livingston line to use standard colour, and have a label

ax.set_yscale('log')
ax.set_xscale('log')
ax.set_ylabel(r'Strain noise [$1/\sqrt{\mathrm{Hz}}$]')
ax.set_xlabel(r'Frequency [Hz]')

ax.legend()
plt.show()

# # The GW190412 data has lesser noise around 100 Hz

# In[ ]:

gw-odw_Day1_with_Solns/Tuto 1.2 Open Data access with GWpy with solutions.py


# python src/chapter3/chapter3_2.py
# python3 src/chapter3/chapter3_2.py
from __future__ import division, absolute_import, print_function
import sys
import math

from copy import copy
from copy import deepcopy

import numpy as np
from numpy import arange

from matplotlib.pyplot import plot
from matplotlib.pyplot import figure
from matplotlib.pyplot import show

class Chapter3_2:
    '''
    CLRS 第三章 3.2 算法函数和笔记
    '''
    def note(self):
        '''
        Summary
        =
        Print chapter3.2 note

        Example
        =
        >>> Chapter3_2().note()
        '''
        print('3.2 标准记号和常用函数')
        print('单调性：一个函数f(n)是单调递增的，若m<=n,则有f(m)<=f(n)，反之单调递减，将小于等于号换成小于号，即变为严格不等式，则函数是严格单调递增的')
        print('下取整(floor)和上取整(ceiling)')
        print('取模运算(modular arithmetic)')
        print('多项式定义及其性质')
        print('指数式定义及其性质')
        print('任何底大于1的指数函数比任何多项式函数增长得更快')
        print('对数定义及其性质')
        print('阶乘定义及其性质')
        print('计算机工作者常常认为对数的底取2最自然，因为很多算法和数据结构都涉及到对问题进行二分')
        print('任意正的多项式函数都比多项对数函数增长得快')
        print('斯特林近似公式：n!=sqrt(2*pi*n)*(n/e)^n*(1+Θ(1/n))')
        print('阶乘函数的一个更紧确的上界和下界：')
        print('n!=o(n^n) n!=ω(2^n) lg(n!)=Θ(nlgn)')
        print('函数迭代的定义和性质')
        print('多重对数函数：用记号lg * n(读作n的log星)来表示多重对数，定义为lg * n=min(i>=0;lg^(i)n<=1)')
        print('多重函数是一种增长很慢的函数')
        print('lg * 2 = 1; lg * 4 = 2; lg * 16 = 3; lg * 65536 = 4; lg * 2^65536 = 5')
        print('宇宙中可以观察到的原子数目估计约有10^80，远远小于2^65536,因此很少会遇到一个使lg * n > 5的一个n输入规模')
        print('斐波那契数列：F0 = 0 F1 = 1 F(i) = F(i-1) + F(i-2),产生的序列为0,1,1,2,3,5,8,13,21,34,55,……')
        print('斐波那契数列和黄金分割率φ以及共轭有关系')
        print('φ=((1+sqrt(5))/2=1.61803 和它的共轭(1-sqrt(5))/2=-0.61803)')
        print('练习题和思考题略')
        # python src/chapter3/chapter3_2.py
        # python3 src/chapter3/chapter3_2.py
        return self

if __name__ == '__main__':
    print('Run main : single chapter three!')
    Chapter3_2().note()
else:
    pass