#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

from __future__ import print_function
import sys
import warnings
from functools import reduce
from threading import RLock

if sys.version >= '3':
    basestring = unicode = str
else:
    from itertools import imap as map

from pyspark import since
from pyspark.rdd import RDD, ignore_unicode_prefix
from pyspark.sql.catalog import Catalog
from pyspark.sql.conf import RuntimeConfig
from pyspark.sql.dataframe import DataFrame
from pyspark.sql.readwriter import DataFrameReader
from pyspark.sql.streaming import DataStreamReader
from pyspark.sql.types import Row, DataType, StringType, StructType, _verify_type, \
    _infer_schema, _has_nulltype, _merge_type, _create_converter, _parse_datatype_string
from pyspark.sql.utils import install_exception_handler

__all__ = ["SparkSession"]

def _monkey_patch_RDD(sparkSession):
    def toDF(self, schema=None, sampleRatio=None):
        """
        Converts current :class:`RDD` into a :class:`DataFrame`

        This is a shorthand for ``spark.createDataFrame(rdd, schema, sampleRatio)``

        :param schema: a :class:`pyspark.sql.types.StructType` or list of names of columns
        :param samplingRatio: the sample ratio of rows used for inferring
        :return: a DataFrame

        >>> rdd.toDF().collect()
        [Row(name=u'Alice', age=1)]
        """
        return sparkSession.createDataFrame(self, schema, sampleRatio)

    RDD.toDF = toDF

class SparkSession(object):
    """The entry point to programming Spark with the Dataset and DataFrame API.

    A SparkSession can be used create :class:`DataFrame`, register :class:`DataFrame` as
    tables, execute SQL over tables, cache tables, and read parquet files.
    To create a SparkSession, use the following builder pattern:

    >>> spark = SparkSession.builder \\
    ...     .master("local") \\
    ...     .appName("Word Count") \\
    ...     .config("spark.some.config.option", "some-value") \\
    ...     .getOrCreate()
    """

    class Builder(object):
        """Builder for :class:`SparkSession`.
        """

        _lock = RLock()
        _options = {}

        @since(2.0)
        def config(self, key=None, value=None, conf=None):
            """Sets a config option. Options set using this method are automatically propagated to
            both :class:`SparkConf` and :class:`SparkSession`'s own configuration.

            For an existing SparkConf, use `conf` parameter.

            >>> from pyspark.conf import SparkConf
            >>> SparkSession.builder.config(conf=SparkConf())
            <pyspark.sql.session...

            For a (key, value) pair, you can omit parameter names.

            >>> SparkSession.builder.config("spark.some.config.option", "some-value")
            <pyspark.sql.session...

            :param key: a key name string for configuration property
            :param value: a value for configuration property
            :param conf: an instance of :class:`SparkConf`
            """
            with self._lock:
                if conf is None:
                    self._options[key] = str(value)
                else:
                    for (k, v) in conf.getAll():
                        self._options[k] = v
                return self

        @since(2.0)
        def master(self, master):
            """Sets the Spark master URL to connect to, such as "local" to run locally, "local[4]"
            to run locally with 4 cores, or "spark://master:7077" to run on a Spark standalone
            cluster.

            :param master: a url for spark master
            """
            return self.config("spark.master", master)

        @since(2.0)
        def appName(self, name):
            """Sets a name for the application, which will be shown in the Spark web UI.

            If no application name is set, a randomly generated name will be used.

            :param name: an application name
            """
            return self.config("spark.app.name", name)

        @since(2.0)
        def enableHiveSupport(self):
            """Enables Hive support, including connectivity to a persistent Hive metastore, support
            for Hive serdes, and Hive user-defined functions.
            """
            return self.config("spark.sql.catalogImplementation", "hive")

        @since(2.0)
        def getOrCreate(self):
            """Gets an existing :class:`SparkSession` or, if there is no existing one, creates a
            new one based on the options set in this builder.

            This method first checks whether there is a valid global default SparkSession, and if
            yes, return that one. If no valid global default SparkSession exists, the method
            creates a new SparkSession and assigns the newly created SparkSession as the global
            default.

            >>> s1 = SparkSession.builder.config("k1", "v1").getOrCreate()
            >>> s1.conf.get("k1") == s1.sparkContext.getConf().get("k1") == "v1"
            True

            In case an existing SparkSession is returned, the config options specified
            in this builder will be applied to the existing SparkSession.

            >>> s2 = SparkSession.builder.config("k2", "v2").getOrCreate()
            >>> s1.conf.get("k1") == s2.conf.get("k1")
            True
            >>> s1.conf.get("k2") == s2.conf.get("k2")
            True
            """
            with self._lock:
                from pyspark.context import SparkContext
                from pyspark.conf import SparkConf
                session = SparkSession._instantiatedSession
                if session is None or session._sc._jsc is None:
                    sparkConf = SparkConf()
                    for key, value in self._options.items():
                        sparkConf.set(key, value)
                    sc = SparkContext.getOrCreate(sparkConf)
                    # This SparkContext may be an existing one.
                    for key, value in self._options.items():
                        # we need to propagate the confs
                        # before we create the SparkSession. Otherwise, confs like
                        # warehouse path and metastore url will not be set correctly (
                        # these confs cannot be changed once the SparkSession is created).
                        sc._conf.set(key, value)
                    session = SparkSession(sc)
                for key, value in self._options.items():
                    session._jsparkSession.sessionState().conf().setConfString(key, value)
                for key, value in self._options.items():
                    session.sparkContext._conf.set(key, value)
                return session

    builder = Builder()

    _instantiatedSession = None

    @ignore_unicode_prefix
    def __init__(self, sparkContext, jsparkSession=None):
        """Creates a new SparkSession.

        >>> from datetime import datetime
        >>> spark = SparkSession(sc)
        >>> allTypes = sc.parallelize([Row(i=1, s="string", d=1.0, l=1,
        ...     b=True, list=[1, 2, 3], dict={"s": 0}, row=Row(a=1),
        ...     time=datetime(2014, 8, 1, 14, 1, 5))])
        >>> df = allTypes.toDF()
        >>> df.createOrReplaceTempView("allTypes")
        >>> spark.sql('select i+1, d+1, not b, list[1], dict["s"], time, row.a '
        ...            'from allTypes where b and i > 0').collect()
        [Row((i + CAST(1 AS BIGINT))=2, (d + CAST(1 AS DOUBLE))=2.0, (NOT b)=False, list[1]=2, \
            dict[s]=0, time=datetime.datetime(2014, 8, 1, 14, 1, 5), a=1)]
        >>> df.rdd.map(lambda x: (x.i, x.s, x.d, x.l, x.b, x.time, x.row.a, x.list)).collect()
        [(1, u'string', 1.0, 1, True, datetime.datetime(2014, 8, 1, 14, 1, 5), 1, [1, 2, 3])]
        """
        from pyspark.sql.context import SQLContext
        self._sc = sparkContext
        self._jsc = self._sc._jsc
        self._jvm = self._sc._jvm
        if jsparkSession is None:
            jsparkSession = self._jvm.SparkSession(self._jsc.sc())
        self._jsparkSession = jsparkSession
        self._jwrapped = self._jsparkSession.sqlContext()
        self._wrapped = SQLContext(self._sc, self, self._jwrapped)
        _monkey_patch_RDD(self)
        install_exception_handler()
        # If we had an instantiated SparkSession attached with a SparkContext
        # which is stopped now, we need to renew the instantiated SparkSession.
        # Otherwise, we will use invalid SparkSession when we call Builder.getOrCreate.
        if SparkSession._instantiatedSession is None \
                or SparkSession._instantiatedSession._sc._jsc is None:
            SparkSession._instantiatedSession = self

    def _repr_html_(self):
        return """

SparkSession - {catalogImplementation}

                {sc_HTML}

        """.format(
            catalogImplementation=self.conf.get("spark.sql.catalogImplementation"),
            sc_HTML=self.sparkContext._repr_html_()
        )

    @since(2.0)
    def newSession(self):
        """
        Returns a new SparkSession as new session, that has separate SQLConf,
        registered temporary views and UDFs, but shared SparkContext and
        table cache.
        """
        return self.__class__(self._sc, self._jsparkSession.newSession())

    @property
    @since(2.0)
    def sparkContext(self):
        """Returns the underlying :class:`SparkContext`."""
        return self._sc

    @property
    @since(2.0)
    def version(self):
        """The version of Spark on which this application is running."""
        return self._jsparkSession.version()

    @property
    @since(2.0)
    def conf(self):
        """Runtime configuration interface for Spark.

        This is the interface through which the user can get and set all Spark and Hadoop
        configurations that are relevant to Spark SQL. When getting the value of a config,
        this defaults to the value set in the underlying :class:`SparkContext`, if any.
        """
        if not hasattr(self, "_conf"):
            self._conf = RuntimeConfig(self._jsparkSession.conf())
        return self._conf

    @property
    @since(2.0)
    def catalog(self):
        """Interface through which the user may create, drop, alter or query underlying
        databases, tables, functions etc.
        """
        if not hasattr(self, "_catalog"):
            self._catalog = Catalog(self)
        return self._catalog

    @property
    @since(2.0)
    def udf(self):
        """Returns a :class:`UDFRegistration` for UDF registration.

        :return: :class:`UDFRegistration`
        """
        from pyspark.sql.context import UDFRegistration
        return UDFRegistration(self._wrapped)

    @since(2.0)
    def range(self, start, end=None, step=1, numPartitions=None):
        """
        Create a :class:`DataFrame` with single :class:`pyspark.sql.types.LongType` column named
        ``id``, containing elements in a range from ``start`` to ``end`` (exclusive) with
        step value ``step``.

        :param start: the start value
        :param end: the end value (exclusive)
        :param step: the incremental step (default: 1)
        :param numPartitions: the number of partitions of the DataFrame
        :return: :class:`DataFrame`

        >>> spark.range(1, 7, 2).collect()
        [Row(id=1), Row(id=3), Row(id=5)]

        If only one argument is specified, it will be used as the end value.

        >>> spark.range(3).collect()
        [Row(id=0), Row(id=1), Row(id=2)]
        """
        if numPartitions is None:
            numPartitions = self._sc.defaultParallelism

        if end is None:
            jdf = self._jsparkSession.range(0, int(start), int(step), int(numPartitions))
        else:
            jdf = self._jsparkSession.range(int(start), int(end), int(step), int(numPartitions))

        return DataFrame(jdf, self._wrapped)

    def _inferSchemaFromList(self, data):
        """
        Infer schema from list of Row or tuple.

        :param data: list of Row or tuple
        :return: :class:`pyspark.sql.types.StructType`
        """
        if not data:
            raise ValueError("can not infer schema from empty dataset")
        first = data[0]
        if type(first) is dict:
            warnings.warn("inferring schema from dict is deprecated,"
                          "please use pyspark.sql.Row instead")
        schema = reduce(_merge_type, map(_infer_schema, data))
        if _has_nulltype(schema):
            raise ValueError("Some of types cannot be determined after inferring")
        return schema

    def _inferSchema(self, rdd, samplingRatio=None):
        """
        Infer schema from an RDD of Row or tuple.

        :param rdd: an RDD of Row or tuple
        :param samplingRatio: sampling ratio, or no sampling (default)
        :return: :class:`pyspark.sql.types.StructType`
        """
        first = rdd.first()
        if not first:
            raise ValueError("The first row in RDD is empty, "
                             "can not infer schema")
        if type(first) is dict:
            warnings.warn("Using RDD of dict to inferSchema is deprecated. "
                          "Use pyspark.sql.Row instead")

        if samplingRatio is None:
            schema = _infer_schema(first)
            if _has_nulltype(schema):
                for row in rdd.take(100)[1:]:
                    schema = _merge_type(schema, _infer_schema(row))
                    if not _has_nulltype(schema):
                        break
                else:
                    raise ValueError("Some of types cannot be determined by the "
                                     "first 100 rows, please try again with sampling")
        else:
            if samplingRatio < 0.99:
                rdd = rdd.sample(False, float(samplingRatio))
            schema = rdd.map(_infer_schema).reduce(_merge_type)
        return schema

    def _createFromRDD(self, rdd, schema, samplingRatio):
        """
        Create an RDD for DataFrame from an existing RDD, returns the RDD and schema.
        """
        if schema is None or isinstance(schema, (list, tuple)):
            struct = self._inferSchema(rdd, samplingRatio)
            converter = _create_converter(struct)
            rdd = rdd.map(converter)
            if isinstance(schema, (list, tuple)):
                for i, name in enumerate(schema):
                    struct.fields[i].name = name
                    struct.names[i] = name
            schema = struct

        elif not isinstance(schema, StructType):
            raise TypeError("schema should be StructType or list or None, but got: %s" % schema)

        # convert python objects to sql data
        rdd = rdd.map(schema.toInternal)
        return rdd, schema

    def _createFromLocal(self, data, schema):
        """
        Create an RDD for DataFrame from a list or pandas.DataFrame, returns
        the RDD and schema.
        """
        # make sure data could consumed multiple times
        if not isinstance(data, list):
            data = list(data)

        if schema is None or isinstance(schema, (list, tuple)):
            struct = self._inferSchemaFromList(data)
            converter = _create_converter(struct)
            data = map(converter, data)
            if isinstance(schema, (list, tuple)):
                for i, name in enumerate(schema):
                    struct.fields[i].name = name
                    struct.names[i] = name
            schema = struct

        elif not isinstance(schema, StructType):
            raise TypeError("schema should be StructType or list or None, but got: %s" % schema)

        # convert python objects to sql data
        data = [schema.toInternal(row) for row in data]
        return self._sc.parallelize(data), schema

    @since(2.0)
    @ignore_unicode_prefix
    def createDataFrame(self, data, schema=None, samplingRatio=None, verifySchema=True):
        """
        Creates a :class:`DataFrame` from an :class:`RDD`, a list or a :class:`pandas.DataFrame`.

        When ``schema`` is a list of column names, the type of each column
        will be inferred from ``data``.

        When ``schema`` is ``None``, it will try to infer the schema (column names and types)
        from ``data``, which should be an RDD of :class:`Row`,
        or :class:`namedtuple`, or :class:`dict`.

        When ``schema`` is :class:`pyspark.sql.types.DataType` or a datatype string, it must match
        the real data, or an exception will be thrown at runtime. If the given schema is not
        :class:`pyspark.sql.types.StructType`, it will be wrapped into a
        :class:`pyspark.sql.types.StructType` as its only field, and the field name will be "value",
        each record will also be wrapped into a tuple, which can be converted to row later.

        If schema inference is needed, ``samplingRatio`` is used to determined the ratio of
        rows used for schema inference. The first row will be used if ``samplingRatio`` is ``None``.

        :param data: an RDD of any kind of SQL data representation(e.g. row, tuple, int, boolean,
            etc.), or :class:`list`, or :class:`pandas.DataFrame`.
        :param schema: a :class:`pyspark.sql.types.DataType` or a datatype string or a list of
            column names, default is ``None``.  The data type string format equals to
            :class:`pyspark.sql.types.DataType.simpleString`, except that top level struct type can
            omit the ``struct<>`` and atomic types use ``typeName()`` as their format, e.g. use
            ``byte`` instead of ``tinyint`` for :class:`pyspark.sql.types.ByteType`. We can also use
            ``int`` as a short name for ``IntegerType``.
        :param samplingRatio: the sample ratio of rows used for inferring
        :param verifySchema: verify data types of every row against schema.
        :return: :class:`DataFrame`

        .. versionchanged:: 2.1
           Added verifySchema.

        >>> l = [('Alice', 1)]
        >>> spark.createDataFrame(l).collect()
        [Row(_1=u'Alice', _2=1)]
        >>> spark.createDataFrame(l, ['name', 'age']).collect()
        [Row(name=u'Alice', age=1)]

        >>> d = [{'name': 'Alice', 'age': 1}]
        >>> spark.createDataFrame(d).collect()
        [Row(age=1, name=u'Alice')]

        >>> rdd = sc.parallelize(l)
        >>> spark.createDataFrame(rdd).collect()
        [Row(_1=u'Alice', _2=1)]
        >>> df = spark.createDataFrame(rdd, ['name', 'age'])
        >>> df.collect()
        [Row(name=u'Alice', age=1)]

        >>> from pyspark.sql import Row
        >>> Person = Row('name', 'age')
        >>> person = rdd.map(lambda r: Person(*r))
        >>> df2 = spark.createDataFrame(person)
        >>> df2.collect()
        [Row(name=u'Alice', age=1)]

        >>> from pyspark.sql.types import *
        >>> schema = StructType([
        ...    StructField("name", StringType(), True),
        ...    StructField("age", IntegerType(), True)])
        >>> df3 = spark.createDataFrame(rdd, schema)
        >>> df3.collect()
        [Row(name=u'Alice', age=1)]

        >>> spark.createDataFrame(df.toPandas()).collect()  # doctest: +SKIP
        [Row(name=u'Alice', age=1)]
        >>> spark.createDataFrame(pandas.DataFrame([[1, 2]])).collect()  # doctest: +SKIP
        [Row(0=1, 1=2)]

        >>> spark.createDataFrame(rdd, "a: string, b: int").collect()
        [Row(a=u'Alice', b=1)]
        >>> rdd = rdd.map(lambda row: row[1])
        >>> spark.createDataFrame(rdd, "int").collect()
        [Row(value=1)]
        >>> spark.createDataFrame(rdd, "boolean").collect() # doctest: +IGNORE_EXCEPTION_DETAIL
        Traceback (most recent call last):
            ...
        Py4JJavaError: ...
        """
        if isinstance(data, DataFrame):
            raise TypeError("data is already a DataFrame")

        if isinstance(schema, basestring):
            schema = _parse_datatype_string(schema)

        try:
            import pandas
            has_pandas = True
        except Exception:
            has_pandas = False
        if has_pandas and isinstance(data, pandas.DataFrame):
            if schema is None:
                schema = [str(x) for x in data.columns]
            data = [r.tolist() for r in data.to_records(index=False)]

        verify_func = _verify_type if verifySchema else lambda _, t: True
        if isinstance(schema, StructType):
            def prepare(obj):
                verify_func(obj, schema)
                return obj
        elif isinstance(schema, DataType):
            dataType = schema
            schema = StructType().add("value", schema)

            def prepare(obj):
                verify_func(obj, dataType)
                return obj,
        else:
            if isinstance(schema, list):
                schema = [x.encode('utf-8') if not isinstance(x, str) else x for x in schema]
            prepare = lambda obj: obj

        if isinstance(data, RDD):
            rdd, schema = self._createFromRDD(data.map(prepare), schema, samplingRatio)
        else:
            rdd, schema = self._createFromLocal(map(prepare, data), schema)
        jrdd = self._jvm.SerDeUtil.toJavaArray(rdd._to_java_object_rdd())
        jdf = self._jsparkSession.applySchemaToPythonRDD(jrdd.rdd(), schema.json())
        df = DataFrame(jdf, self._wrapped)
        df._schema = schema
        return df

    @ignore_unicode_prefix
    @since(2.0)
    def sql(self, sqlQuery):
        """Returns a :class:`DataFrame` representing the result of the given query.

        :return: :class:`DataFrame`

        >>> df.createOrReplaceTempView("table1")
        >>> df2 = spark.sql("SELECT field1 AS f1, field2 as f2 from table1")
        >>> df2.collect()
        [Row(f1=1, f2=u'row1'), Row(f1=2, f2=u'row2'), Row(f1=3, f2=u'row3')]
        """
        return DataFrame(self._jsparkSession.sql(sqlQuery), self._wrapped)

    @since(2.0)
    def table(self, tableName):
        """Returns the specified table as a :class:`DataFrame`.

        :return: :class:`DataFrame`

        >>> df.createOrReplaceTempView("table1")
        >>> df2 = spark.table("table1")
        >>> sorted(df.collect()) == sorted(df2.collect())
        True
        """
        return DataFrame(self._jsparkSession.table(tableName), self._wrapped)

    @property
    @since(2.0)
    def read(self):
        """
        Returns a :class:`DataFrameReader` that can be used to read data
        in as a :class:`DataFrame`.

        :return: :class:`DataFrameReader`
        """
        return DataFrameReader(self._wrapped)

    @property
    @since(2.0)
    def readStream(self):
        """
        Returns a :class:`DataStreamReader` that can be used to read data streams
        as a streaming :class:`DataFrame`.

        .. note:: Evolving.

        :return: :class:`DataStreamReader`
        """
        return DataStreamReader(self._wrapped)

    @property
    @since(2.0)
    def streams(self):
        """Returns a :class:`StreamingQueryManager` that allows managing all the
        :class:`StreamingQuery` StreamingQueries active on `this` context.

        .. note:: Evolving.

        :return: :class:`StreamingQueryManager`
        """
        from pyspark.sql.streaming import StreamingQueryManager
        return StreamingQueryManager(self._jsparkSession.streams())

    @since(2.0)
    def stop(self):
        """Stop the underlying :class:`SparkContext`.
        """
        self._sc.stop()
        SparkSession._instantiatedSession = None

    @since(2.0)
    def __enter__(self):
        """
        Enable 'with SparkSession.builder.(...).getOrCreate() as session: app' syntax.
        """
        return self

    @since(2.0)
    def __exit__(self, exc_type, exc_val, exc_tb):
        """
        Enable 'with SparkSession.builder.(...).getOrCreate() as session: app' syntax.

        Specifically stop the SparkSession on exit of the with block.
        """
        self.stop()

def _test():
    import os
    import doctest
    from pyspark.context import SparkContext
    from pyspark.sql import Row
    import pyspark.sql.session

    os.chdir(os.environ["SPARK_HOME"])

    globs = pyspark.sql.session.__dict__.copy()
    sc = SparkContext('local[4]', 'PythonTest')
    globs['sc'] = sc
    globs['spark'] = SparkSession(sc)
    globs['rdd'] = rdd = sc.parallelize(
        [Row(field1=1, field2="row1"),
         Row(field1=2, field2="row2"),
         Row(field1=3, field2="row3")])
    globs['df'] = rdd.toDF()
    (failure_count, test_count) = doctest.testmod(
        pyspark.sql.session, globs=globs,
        optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE)
    globs['sc'].stop()
    if failure_count:
        exit(-1)

if __name__ == "__main__":
    _test()


"""
==========
Kernel PCA
==========

This example shows that Kernel PCA is able to find a projection of the data
that makes data linearly separable.
"""
print(__doc__)

# Authors: Mathieu Blondel
#          Andreas Mueller
# License: BSD 3 clause

import numpy as np
import matplotlib.pyplot as plt

from sklearn.decomposition import PCA, KernelPCA
from sklearn.datasets import make_circles

np.random.seed(0)

X, y = make_circles(n_samples=400, factor=.3, noise=.05)

kpca = KernelPCA(kernel="rbf", fit_inverse_transform=True, gamma=10)
X_kpca = kpca.fit_transform(X)
X_back = kpca.inverse_transform(X_kpca)
pca = PCA()
X_pca = pca.fit_transform(X)

# Plot results

plt.figure()
plt.subplot(2, 2, 1, aspect='equal')
plt.title("Original space")
reds = y == 0
blues = y == 1

plt.plot(X[reds, 0], X[reds, 1], "ro")
plt.plot(X[blues, 0], X[blues, 1], "bo")
plt.xlabel("$x_1$")
plt.ylabel("$x_2$")

X1, X2 = np.meshgrid(np.linspace(-1.5, 1.5, 50), np.linspace(-1.5, 1.5, 50))
X_grid = np.array([np.ravel(X1), np.ravel(X2)]).T
# projection on the first principal component (in the phi space)
Z_grid = kpca.transform(X_grid)[:, 0].reshape(X1.shape)
plt.contour(X1, X2, Z_grid, colors='grey', linewidths=1, origin='lower')

plt.subplot(2, 2, 2, aspect='equal')
plt.plot(X_pca[reds, 0], X_pca[reds, 1], "ro")
plt.plot(X_pca[blues, 0], X_pca[blues, 1], "bo")
plt.title("Projection by PCA")
plt.xlabel("1st principal component")
plt.ylabel("2nd component")

plt.subplot(2, 2, 3, aspect='equal')
plt.plot(X_kpca[reds, 0], X_kpca[reds, 1], "ro")
plt.plot(X_kpca[blues, 0], X_kpca[blues, 1], "bo")
plt.title("Projection by KPCA")
plt.xlabel("1st principal component in space induced by $\phi$")
plt.ylabel("2nd component")

plt.subplot(2, 2, 4, aspect='equal')
plt.plot(X_back[reds, 0], X_back[reds, 1], "ro")
plt.plot(X_back[blues, 0], X_back[blues, 1], "bo")
plt.title("Original space after inverse transform")
plt.xlabel("$x_1$")
plt.ylabel("$x_2$")

plt.subplots_adjust(0.02, 0.10, 0.98, 0.94, 0.04, 0.35)

plt.show()


"""
StructureFactorConstraints contains classes for all constraints related experimental static structure factor functions.

.. inheritance-diagram:: fullrmc.Constraints.StructureFactorConstraints
    :parts: 1
"""
# standard libraries imports
from __future__ import print_function
import itertools, re

# external libraries imports
import numpy as np
from pdbparser.Utilities.Database import is_element_property, get_element_property
from pdbparser.Utilities.Collection import get_normalized_weighting

# fullrmc imports
from ..Globals import INT_TYPE, FLOAT_TYPE, PI, PRECISION, LOGGER
from ..Globals import str, long, unicode, bytes, basestring, range, xrange, maxint
from ..Core.Collection import is_number, is_integer, get_path
from ..Core.Collection import reset_if_collected_out_of_date, get_real_elements_weight
from ..Core.Collection import get_caller_frames
from ..Core.Constraint import Constraint, ExperimentalConstraint
from ..Core.pairs_histograms import multiple_pairs_histograms_coords, full_pairs_histograms_coords

class StructureFactorConstraint(ExperimentalConstraint):
    """
    Controls the Structure Factor noted as S(Q) and also called
    total-scattering structure function or Static Structure Factor.
    S(Q) is a dimensionless quantity and normalized such as the average
    value :math:`<S(Q)>=1`.

    It is worth mentioning that S(Q) is nothing other than the normalized and
    corrected diffraction pattern if all experimental artefacts powder.

    The computation of S(Q) is done through an inverse Sine Fourier transform
    of the computed pair distribution function G(r).

    .. math::

        S(Q) = 1+ \\frac{1}{Q} \\int_{0}^{\\infty} G(r) sin(Qr) dr

    From an atomistic model and histogram point of view, G(r) is computed as
    the following:

    .. math::

        G(r) = 4 \\pi r (\\rho_{r} - \\rho_{0})
             = 4 \\pi \\rho_{0} r (g(r)-1)
             = \\frac{R(r)}{r} - 4 \\pi \\rho_{0}

    g(r) is calculated after binning all pair atomic distances into a
    weighted histograms as the following:

    .. math::
        g(r) = \\sum \\limits_{i,j}^{N} w_{i,j} \\frac{\\rho_{i,j}(r)}{\\rho_{0}}
             = \\sum \\limits_{i,j}^{N} w_{i,j} \\frac{n_{i,j}(r) / v(r)}{N_{i,j} / V}

    Where:\n
    :math:`Q` is the momentum transfer. \n
    :math:`r` is the distance between two atoms. \n
    :math:`\\rho_{i,j}(r)` is the pair density function of atoms i and j. \n
    :math:`\\rho_{0}` is the  average number density of the system. \n
    :math:`w_{i,j}` is the relative weighting of atom types i and j. \n
    :math:`R(r)` is the radial distribution function (rdf). \n
    :math:`N` is the total number of atoms. \n
    :math:`V` is the volume of the system. \n
    :math:`n_{i,j}(r)` is the number of atoms i neighbouring j at a distance r. \n
    :math:`v(r)` is the annulus volume at distance r and of thickness dr. \n
    :math:`N_{i,j}` is the total number of atoms i and j in the system. \n

    +----------------------------------------------------------------------+
    |.. figure:: reduced_structure_factor_constraint_plot_method.png       |
    |   :width: 530px                                                      |
    |   :height: 400px                                                     |
    |   :align: left                                                       |
    |                                                                      |
    |   Reduced structure factor of memory shape Nickel-Titanium alloy.    |
    +----------------------------------------------------------------------+

    :Parameters:
        #. experimentalData (numpy.ndarray, string): Experimental data as
           numpy.ndarray or string path to load data using numpy.loadtxt
           method.
        #. dataWeights (None, numpy.ndarray): Weights array of the same number
           of points of experimentalData used in the constraint's standard
           error computation. Therefore particular fitting emphasis can be
           put on different data points that might be considered as more or less
           important in order to get a reasonable and plausible modal.\n
           If None is given, all data points are considered of the same
           importance in the computation of the constraint's standard error.\n
           If numpy.ndarray is given, all weights must be positive and all
           zeros weighted data points won't contribute to the total
           constraint's standard error. At least a single weight point is
           required to be non-zeros and the weights array will be automatically
           scaled upon setting such as the the sum of all the weights
           is equal to the number of data points.
        #. weighting (string): The elements weighting scheme. It must be any
           atomic attribute (atomicNumber, neutronCohb, neutronIncohb,
           neutronCohXs, neutronIncohXs, atomicWeight, covalentRadius) defined
           in pdbparser database. In case of xrays or neutrons experimental
           weights, one can simply set weighting to 'xrays' or 'neutrons'
           and the value will be automatically adjusted to respectively
           'atomicNumber' and 'neutronCohb'. If attribute values are
           missing in the pdbparser database, atomic weights must be
           given in atomsWeight dictionary argument.
        #. atomsWeight (None, dict): Atoms weight dictionary where keys are
           atoms element and values are custom weights. If None is given
           or partially given, missing elements weighting will be fully set
           given weighting scheme.
        #. rmin (None, number): The minimum distance value to compute G(r)
           histogram. If None is given, rmin is computed as
           :math:`2 \\pi / Q_{max}`.
        #. rmax (None, number): The maximum distance value to compute G(r)
           histogram. If None is given, rmax is computed as
           :math:`2 \\pi / dQ`.
        #. dr (None, number): The distance bin value to compute G(r)
           histogram. If None is given, bin is computed as
           :math:`2 \\pi / (Q_{max}-Q_{min})`.
        #. scaleFactor (number): A normalization scale factor used to normalize
           the computed data to the experimental ones.
        #. adjustScaleFactor (list, tuple): Used to adjust fit or guess
           the best scale factor during stochastic engine runtime.
           It must be a list of exactly three entries.\n
           #. The frequency in number of generated moves of finding the best
              scale factor. If 0 frequency is given, it means that the scale
              factor is fixed.
           #. The minimum allowed scale factor value.
           #. The maximum allowed scale factor value.
        #. windowFunction (None, numpy.ndarray): The window function to
           convolute with the computed pair distribution function of the
           system prior to comparing it with the experimental data. In
           general, the experimental pair distribution function G(r) shows
           artificial wrinkles, among others the main reason is because
           G(r) is computed by applying a sine Fourier transform to the
           experimental structure factor S(q). Therefore window function is
           used to best imitate the numerical artefacts in the experimental
           data.
        #. limits (None, tuple, list): The distance limits to compute the
           histograms. If None is given, the limits will be automatically
           set the the min and max distance of the experimental data.
           Otherwise, a tuple of exactly two items where the first is the
           minimum distance or None and the second is the maximum distance
           or None.

    **NB**: If adjustScaleFactor first item (frequency) is 0, the scale factor
    will remain untouched and the limits minimum and maximum won't be checked.

    .. code-block:: python

        # import fullrmc modules
        from fullrmc.Engine import Engine
        from fullrmc.Constraints.StructureFactorConstraints import StructureFactorConstraint

        # create engine
        ENGINE = Engine(path='my_engine.rmc')

        # set pdb file
        ENGINE.set_pdb('system.pdb')

        # create and add constraint
        SFC = StructureFactorConstraint(experimentalData="sq.dat", weighting="atomicNumber")
        ENGINE.add_constraints(SFC)

    """
    def __init__(self, experimentalData, dataWeights=None,
                       weighting="atomicNumber", atomsWeight=None,
                       rmin=None, rmax=None, dr=None,
                       scaleFactor=1.0, adjustScaleFactor=(0, 0.8, 1.2),
                       windowFunction=None, limits=None):
        # initialize variables
        self.__experimentalQValues = None
        self.__experimentalSF      = None
        self.__rmin                = None
        self.__rmax                = None
        self.__dr                  = None
        self.__minimumDistance     = None
        self.__maximumDistance     = None
        self.__bin                 = None
        self.__shellCenters        = None
        self.__histogramSize       = None
        self.__shellVolumes        = None
        self.__Gr2SqMatrix         = None
        # initialize constraint
        super(StructureFactorConstraint, self).__init__( experimentalData=experimentalData, dataWeights=dataWeights, scaleFactor=scaleFactor, adjustScaleFactor=adjustScaleFactor)
        # set atomsWeight
        self.set_atoms_weight(atomsWeight)
        # set elements weighting
        self.set_weighting(weighting)
        self.__set_weighting_scheme()
        # set window function
        self.set_window_function(windowFunction)
        # set r parameters
        self.set_rmin(rmin)
        self.set_rmax(rmax)
        self.set_dr(dr)

        # set frame data
        FRAME_DATA = [d for d in self.FRAME_DATA]
        FRAME_DATA.extend(['_StructureFactorConstraint__experimentalQValues',
                           '_StructureFactorConstraint__experimentalSF',
                           '_StructureFactorConstraint__elementsPairs',
                           '_StructureFactorConstraint__weightingScheme',
                           '_StructureFactorConstraint__atomsWeight',
                           '_StructureFactorConstraint__qmin',
                           '_StructureFactorConstraint__qmax',
                           '_StructureFactorConstraint__rmin',
                           '_StructureFactorConstraint__rmax',
                           '_StructureFactorConstraint__dr',
                           '_StructureFactorConstraint__minimumDistance',
                           '_StructureFactorConstraint__maximumDistance',
                           '_StructureFactorConstraint__bin',
                           '_StructureFactorConstraint__shellCenters',
                           '_StructureFactorConstraint__histogramSize',
                           '_StructureFactorConstraint__shellVolumes',
                           '_StructureFactorConstraint__Gr2SqMatrix',
                           '_StructureFactorConstraint__windowFunction',
                           '_elementsWeight',] )
        RUNTIME_DATA = [d for d in self.RUNTIME_DATA]
        RUNTIME_DATA.extend( [] )
        object.__setattr__(self, 'FRAME_DATA',   tuple(FRAME_DATA)   )
        object.__setattr__(self, 'RUNTIME_DATA', tuple(RUNTIME_DATA) )

    def _codify_update__(self, name='constraint', addDependencies=True):
        dependencies = []
        code         = []
        if addDependencies:
            code.extend(dependencies)
        dw = self.dataWeights
        if dw is not None:
            dw = list(dw)
        code.append("dw = {dw}".format(dw=dw))
        wf = self.windowFunction
        if isinstance(wf, np.ndarray):
            code.append("wf = np.array({wf})".format(wf=list(wf)))
        else:
            code.append("wf = {wf}".format(wf=wf))

        code.append("{name}.set_used({val})".format(name=name, val=self.used))
        code.append("{name}.set_scale_factor({val})".format(name=name, val=self.scaleFactor))
        code.append("{name}.set_adjust_scale_factor({val})".format(name=name, val=self.adjustScaleFactor))
        code.append("{name}.set_data_weights(dw)".format(name=name))
        code.append("{name}.set_atoms_weight({val})".format(name=name, val=self.atomsWeight))
        code.append("{name}.set_window_function(wf)".format(name=name))
        code.append("{name}.set_rmin({val})".format(name=name, val=self.rmin))
        code.append("{name}.set_rmax({val})".format(name=name, val=self.rmax))
        code.append("{name}.set_dr({val})".format(name=name, val=self.dr))
        code.append("{name}.set_limits({val})".format(name=name, val=self.limits))
        # return
        return dependencies, '\n'.join(code)

    def _codify__(self, engine, name='constraint', addDependencies=True):
        assert isinstance(name, basestring), LOGGER.error("name must be a string")
        assert re.match('[a-zA-Z_][a-zA-Z0-9_]*$', name) is not None, LOGGER.error("given name '%s' can't be used as a variable name"%name)
        klass        = self.__class__.__name__
        dependencies = ['import numpy as np','from fullrmc.Constraints import StructureFactorConstraints']
        code         = []
        if addDependencies:
            code.extend(dependencies)
        x = list(self.experimentalData[:,0])
        y = list(self.experimentalData[:,1])
        code.append("x = {x}".format(x=x))
        code.append("y = {y}".format(y=y))
        code.append("d = np.transpose([x,y]).astype(np.float32)")
        dw = self.dataWeights
        if dw is not None:
            dw = list(dw)
        code.append("dw = {dw}".format(dw=dw))
        wf = self.windowFunction
        if isinstance(wf, np.ndarray):
            code.append("wf = np.array({wf})".format(wf=list(wf)))
        else:
            code.append("wf = {wf}".format(wf=wf))
        code.append("{name} = {klass}s.{klass}\
(experimentalData=d, dataWeights=dw, weighting='{weighting}', atomsWeight={atomsWeight}, \
rmin={rmin}, rmax={rmax}, dr={dr}, scaleFactor={scaleFactor}, adjustScaleFactor={adjustScaleFactor}, \
shapeFuncParams=sfp, windowFunction=wf, limits={limits})".format(name=name, klass=klass,
                weighting=self.weighting, atomsWeight=self.atomsWeight, rmin=self.rmin,
                rmax=self.rmax, dr=self.dr, scaleFactor=self.scaleFactor,
                adjustScaleFactor=self.adjustScaleFactor, limits=self.limits))
        code.append("{engine}.add_constraints([{name}])".format(engine=engine, name=name))
        # return
        return dependencies, '\n'.join(code)

    #def __getstate__(self):
    #    # make sure that __Gr2SqMatrix is not pickled but saved to the disk as None
    #    state = super(StructureFactorConstraint, self).__getstate__()
    #    state["_StructureFactorConstraint__Gr2SqMatrix"] = None
    #    return state
    #
    #def __setstate__(self, state):
    #    # make sure to regenerate G(r) to S(q) matrix at loading time
    #    self.__dict__.update( state )
    #    self.__set_Gr_2_Sq_matrix()
    #

    def __set_Gr_2_Sq_matrix(self):
        if self.__experimentalQValues is None or self.__shellCenters is None:
            self.__Gr2SqMatrix = None
        else:
            Qs = self.__experimentalQValues
            Rs = self.__shellCenters
            dr = self.__shellCenters[1]-self.__shellCenters[0]
            qr = Rs.reshape((-1,1))*(np.ones((len(Rs),1), dtype=FLOAT_TYPE)*Qs)
            sinqr = np.sin(qr)
            sinqr_q = sinqr/Qs
            self.__Gr2SqMatrix = dr*sinqr_q

    def __set_weighting_scheme(self):
        if self.engine is not None:
            self.__elementsPairs   = sorted(itertools.combinations_with_replacement(self.engine.elements,2))
            #elementsWeight        = dict([(el,float(get_element_property(el,self.__weighting))) for el in self.engine.elements])
            #self._elementsWeight  = dict([(el,self.__atomsWeight.get(el, float(get_element_property(el,self.__weighting)))) for el in self.engine.elements])
            self._elementsWeight   = get_real_elements_weight(elements=self.engine.elements, weightsDict=self.__atomsWeight, weighting=self.__weighting)
            self.__weightingScheme = get_normalized_weighting(numbers=self.engine.numberOfAtomsPerElement, weights=self._elementsWeight)
            for k in self.__weightingScheme:
                self.__weightingScheme[k] = FLOAT_TYPE(self.__weightingScheme[k])
        else:
            self.__elementsPairs   = None
            self.__weightingScheme = None
        # dump to repository
        self._dump_to_repository({'_StructureFactorConstraint__elementsPairs'  : self.__elementsPairs,
                                  '_StructureFactorConstraint__weightingScheme': self.__weightingScheme})

    def __set_histogram(self):
        if self.__minimumDistance is None or self.__maximumDistance is None or self.__bin is None:
            self.__shellCenters  = None
            self.__histogramSize = None
            self.__shellVolumes  = None
        else:
            # compute edges
            if self.engine is not None and self.rmax is None:
                minHalfBox = np.min( [np.linalg.norm(v)/2. for v in self.engine.basisVectors])
                self.__edges = np.arange(self.__minimumDistance,minHalfBox, self.__bin).astype(FLOAT_TYPE)
            else:
                self.__edges = np.arange(self.__minimumDistance, self.__maximumDistance+self.__bin, self.__bin).astype(FLOAT_TYPE)
            # adjust rmin and rmax
            self.__minimumDistance  = self.__edges[0]
            self.__maximumDistance  = self.__edges[-1]
            # compute shellCenters
            self.__shellCenters = (self.__edges[0:-1]+self.__edges[1:])/FLOAT_TYPE(2.)
            # set histogram size
            self.__histogramSize = INT_TYPE( len(self.__edges)-1 )
            # set shell centers and volumes
            self.__shellVolumes = FLOAT_TYPE(4.0/3.)*PI*((self.__edges[1:])**3 - self.__edges[0:-1]**3)
        # dump to repository
        self._dump_to_repository({'_StructureFactorConstraint__minimumDistance': self.__minimumDistance,
                                  '_StructureFactorConstraint__maximumDistance': self.__maximumDistance,
                                  '_StructureFactorConstraint__shellCenters'   : self.__shellCenters,
                                  '_StructureFactorConstraint__histogramSize'  : self.__histogramSize,
                                  '_StructureFactorConstraint__shellVolumes'   : self.__shellVolumes})
        # reset constraint
        self.reset_constraint()
        # reset sq matrix
        self.__set_Gr_2_Sq_matrix()

    def _on_collector_reset(self):
        pass

    @property
    def rmin(self):
        """ Histogram minimum distance. """
        return self.__rmin

    @property
    def rmax(self):
        """ Histogram maximum distance. """
        return self.__rmax

    @property
    def dr(self):
        """ Histogram bin size."""
        return self.__dr

    @property
    def bin(self):
        """ Computed histogram distance bin size."""
        return self.__bin

    @property
    def minimumDistance(self):
        """ Computed histogram minimum distance. """
        return self.__minimumDistance

    @property
    def maximumDistance(self):
        """ Computed histogram maximum distance. """
        return self.__maximumDistance

    @property
    def qmin(self):
        """ Experimental data reciprocal distances minimum. """
        return self.__qmin

    @property
    def qmax(self):
        """ Experimental data reciprocal distances maximum. """
        return self.__qmax

    @property
    def dq(self):
        """ Experimental data reciprocal distances bin size. """
        return self.__experimentalQValues[1]-self.__experimentalQValues[0]

    @property
    def experimentalQValues(self):
        """ Experimental data used q values. """
        return self.__experimentalQValues

    @property
    def histogramSize(self):
        """ Histogram size"""
        return self.__histogramSize

    @property
    def shellCenters(self):
        """ Shells center array"""
        return self.__shellCenters

    @property
    def shellVolumes(self):
        """ Shells volume array"""
        return self.__shellVolumes

    @property
    def experimentalSF(self):
        """ Experimental Structure Factor or S(q)"""
        return self.__experimentalSF

    @property
    def elementsPairs(self):
        """ Elements pairs """
        return self.__elementsPairs

    @property
    def atomsWeight(self):
        """Custom atoms weight"""
        return self.__atomsWeight

    @property
    def weighting(self):
        """ Elements weighting definition. """
        return self.__weighting

    @property
    def weightingScheme(self):
        """ Elements weighting scheme. """
        return self.__weightingScheme

    @property
    def windowFunction(self):
        """ Convolution window function. """
        return self.__windowFunction

    @property
    def Gr2SqMatrix(self):
        """ G(r) to S(q) transformation matrix."""
        return self.__Gr2SqMatrix

    @property
    def _experimentalX(self):
        """For internal use only to interface
        ExperimentalConstraint.get_constraints_properties"""
        return self.__experimentalQValues

    @property
    def _experimentalY(self):
        """For internal use only to interface
        ExperimentalConstraint.get_constraints_properties"""
        return self.__experimentalSF

    @property
    def _modelX(self):
        """For internal use only to interface
        ExperimentalConstraint.get_constraints_properties"""
        return self.__experimentalQValues

    def listen(self, message, argument=None):
        """
        Listens to any message sent from the Broadcaster.

        :Parameters:
            #. message (object): Any python object to send to constraint's
               listen method.
            #. argument (object): Any type of argument to pass to the
               listeners.
        """
        if message in ("engine set","update pdb","update molecules indexes","update elements indexes","update names indexes"):
            self.__set_weighting_scheme()
            # reset histogram
            if self.engine is not None:
                self.__set_histogram()
            self.reset_constraint() # ADDED 2017-JAN-08
        elif message in("update boundary conditions",):
            self.reset_constraint()

    def set_rmin(self, rmin):
        """
        Set rmin value.

        :parameters:
            #. rmin (None, number): The minimum distance value to compute G(r)
               histogram. If None is given, rmin is computed as
               :math:`2 \\pi / Q_{max}`.
        """
        if rmin is None:
            minimumDistance = FLOAT_TYPE( 2.*PI/self.__qmax )
        else:
            assert is_number(rmin), LOGGER.error("rmin must be None or a number")
            minimumDistance = FLOAT_TYPE(rmin)
        if self.__maximumDistance is not None:
            assert minimumDistance<self.__maximumDistance, LOGGER.error("rmin must be smaller than rmax %s"%self.__maximumDistance)
        self.__rmin = rmin
        self.__minimumDistance = minimumDistance
        # dump to repository
        self._dump_to_repository({'_StructureFactorConstraint__rmin': self.__rmin,
                                  '_StructureFactorConstraint__minimumDistance': self.__minimumDistance})
        # reset histogram
        self.__set_histogram()

    def set_rmax(self, rmax):
        """
        Set rmax value.

        :Parameters:
            #. rmax (None, number): The maximum distance value to compute G(r)
               histogram. If None is given, rmax is computed as
               :math:`2 \\pi / dQ`.
        """
        if rmax is None:
            dq = self.__experimentalQValues[1]-self.__experimentalQValues[0]
            maximumDistance = FLOAT_TYPE( 2.*PI/dq )
        else:
            assert is_number(rmax), LOGGER.error("rmax must be None or a number")
            maximumDistance = FLOAT_TYPE(rmax)
        if self.__minimumDistance is not None:
            assert maximumDistance>self.__minimumDistance, LOGGER.error("rmax must be bigger than rmin %s"%self.__minimumDistance)
        self.__rmax = rmax
        self.__maximumDistance = maximumDistance
        # dump to repository
        self._dump_to_repository({'_StructureFactorConstraint__rmax': self.__rmax,
                                  '_StructureFactorConstraint__maximumDistance': self.__maximumDistance})
        # reset histogram
        self.__set_histogram()

    def set_dr(self, dr):
        """
        Set dr value.

        :Parameters:
            #. dr (None, number): The distance bin value to compute G(r)
               histogram. If None is given, bin is computed as
               :math:`2 \\pi / (Q_{max}-Q_{min})`.
        """
        if dr is None:
            bin  = 2.*PI/self.__qmax
            rbin = round(bin,1)
            if rbin>bin:
                rbin -= 0.1
            bin = FLOAT_TYPE( rbin  )
        else:
            assert is_number(dr), LOGGER.error("dr must be None or a number")
            bin = FLOAT_TYPE(dr)
        self.__dr = dr
        self.__bin = bin
        # dump to repository
        self._dump_to_repository({'_StructureFactorConstraint__dr': self.__dr,
                                  '_StructureFactorConstraint__bin': self.__bin})
        # reset histogram
        self.__set_histogram()

    def set_weighting(self, weighting):
        """
        Set elements weighting. It must be a valid entry of pdbparser atom's
        database.

        :Parameters:
            #. weighting (string): The elements weighting scheme. It must be
               any atomic attribute (atomicNumber, neutronCohb, neutronIncohb,
               neutronCohXs, neutronIncohXs, atomicWeight, covalentRadius)
               defined in pdbparser database. In case of xrays or neutrons
               experimental weights, one can simply set weighting to 'xrays'
               or 'neutrons' and the value will be automatically adjusted to
               respectively 'atomicNumber' and 'neutronCohb'. If attribute
               values are  missing in the pdbparser database, atomic weights
               must be given in atomsWeight dictionary argument.
        """
        if weighting.lower() in ["xrays","x-rays","xray","x-ray"]:
            LOGGER.fixed("'%s' weighting is set to atomicNumber"%weighting)
            weighting = "atomicNumber"
        elif weighting.lower() in ["neutron","neutrons"]:
            LOGGER.fixed("'%s' weighting is set to neutronCohb"%weighting)
            weighting = "neutronCohb"
        assert is_element_property(weighting),LOGGER.error( "weighting is not a valid pdbparser atoms database entry")
        assert weighting != "atomicFormFactor", LOGGER.error("atomicFormFactor weighting is not allowed")
        self.__weighting = weighting
        # dump to repository
        self._dump_to_repository({'_StructureFactorConstraint__weighting': self.__weighting})

    def set_atoms_weight(self, atomsWeight):
        """
        Custom set atoms weight. This is the way to setting a atoms weights
        different than the given weighting scheme.

        :Parameters:
            #. atomsWeight (None, dict): Atoms weight dictionary where keys are
               atoms element and values are custom weights. If None is given
               or partially given, missing elements weighting will be fully set
               given weighting scheme.
        """
        if atomsWeight is None:
            AW = {}
        else:
            assert isinstance(atomsWeight, dict),LOGGER.error("atomsWeight must be None or a dictionary")
            AW = {}
            for k in atomsWeight:
                assert isinstance(k, basestring),LOGGER.error("atomsWeight keys must be strings")
                try:
                    val = float(atomsWeight[k])
                except:
                    raise LOGGER.error( "atomsWeight values must be numerical")
                AW[k]=val
        # set atomsWeight
        self.__atomsWeight = AW
        # dump to repository
        self._dump_to_repository({'_StructureFactorConstraint__atomsWeight': self.__atomsWeight})

    def set_window_function(self, windowFunction):
        """
        Set convolution window function.

        :Parameters:
             #. windowFunction (None, numpy.ndarray): The window function to
                convolute with the computed pair distribution function of the
                system prior to comparing it with the experimental data. In
                general, the experimental pair distribution function G(r) shows
                artificial wrinkles, among others the main reason is because
                G(r) is computed by applying a sine Fourier transform to the
                experimental structure factor S(q). Therefore window function is
                used to best imitate the numerical artefacts in the experimental
                data.
        """
        if windowFunction is not None:
            assert isinstance(windowFunction, np.ndarray), LOGGER.error("windowFunction must be a numpy.ndarray")
            assert windowFunction.dtype.type is FLOAT_TYPE, LOGGER.error("windowFunction type must be %s"%FLOAT_TYPE)
            assert len(windowFunction.shape) == 1, LOGGER.error("windowFunction must be of dimension 1")
            assert len(windowFunction) <= self.experimentalData.shape[0], LOGGER.error("windowFunction length must be smaller than experimental data")
            # normalize window function
            windowFunction /= np.sum(windowFunction)
        # check window size
        # set windowFunction
        self.__windowFunction = windowFunction
        # dump to repository
        self._dump_to_repository({'_StructureFactorConstraint__windowFunction': self.__windowFunction})

    def set_experimental_data(self, experimentalData):
        """
        Set constraint's experimental data.

        :Parameters:
            #. experimentalData (numpy.ndarray, string): The experimental
               data as numpy.ndarray or string path to load data using
               numpy.loadtxt function.
        """
        # get experimental data
        super(StructureFactorConstraint, self).set_experimental_data(experimentalData=experimentalData)
        # set limits
        self.set_limits(self.limits)

    def set_limits(self, limits):
        """
        Set the reciprocal distance limits (qmin, qmax).

        :Parameters:
            #. limits (None, tuple, list): Distance limits to bound
               experimental data and compute histograms.
               If None is given, the limits will be automatically set to
               min and max reciprocal distance recorded in experimental data.
               If given, a tuple of minimum reciprocal distance (qmin) or None
               and maximum reciprocal distance (qmax) or None should be given.
        """
        self._ExperimentalConstraint__set_limits(limits)
        # set qvalues
        self.__experimentalQValues = self.experimentalData[self.limitsIndexStart:self.limitsIndexEnd+1,0].astype(FLOAT_TYPE)
        self.__experimentalSF      = self.experimentalData[self.limitsIndexStart:self.limitsIndexEnd+1,1].astype(FLOAT_TYPE)
        # set qmin and qmax
        self.__qmin = self.__experimentalQValues[0]
        self.__qmax = self.__experimentalQValues[-1]
        assert self.__qmin>0, LOGGER.error("qmin must be bigger than 0. Experimental null q values are ambigous. Try setting limits.")
        # dump to repository
        self._dump_to_repository({'_StructureFactorConstraint__experimentalQValues': self.__experimentalQValues,
                                  '_StructureFactorConstraint__experimentalSF'     : self.__experimentalSF,
                                  '_StructureFactorConstraint__qmin'               : self.__qmin,
                                  '_StructureFactorConstraint__qmax'               : self.__qmax})
        # set used dataWeights
        self._set_used_data_weights(limitsIndexStart=self.limitsIndexStart, limitsIndexEnd=self.limitsIndexEnd)
        # reset constraint
        self.reset_constraint()
        # reset sq matrix
        self.__set_Gr_2_Sq_matrix()

    def update_standard_error(self):
        """ Compute and set constraint's standardError."""
        # set standardError
        totalSQ = self.get_constraint_value()["total_no_window"]
        self.set_standard_error(self.compute_standard_error(modelData = totalSQ))

    def check_experimental_data(self, experimentalData):
        """
        Check whether experimental data is correct.

        :Parameters:
            #. experimentalData (object): The experimental data to check.

        :Returns:
            #. result (boolean): Whether it is correct or not.
            #. message (str): Checking message that explains whats's wrong
               with the given data
        """
        if not isinstance(experimentalData, np.ndarray):
            return False, "experimentalData must be a numpy.ndarray"
        if experimentalData.dtype.type is not FLOAT_TYPE:
            return False, "experimentalData type must be %s"%FLOAT_TYPE
        if len(experimentalData.shape) !=2:
            return False, "experimentalData must be of dimension 2"
        if experimentalData.shape[1] !=2:
            return False, "experimentalData must have only 2 columns"
        # check distances order
        inOrder = (np.array(sorted(experimentalData[:,0]), dtype=FLOAT_TYPE)-experimentalData[:,0])<=PRECISION
        if not np.all(inOrder):
            return False, "experimentalData distances are not sorted in order"
        if experimentalData[0][0]<0:
            return False, "experimentalData distances min value is found negative"
        # data format is correct
        return True, ""

    def compute_standard_error(self, modelData):
        """
        Compute the standard error (StdErr) as the squared deviations
        between model computed data and the experimental ones.

        .. math::
            StdErr = \\sum \\limits_{i}^{N} W_{i}(Y(X_{i})-F(X_{i}))^{2}

        Where:\n
        :math:`N` is the total number of experimental data points. \n
        :math:`W_{i}` is the data point weight. It becomes equivalent to 1 when dataWeights is set to None. \n
        :math:`Y(X_{i})` is the experimental data point :math:`X_{i}`. \n
        :math:`F(X_{i})` is the computed from the model data  :math:`X_{i}`. \n

        :Parameters:
            #. modelData (numpy.ndarray): The data to compare with the
               experimental one and compute the squared deviation.

        :Returns:
            #. standardError (number): The calculated constraint's
               standardError.
        """
        # compute difference
        diff = self.__experimentalSF-modelData
        # return standard error
        if self._usedDataWeights is None:
            return np.add.reduce((diff)**2)
        else:
            return np.add.reduce(self._usedDataWeights*((diff)**2))

    def _get_Sq_from_Gr(self, Gr):
        return np.sum(Gr.reshape((-1,1))*self.__Gr2SqMatrix, axis=0)+1

    def _apply_scale_factor(self, Sq, scaleFactor):
        if scaleFactor != 1:
            Sq = scaleFactor*(Sq-1) + 1
        return Sq

    def __get_total_Sq(self, data, rho0):
        """This method is created just to speed up the computation of
        the total Sq upon fitting."""
        Gr = np.zeros(self.__histogramSize, dtype=FLOAT_TYPE)
        for pair in self.__elementsPairs:
            # get weighting scheme
            wij = self.__weightingScheme.get(pair[0]+"-"+pair[1], None)
            if wij is None:
                wij = self.__weightingScheme[pair[1]+"-"+pair[0]]
            # get number of atoms per element
            ni = self.engine.numberOfAtomsPerElement[pair[0]]
            nj = self.engine.numberOfAtomsPerElement[pair[1]]
            # get index of element
            idi = self.engine.elements.index(pair[0])
            idj = self.engine.elements.index(pair[1])
            # get Nij
            if idi == idj:
                Nij = ni*(ni-1)/2.0
                Dij = Nij/self.engine.volume
                nij = data["intra"][idi,idj,:]+data["inter"][idi,idj,:]
                Gr += wij*nij/Dij
            else:
                Nij = ni*nj
                Dij = Nij/self.engine.volume
                nij = data["intra"][idi,idj,:]+data["intra"][idj,idi,:] + data["inter"][idi,idj,:]+data["inter"][idj,idi,:]
                Gr += wij*nij/Dij
        # Devide by shells volume
        Gr /= self.shellVolumes
        # compute total G(r)
        #rho0 = (self.engine.numberOfAtoms/self.engine.volume).astype(FLOAT_TYPE)
        Gr   = (FLOAT_TYPE(4.)*PI*self.__shellCenters*rho0)*(Gr-1)
        # Compute S(q) from G(r)
        Sq = self._get_Sq_from_Gr(Gr)
        # Multiply by scale factor
        self._fittedScaleFactor = self.get_adjusted_scale_factor(self.__experimentalSF, Sq, self._usedDataWeights)
        # apply scale factor
        Sq = self._apply_scale_factor(Sq, self._fittedScaleFactor)
        # apply multiframe prior and weight
        Sq = self._apply_multiframe_prior(Sq)
        # convolve total with window function
        if self.__windowFunction is not None:
            Sq = np.convolve(Sq, self.__windowFunction, 'same')
        return Sq

    def get_adjusted_scale_factor(self, experimentalData, modelData, dataWeights):
        """Overload to reduce S(q) prior to fitting scale factor.
        S(q) -> 1 at high q and this will create a wrong scale factor.
        Overloading can be avoided but it's better to for performance reasons
        """
        SF = self.scaleFactor
        # check to update scaleFactor
        if self.adjustScaleFactorFrequency:
            if not self.engine.accepted%self.adjustScaleFactorFrequency:
                SF = self.fit_scale_factor(experimentalData-1, modelData-1, dataWeights)
        return SF

    def _get_constraint_value(self, data, applyMultiframePrior=True):
        # http://erice2011.docking.org/upload/Other/Billinge_PDF/03-ReadingMaterial/BillingePDF2011.pdf    page 6
        #import time
        #startTime = time.clock()
        output = {}
        for pair in self.__elementsPairs:
            output["sf_intra_%s-%s" % pair] = np.zeros(self.__histogramSize, dtype=FLOAT_TYPE)
            output["sf_inter_%s-%s" % pair] = np.zeros(self.__histogramSize, dtype=FLOAT_TYPE)
            output["sf_total_%s-%s" % pair] = np.zeros(self.__histogramSize, dtype=FLOAT_TYPE)
        gr = np.zeros(self.__histogramSize, dtype=FLOAT_TYPE)
        for pair in self.__elementsPairs:
            # get weighting scheme
            wij = self.__weightingScheme.get(pair[0]+"-"+pair[1], None)
            if wij is None:
                wij = self.__weightingScheme[pair[1]+"-"+pair[0]]
            # get number of atoms per element
            ni = self.engine.numberOfAtomsPerElement[pair[0]]
            nj = self.engine.numberOfAtomsPerElement[pair[1]]
            # get index of element
            idi = self.engine.elements.index(pair[0])
            idj = self.engine.elements.index(pair[1])
            # get Nij
            if idi == idj:
                Nij = ni*(ni-1)/2.0
                output["sf_intra_%s-%s" % pair] += data["intra"][idi,idj,:]
                output["sf_inter_%s-%s" % pair] += data["inter"][idi,idj,:]
            else:
                Nij = ni*nj
                output["sf_intra_%s-%s" % pair] += data["intra"][idi,idj,:] + data["intra"][idj,idi,:]
                output["sf_inter_%s-%s" % pair] += data["inter"][idi,idj,:] + data["inter"][idj,idi,:]
            # compute g(r)
            nij = output["sf_intra_%s-%s" % pair] + output["sf_inter_%s-%s" % pair]
            dij = nij/self.__shellVolumes
            Dij = Nij/self.engine.volume
            gr += wij*dij/Dij
            # calculate intensityFactor
            intensityFactor = (self.engine.volume*wij)/(Nij*self.__shellVolumes)
            # divide by factor
            output["sf_intra_%s-%s" % pair] *= intensityFactor
            output["sf_inter_%s-%s" % pair] *= intensityFactor
            output["sf_total_%s-%s" % pair]  = output["sf_intra_%s-%s" % pair] + output["sf_inter_%s-%s" % pair]
            # Compute S(q) from G(r)
            output["sf_intra_%s-%s" % pair] = self._get_Sq_from_Gr(output["sf_intra_%s-%s" % pair])
            output["sf_inter_%s-%s" % pair] = self._get_Sq_from_Gr(output["sf_inter_%s-%s" % pair])
            output["sf_total_%s-%s" % pair] = self._get_Sq_from_Gr(output["sf_total_%s-%s" % pair])
        # compute total G(r)
        rho0 = (self.engine.numberOfAtoms/self.engine.volume).astype(FLOAT_TYPE)
        Gr = (FLOAT_TYPE(4.)*PI*self.__shellCenters*rho0) * (gr-1)
        # Compute S(q) from G(r)
        Sq = self._get_Sq_from_Gr(Gr)
        # multiply by scale factor
        output["total_no_window"] = self._apply_scale_factor(Sq, self._fittedScaleFactor)
        # apply multiframe prior and weight
        if applyMultiframePrior:
            output["total_no_window"] = self._apply_multiframe_prior(output["total_no_window"])
        # convolve total with window function
        if self.__windowFunction is not None:
            output["total"] = np.convolve(output["total_no_window"], self.__windowFunction, 'same').astype(FLOAT_TYPE)
        else:
            output["total"] = output["total_no_window"]
        return output

    def get_constraint_value(self, applyMultiframePrior=True):
        """
        Compute all partial Structure Factor (SQs).

        :Parameters:
            #. applyMultiframePrior (boolean): Whether to apply subframe weight
               and prior to the total. This will only have an effect when used
               frame is a subframe and in case subframe weight and prior is
               defined.

        :Returns:
            #. SQs (dictionary): The SQs dictionnary, where keys are the
               element wise intra and inter molecular SQs and values are
               the computed SQs.
        """
        if self.data is None:
            LOGGER.warn("data must be computed first using 'compute_data' method.")
            return {}
        return self._get_constraint_value(self.data, applyMultiframePrior=applyMultiframePrior)

    def get_constraint_original_value(self):
        """
        Compute all partial Pair Distribution Functions (PDFs).

        :Returns:
            #. PDFs (dictionary): The PDFs dictionnary, where keys are the
               element wise intra and inter molecular PDFs and values are the
               computed PDFs.
        """
        if self.originalData is None:
            LOGGER.warn("originalData must be computed first using 'compute_data' method.")
            return {}
        return self._get_constraint_value(self.originalData)

    @reset_if_collected_out_of_date
    def compute_data(self, update=True):
        """ Compute constraint's data.

        :Parameters:
            #. update (boolean): whether to update constraint data and
               standard error with new computation. If data is computed and
               updated by another thread or process while the stochastic
               engine is running, this might lead to a state alteration of
               the constraint which will lead to a no additional accepted
               moves in the run

        :Returns:
            #. data (dict): constraint data dictionary
            #. standardError (float): constraint standard error
        """
        intra,inter = full_pairs_histograms_coords( boxCoords        = self.engine.boxCoordinates,
                                                    basis            = self.engine.basisVectors,
                                                    isPBC            = self.engine.isPBC,
                                                    moleculeIndex    = self.engine.moleculesIndex,
                                                    elementIndex     = self.engine.elementsIndex,
                                                    numberOfElements = self.engine.numberOfElements,
                                                    minDistance      = self.__minimumDistance,
                                                    maxDistance      = self.__maximumDistance,
                                                    histSize         = self.__histogramSize,
                                                    bin              = self.__bin,
                                                    ncores           = self.engine._runtime_ncores  )
        # create data and compute standard error
        data     = {"intra":intra, "inter":inter}
        totalSQ  = self.__get_total_Sq(data, rho0=self.engine.numberDensity)
        stdError = self.compute_standard_error(modelData = totalSQ)
        # update
        if update:
            self.set_data(data)
            self.set_active_atoms_data_before_move(None)
            self.set_active_atoms_data_after_move(None)
            self.set_standard_error(stdError)
            # set original data
            if self.originalData is None:
                self._set_original_data(self.data)
        # return
        return data, stdError

    def compute_before_move(self, realIndexes, relativeIndexes):
        """
        Compute constraint before move is executed

        :Parameters:
            #. realIndexes (numpy.ndarray): Not used here.
            #. relativeIndexes (numpy.ndarray): Group atoms relative index
               the move will be applied to.
        """
        intraM,interM = multiple_pairs_histograms_coords( indexes          = relativeIndexes,
                                                          boxCoords        = self.engine.boxCoordinates,
                                                          basis            = self.engine.basisVectors,
                                                          isPBC            = self.engine.isPBC,
                                                          moleculeIndex    = self.engine.moleculesIndex,
                                                          elementIndex     = self.engine.elementsIndex,
                                                          numberOfElements = self.engine.numberOfElements,
                                                          minDistance      = self.__minimumDistance,
                                                          maxDistance      = self.__maximumDistance,
                                                          histSize         = self.__histogramSize,
                                                          bin              = self.__bin,
                                                          allAtoms         = True,
                                                          ncores           = self.engine._runtime_ncores )
        intraF,interF = full_pairs_histograms_coords( boxCoords        = self.engine.boxCoordinates[relativeIndexes],
                                                      basis            = self.engine.basisVectors,
                                                      isPBC            = self.engine.isPBC,
                                                      moleculeIndex    = self.engine.moleculesIndex[relativeIndexes],
                                                      elementIndex     = self.engine.elementsIndex[relativeIndexes],
                                                      numberOfElements = self.engine.numberOfElements,
                                                      minDistance      = self.__minimumDistance,
                                                      maxDistance      = self.__maximumDistance,
                                                      histSize         = self.__histogramSize,
                                                      bin              = self.__bin,
                                                      ncores           = self.engine._runtime_ncores )
        self.set_active_atoms_data_before_move( {"intra":intraM-intraF, "inter":interM-interF} )
        self.set_active_atoms_data_after_move(None)

    def compute_after_move(self, realIndexes, relativeIndexes, movedBoxCoordinates):
        """
        Compute constraint after move is executed

        :Parameters:
            #. realIndexes (numpy.ndarray): Not used here.
            #. relativeIndexes (numpy.ndarray): Group atoms relative index
               the move will be applied to.
            #. movedBoxCoordinates (numpy.ndarray): The moved atoms new coordinates.
        """
        # change coordinates temporarily
        boxData = np.array(self.engine.boxCoordinates[relativeIndexes], dtype=FLOAT_TYPE)
        self.engine.boxCoordinates[relativeIndexes] = movedBoxCoordinates
        # calculate pair distribution function
        intraM,interM = multiple_pairs_histograms_coords( indexes          = relativeIndexes,
                                                          boxCoords        = self.engine.boxCoordinates,
                                                          basis            = self.engine.basisVectors,
                                                          isPBC            = self.engine.isPBC,
                                                          moleculeIndex    = self.engine.moleculesIndex,
                                                          elementIndex     = self.engine.elementsIndex,
                                                          numberOfElements = self.engine.numberOfElements,
                                                          minDistance      = self.__minimumDistance,
                                                          maxDistance      = self.__maximumDistance,
                                                          histSize         = self.__histogramSize,
                                                          bin              = self.__bin,
                                                          allAtoms         = True,
                                                          ncores           = self.engine._runtime_ncores )
        intraF,interF = full_pairs_histograms_coords( boxCoords        = self.engine.boxCoordinates[relativeIndexes],
                                                      basis            = self.engine.basisVectors,
                                                      isPBC            = self.engine.isPBC,
                                                      moleculeIndex    = self.engine.moleculesIndex[relativeIndexes],
                                                      elementIndex     = self.engine.elementsIndex[relativeIndexes],
                                                      numberOfElements = self.engine.numberOfElements,
                                                      minDistance      = self.__minimumDistance,
                                                      maxDistance      = self.__maximumDistance,
                                                      histSize         = self.__histogramSize,
                                                      bin              = self.__bin,
                                                      ncores           = self.engine._runtime_ncores  )
        # set active atoms data
        self.set_active_atoms_data_after_move( {"intra":intraM-intraF, "inter":interM-interF} )
        # reset coordinates
        self.engine.boxCoordinates[relativeIndexes] = boxData
        # compute standardError after move
        dataIntra = self.data["intra"]-self.activeAtomsDataBeforeMove["intra"]+self.activeAtomsDataAfterMove["intra"]
        dataInter = self.data["inter"]-self.activeAtomsDataBeforeMove["inter"]+self.activeAtomsDataAfterMove["inter"]
        totalSQ   = self.__get_total_Sq({"intra":dataIntra, "inter":dataInter}, rho0=self.engine.numberDensity)
        self.set_after_move_standard_error( self.compute_standard_error(modelData = totalSQ) )
        # increment tried
        self.increment_tried()

    def accept_move(self, realIndexes, relativeIndexes):
        """
        Accept move

        :Parameters:
            #. realIndexes (numpy.ndarray): Not used here.
            #. relativeIndexes (numpy.ndarray): Not used here.
        """
        dataIntra = self.data["intra"]-self.activeAtomsDataBeforeMove["intra"]+self.activeAtomsDataAfterMove["intra"]
        dataInter = self.data["inter"]-self.activeAtomsDataBeforeMove["inter"]+self.activeAtomsDataAfterMove["inter"]
        # change permanently _data
        self.set_data( {"intra":dataIntra, "inter":dataInter} )
        # reset activeAtoms data
        self.set_active_atoms_data_before_move(None)
        self.set_active_atoms_data_after_move(None)
        # update standardError
        self.set_standard_error( self.afterMoveStandardError )
        self.set_after_move_standard_error( None )
        # set new scale factor
        self._set_fitted_scale_factor_value(self._fittedScaleFactor)
        # increment accepted
        self.increment_accepted()

    def reject_move(self, realIndexes, relativeIndexes):
        """
        Reject move

        :Parameters:
            #. realIndexes (numpy.ndarray): Not used here.
            #. relativeIndexes (numpy.ndarray): Not used here.
        """
        # reset activeAtoms data
        self.set_active_atoms_data_before_move(None)
        self.set_active_atoms_data_after_move(None)
        # update standardError
        self.set_after_move_standard_error( None )

    def compute_as_if_amputated(self, realIndex, relativeIndex):
        """
        Compute and return constraint's data and standard error as if
        given atom is amputated.

        :Parameters:
            #. realIndex (numpy.ndarray): Atom's index as a numpy array
               of a single element.
            #. relativeIndex (numpy.ndarray): Atom's relative index as a
               numpy array of a single element.
        """
        # compute data
        self.compute_before_move(realIndexes=realIndex, relativeIndexes=relativeIndex)
        dataIntra = self.data["intra"]-self.activeAtomsDataBeforeMove["intra"]
        dataInter = self.data["inter"]-self.activeAtomsDataBeforeMove["inter"]
        data      = {"intra":dataIntra, "inter":dataInter}
        # temporarily adjust self.__weightingScheme
        weightingScheme = self.__weightingScheme
        relativeIndex = relativeIndex[0]
        selectedElement = self.engine.allElements[relativeIndex]
        self.engine.numberOfAtomsPerElement[selectedElement] -= 1
        self.__weightingScheme = get_normalized_weighting(numbers=self.engine.numberOfAtomsPerElement, weights=self._elementsWeight )
        for k in self.__weightingScheme:
            self.__weightingScheme[k] = FLOAT_TYPE(self.__weightingScheme[k])
        ## END OF ADDED 08 FEB 2017
        # compute standard error
        if not self.engine._RT_moveGenerator.allowFittingScaleFactor:
            SF = self.adjustScaleFactorFrequency
            self._set_adjust_scale_factor_frequency(0)
        rho0          = ((self.engine.numberOfAtoms-1)/self.engine.volume).astype(FLOAT_TYPE)
        totalSQ       = self.__get_total_Sq(data, rho0=rho0)
        standardError = self.compute_standard_error(modelData = totalSQ)
        if not self.engine._RT_moveGenerator.allowFittingScaleFactor:
            self._set_adjust_scale_factor_frequency(SF)
        # reset activeAtoms data
        self.set_active_atoms_data_before_move(None)
        # set amputation
        self.set_amputation_data( {'data':data, 'weightingScheme':self.__weightingScheme} )
        # compute standard error
        self.set_amputation_standard_error( standardError )
        # reset weightingScheme and number of atoms per element
        self.__weightingScheme = weightingScheme
        self.engine.numberOfAtomsPerElement[selectedElement] += 1

    def accept_amputation(self, realIndex, relativeIndex):
        """
        Accept amputated atom and sets constraints data and standard error accordingly.

        :Parameters:
            #. realIndex (numpy.ndarray): Not used here.
            #. relativeIndex (numpy.ndarray): Not used here.
        """
        #self.set_data( self.amputationData ) ## COMMENTED 08 FEB 2017
        self.set_data( self.amputationData['data'] )
        self.__weightingScheme = self.amputationData['weightingScheme']
        self.set_standard_error( self.amputationStandardError )
        self.set_amputation_data( None )
        self.set_amputation_standard_error( None )
        # set new scale factor
        self._set_fitted_scale_factor_value(self._fittedScaleFactor)

    def reject_amputation(self, realIndex, relativeIndex):
        """
        Reject amputated atom and set constraint's data and standard
        error accordingly.

        :Parameters:
            #. realIndex (numpy.ndarray): Not used here.
            #. relativeIndex (numpy.ndarray): Not used here.
        """
        self.set_amputation_data( None )
        self.set_amputation_standard_error( None )

    def _on_collector_collect_atom(self, realIndex):
        pass

    def _on_collector_release_atom(self, realIndex):
        pass

    def _constraint_copy_needs_lut(self):
        return {'_StructureFactorConstraint__elementsPairs'       :'_StructureFactorConstraint__elementsPairs',
                '_StructureFactorConstraint__histogramSize'       :'_StructureFactorConstraint__histogramSize',
                '_StructureFactorConstraint__weightingScheme'     :'_StructureFactorConstraint__weightingScheme',
                '_StructureFactorConstraint__shellVolumes'        :'_StructureFactorConstraint__shellVolumes',
                '_StructureFactorConstraint__shellCenters'        :'_StructureFactorConstraint__shellCenters',
                '_StructureFactorConstraint__windowFunction'      :'_StructureFactorConstraint__windowFunction',
                '_StructureFactorConstraint__experimentalQValues' :'_StructureFactorConstraint__experimentalQValues',
                '_StructureFactorConstraint__experimentalSF'      :'_StructureFactorConstraint__experimentalSF',
                '_StructureFactorConstraint__Gr2SqMatrix'         :'_StructureFactorConstraint__Gr2SqMatrix',
                '_StructureFactorConstraint__minimumDistance'     :'_StructureFactorConstraint__minimumDistance',
                '_StructureFactorConstraint__maximumDistance'     :'_StructureFactorConstraint__maximumDistance',
                '_StructureFactorConstraint__bin'                 :'_StructureFactorConstraint__bin',
                '_ExperimentalConstraint__scaleFactor'            :'_ExperimentalConstraint__scaleFactor',
                '_ExperimentalConstraint__dataWeights'            :'_ExperimentalConstraint__dataWeights',
                '_ExperimentalConstraint__multiframePrior'        :'_ExperimentalConstraint__multiframePrior',
                '_ExperimentalConstraint__multiframeWeight'       :'_ExperimentalConstraint__multiframeWeight',
                '_ExperimentalConstraint__limits'                 :'_ExperimentalConstraint__limits',
                '_ExperimentalConstraint__limitsIndexStart'       :'_ExperimentalConstraint__limitsIndexStart',
                '_ExperimentalConstraint__limitsIndexEnd'         :'_ExperimentalConstraint__limitsIndexEnd',
                '_Constraint__used'                               :'_Constraint__used',
                '_Constraint__data'                               :'_Constraint__data',
                '_Constraint__state'                              :'_Constraint__state',
                '_Constraint__standardError'                      :'_Constraint__standardError',
                '_fittedScaleFactor'                              :'_fittedScaleFactor',
                '_usedDataWeights'                                :'_usedDataWeights',
                '_Engine__state'                                  :'_Engine__state',
                '_Engine__boxCoordinates'                         :'_Engine__boxCoordinates',
                '_Engine__basisVectors'                           :'_Engine__basisVectors',
                '_Engine__isPBC'                                  :'_Engine__isPBC',
                '_Engine__moleculesIndex'                         :'_Engine__moleculesIndex',
                '_Engine__elementsIndex'                          :'_Engine__elementsIndex',
                '_Engine__numberOfAtomsPerElement'                :'_Engine__numberOfAtomsPerElement',
                '_Engine__elements'                               :'_Engine__elements',
                '_Engine__numberDensity'                          :'_Engine__numberDensity',
                '_Engine__volume'                                 :'_Engine__volume',
                '_Engine__realCoordinates'                        :'_Engine__realCoordinates',
                '_atomsCollector'                                 :'_atomsCollector',
                ('engine','_atomsCollector')                      :'_atomsCollector',
               }

    def plot(self, xlabelParams={'xlabel':'$Q(\\AA^{-1})$', 'size':10},
                   ylabelParams={'ylabel':'$S(Q)$', 'size':10},
                   **kwargs):
        """
        Alias to ExperimentalConstraint.plot with additional parameters

        :Additional/Adjusted Parameters:
            #. xlabelParams (None, dict): modified matplotlib.axes.Axes.set_xlabel
               parameters.
            #. ylabelParams (None, dict): modified matplotlib.axes.Axes.set_ylabel
               parameters.
        """
        return super(StructureFactorConstraint, self).plot(xlabelParams= xlabelParams,
                                                           ylabelParams= ylabelParams,
                                                           **kwargs)

class ReducedStructureFactorConstraint(StructureFactorConstraint):
    """
    The Reduced Structure Factor that we will also note S(Q)
    is exactly the same quantity as the Structure Factor but with
    the slight difference that it is normalized to 0 rather than 1
    and therefore :math:`<S(Q)>=0`.

    The computation of S(Q) is done through a Sine inverse Fourier transform
    of the computed pair distribution function noted as G(r).

    .. math::

        S(Q) = \\frac{1}{Q} \\int_{0}^{\\infty} G(r) sin(Qr) dr

    The only reason why the Reduced Structure Factor is implemented, is because
    many experimental data are treated in this form. And it is just convenient
    not to manipulate the experimental data every time.
    """
    def _get_Sq_from_Gr(self, Gr):
        return np.sum(Gr.reshape((-1,1))*self.Gr2SqMatrix, axis=0)

    def _apply_scale_factor(self, Sq, scaleFactor):
        if scaleFactor != 1:
            Sq = scaleFactor*Sq
        return Sq

    def get_adjusted_scale_factor(self, experimentalData, modelData, dataWeights):
        """ dummy overload that does exactly the same thing
        """
        SF = self.scaleFactor
        # check to update scaleFactor
        if self.adjustScaleFactorFrequency:
            if not self.engine.accepted%self.adjustScaleFactorFrequency:
                SF = self.fit_scale_factor(experimentalData, modelData, dataWeights)
        return SF

    def plot(self, xlabelParams={'xlabel':'$Q(\\AA^{-1})$', 'size':10},
                   ylabelParams={'ylabel':'$S(Q)-1$', 'size':10},
                   **kwargs):
        """
        Alias to ExperimentalConstraint.plot with additional parameters

        :Additional/Adjusted Parameters:
            #. xlabelParams (None, dict): modified matplotlib.axes.Axes.set_xlabel
               parameters.
            #. ylabelParams (None, dict): modified matplotlib.axes.Axes.set_ylabel
               parameters.
        """
        return super(StructureFactorConstraint, self).plot(xlabelParams= xlabelParams,
                                                           ylabelParams= ylabelParams,
                                                           **kwargs)


import numpy as np
from numpy import concatenate as cat
from scipy.sparse import csr_matrix
import scipy.sparse.linalg as spla
from copy import copy
import matplotlib.pyplot as plt
import warnings
from .preprocess import shape, discretization, boundaryCondition

plt.rc('text', usetex=True)
plt.rc('font', family='serif')
# supress the deprecation warning
warnings.filterwarnings("ignore", ".*GUI is implemented.*")

class hdpg1d(object):
    """
    1D HDG solver
    """

    def __init__(self, coeff):
        self.numEle = coeff.numEle
        self.numBasisFuncs = coeff.pOrder + 1
        self.coeff = coeff
        self.mesh = np.linspace(0, 1, self.numEle + 1)
        self.enrichOrder = 1
        self.primalSoln = None
        self.adjointSoln = None
        self.estErrorList = [[], []]
        self.trueErrorList = [[], []]

    def separateSoln(self, soln):
        """Separate gradState (q and u), stateFace from the given soln"""
        gradState, stateFace = np.split(
            soln, [len(soln) - self.numEle + 1])
        return gradState, stateFace

    def plotState(self, counter):
        """Plot solution u with smooth higher oredr quadrature"""
        stateSmooth = np.array([])
        stateNode = np.zeros(self.numEle + 1)
        xSmooth = np.array([])
        gradState, _ = self.separateSoln(self.primalSoln)
        halfLenState = int(len(gradState) / 2)
        state = gradState[halfLenState:2 * halfLenState]
        # quadrature rule
        gorder = 10 * self.numBasisFuncs
        xi, wi = np.polynomial.legendre.leggauss(gorder)
        shp, shpx = shape(xi, self.numBasisFuncs)
        for j in range(1, self.numEle + 1):
            xSmooth = np.hstack((xSmooth, (self.mesh[(j - 1)] + self.mesh[j]) / 2 + (
                self.mesh[j] - self.mesh[j - 1]) / 2 * xi))
            stateSmooth = np.hstack(
                (stateSmooth, shp.T.dot(state[(j - 1) * self.numBasisFuncs:j * self.numBasisFuncs])))
            stateNode[j - 1] = state[(j - 1) * self.numBasisFuncs]
        stateNode[-1] = state[-1]
        plt.figure(1)
        plt.plot(xSmooth, stateSmooth, '-', color='C3')
        plt.plot(self.mesh, stateNode, 'C3.')
        plt.xlabel('$x$', fontsize=17)
        plt.ylabel('$u$', fontsize=17)
        # plt.axis([-0.05, 1.05, 0, 1.3])
        plt.grid()
        plt.pause(5e-1)

    def meshAdapt(self, index):
        """Given the index list, adapt the mesh"""
        inValue = np.zeros(len(index))
        for i in np.arange(len(index)):
            inValue[i] = (self.mesh[index[i]] +
                          self.mesh[index[i] - 1]) / 2
        self.mesh = np.sort(np.insert(self.mesh, 0, inValue))

    def solvePrimal(self):
        """Solve the primal problem"""
        if 'matLocal' in locals():
            # if matLocal exists,
            # only change the mesh instead of initializing again
            matLocal.mesh = self.mesh
        else:
            matLocal = discretization(self.coeff, self.mesh)
        matGroup = matLocal.matGroup()
        A, B, _, C, D, E, F, G, H, L, R = matGroup
        # solve by exploiting the local global separation
        K = -cat((C.T, G), axis=1)\
            .dot(np.linalg.inv(np.bmat([[A, -B], [B.T, D]]))
                 .dot(cat((C, E)))) + H
        sK = csr_matrix(K)
        F_hat = np.array([L]).T - cat((C.T, G), axis=1)\
            .dot(np.linalg.inv(np.bmat([[A, -B], [B.T, D]])))\
            .dot(np.array([cat((R, F))]).T)

        def invRHS(vec):
            """Construct preconditioner"""
            matVec = spla.spsolve(sK, vec)
            return matVec
        n = len(F_hat)
        preconditioner = spla.LinearOperator((n, n), invRHS)
        stateFace = spla.gmres(sK, F_hat, M=preconditioner)[0]
        # stateFace = np.linalg.solve(K, F_hat)
        gradState = np.linalg.inv(np.asarray(np.bmat([[A, -B], [B.T, D]]))).dot(
            cat((R, F)) - cat((C, E)).dot(stateFace))
        self.primalSoln = cat((gradState, stateFace))

    def solveAdjoint(self):
        """Solve the adjoint problem"""
        # solve in the enriched space
        _coeff = copy(self.coeff)
        _coeff.pOrder = _coeff.pOrder + 1
        if 'matAdjoint' in locals():
            matAdjoint.mesh = self.mesh
        else:
            matAdjoint = discretization(_coeff, self.mesh)
        matGroup = matAdjoint.matGroup()
        A, B, _, C, D, E, F, G, H, L, R = matGroup
        # add adjoint LHS conditions
        F = np.zeros(len(F))
        R[-1] = -boundaryCondition('adjoint')[1]
        # assemble global matrix LHS
        LHS = np.bmat([[A, -B, C],
                       [B.T, D, E],
                       [C.T, G, H]])
        sLHS = csr_matrix(LHS)
        RHS = cat((R, F, L))

        # solve in one shoot using GMRES
        def invRHS(vec):
            """Construct preconditioner"""
            matVec = spla.spsolve(sLHS, vec)
            return matVec
        n = len(RHS)
        preconditioner = spla.LinearOperator((n, n), invRHS)
        soln = spla.gmres(sLHS, RHS, M=preconditioner)[0]
        # soln = np.linalg.solve(LHS.T, RHS)
        self.adjointSoln = soln

    def DWResidual(self):
        if 'matResidual' in locals():
            matResidual.mesh = self.mesh
        else:
            matResidual = discretization(
                self.coeff, self.mesh, self.enrichOrder)
        matGroup = matResidual.matGroup()
        A, B, BonQ, C, D, E, F, G, H, L, R = matGroup
        LHS = np.bmat([[A, -B, C],
                       [BonQ, D, E]])
        RHS = cat((R, F))
        residual = np.zeros(self.numEle)
        numEnrich = self.numBasisFuncs + self.enrichOrder
        adjointGradState, adjointStateFace = self.separateSoln(
            self.adjointSoln)
        for i in np.arange(self.numEle):
            primalResidual = (LHS.dot(self.primalSoln) - RHS).A1
            uLength = self.numEle * numEnrich
            stepLength = i * numEnrich
            uDWR = primalResidual[stepLength:stepLength + numEnrich].dot(
                (1 - adjointGradState)[stepLength:stepLength + numEnrich])
            qDWR = primalResidual[uLength + stepLength:uLength +
                                  stepLength + numEnrich]\
                .dot((1 - adjointGradState)[uLength + stepLength:uLength +
                                            stepLength + numEnrich])
            residual[i] = uDWR + qDWR
        # sort residual index
        residualIndex = np.argsort(np.abs(residual))
        # select top \theta% elements with the largest error
        theta = 0.15
        refineIndex = residualIndex[
            int(self.numEle * (1 - theta)):len(residual)] + 1
        return np.abs(np.sum(residual)), refineIndex

    def adaptive(self):
        TOL = self.coeff.TOL
        estError = 10
        nodeCount = 0
        maxCount = self.coeff.MAXIT
        while estError > TOL and nodeCount < maxCount:
            # solve
            self.solvePrimal()
            self.solveAdjoint()
            # plot the solution at certain counter
            if nodeCount in [0, 4, 9, 19, maxCount]:
                plt.clf()
                self.plotState(nodeCount)
            # record error
            self.trueErrorList[0].append(self.numEle)
            self.trueErrorList[1].append(
                self.primalSoln[self.numEle * self.numBasisFuncs - 1])
            estError, index = self.DWResidual()
            self.estErrorList[0].append(self.numEle)
            self.estErrorList[1].append(estError)
            # adapt
            index = index.tolist()
            self.meshAdapt(index)
            self.numEle = self.numEle + len(index)
            nodeCount += 1
            print("Iteration {}. Estimated target function error {:.3e}."
                  .format(nodeCount, estError))
            if nodeCount == maxCount:
                print("Max iteration number is reached "
                      "while the convergence criterion is not satisfied.\n"
                      "Check the problem statement or "
                      "raise the max iteration number, then try again.\n")