import numpy as np
from .cahn_hilliard_simulation import CahnHilliardSimulation
from .microstructure_generator import MicrostructureGenerator
from pymks import DiscreteIndicatorBasis, MKSRegressionModel

__all__ = ['make_delta_microstructures', 'make_elastic_FE_strain_delta',
           'make_elastic_FE_strain_random', 'make_cahn_hilliard',
           'make_microstructure', 'make_checkerboard_microstructure',
           'make_elastic_stress_random']


def make_elastic_FE_strain_delta(elastic_modulus=(100, 150),
                                 poissons_ratio=(0.3, 0.3),
                                 size=(21, 21), macro_strain=0.01):
    """Generate delta microstructures and responses

    Simple interface to generate delta microstructures and their
    strain response fields that can be used for the fit method in the
    `MKSRegressionModel`. The length of `elastic_modulus` and
    `poissons_ratio` indicates the number of phases in the
    microstructure. The following example is or a two phase
    microstructure with dimensions of `(5, 5)`.

    Args:
        elastic_modulus (list, optional): elastic moduli for the phases
        poissons_ratio (list, optional): Poisson's ratios for the phases
        size (tuple, optional): size of the microstructure
        macro_strain (float, optional): Scalar for macroscopic strain applied
        strain_index (int, optional): interger value to return a particular
            strain field. 0 returns exx, 1 returns eyy, etc. To return all
            strain fields set strain_index equal to slice(None).

    Returns:
        tuple containing delta microstructures and their strain fields

    Example

    >>> elastic_modulus = (1., 2.)
    >>> poissons_ratio = (0.3, 0.3)
    >>> X, y = make_elastic_FE_strain_delta(elastic_modulus=elastic_modulus,
    ...                                     poissons_ratio=poissons_ratio,
    ...                                     size=(5, 5))

    `X` is the delta microstructures, and `y` is the
    strain response fields.

    """
    from .elastic_FE_simulation import ElasticFESimulation

    FEsim = ElasticFESimulation(elastic_modulus=elastic_modulus,
                                poissons_ratio=poissons_ratio,
                                macro_strain=macro_strain)

    X = make_delta_microstructures(len(elastic_modulus), size=size)
    FEsim.run(X)
    return X, FEsim.response


def make_delta_microstructures(n_phases=2, size=(21, 21)):
    """Constructs delta microstructures

    Constructs delta microstructures for an arbitrary number of phases
    given the size of the domain.

    Args:
        n_phases (int, optional): number of phases
        size (tuple, optional): dimension of microstructure

    Returns:
        delta microstructures for the system of shape
        (n_samples, n_x, ...)

    Example

    >>> X = np.array([[[[0, 0, 0],
    ...                 [0, 0, 0],
    ...                 [0, 0, 0]],
    ...                [[0, 0, 0],
    ...                 [0, 1, 0],
    ...                 [0, 0, 0]],
    ...                [[0, 0, 0],
    ...                 [0, 0, 0],
    ...                 [0, 0, 0]]],
    ...               [[[1, 1, 1],
    ...                 [1, 1, 1],
    ...                 [1, 1, 1]],
    ...                [[1, 1, 1],
    ...                 [1, 0, 1],
    ...                 [1, 1, 1]],
    ...                [[1, 1, 1],
    ...                 [1, 1, 1],
    ...                 [1, 1, 1]]]])

    >>> assert(np.allclose(X, make_delta_microstructures(2, size=(3, 3, 3))))

    """
    shape = (n_phases, n_phases) + size
    center = tuple((np.array(size) - 1) / 2)
    X = np.zeros(shape=shape, dtype=int)
    X[:] = np.arange(n_phases)[(slice(None), None) + (None,) * len(size)]
    X[(slice(None), slice(None)) + center] = np.arange(n_phases)
    mask = ~np.identity(n_phases, dtype=bool)
    return X[mask]


def make_elastic_FE_strain_random(n_samples=1, elastic_modulus=(100, 150),
                                  poissons_ratio=(0.3, 0.3), size=(21, 21),
                                  macro_strain=0.01):
    """Generate random microstructures and responses

    Simple interface to generate random microstructures and their