"""
physical_models_vec.py

    A module for material strength behavior to be imported into python scripts for
    optimizaton or training emulators.  Adapted from strength_models_add_ptw.py

    Authors:
        DJ Luscher,    djl@lanl.gov
        Peter Trubey,  ptrubey@lanl.gov
        Devin Francom, dfrancom@lanl.gov
        JeeYeon Plohr, jplohr@lanl.gov
        Sky Sjue, sjue@lanl.gov
        Lauren VanDervort, @lvandervort@lanl.gov
        Daniel N Blaschke, dblaschke@lanl.gov
"""

import math
from math import pi

import numpy as np

np.seterr(all="raise")

try:
    from numba import jit

    @jit(nopython=True)
    def erf(x):
        """numba.jit-compatible erf."""
        out = np.empty(x.shape)
        for i, xi in enumerate(x):
            out[i] = math.erf(xi)
        return out

except ImportError:
    from functools import partial

    from scipy.special import erf

    def jit(func=None, forceobj=True, nopython=False):
        """Dummy decorator if numba is unavailable at runtime."""
        return func or partial(jit, forceobj=forceobj, nopython=nopython)

## Error Definitions


class ConstraintError(ValueError):
    pass


class PTWStressError(FloatingPointError):
    pass


## Model Definitions


class BaseModel:
    """
    Base Class for property Models (flow stress, specific heat, melt, density,
    etc.).  Must be instantiated as a child of MaterialModel
    """

    params = []
    consts = []

    def value(self, *args):
        pass

    def update_parameters(self, x):
        self.parent.parameters.update_parameters(x, self.params)

    def __init__(self, parent):
        self.parent = parent


########################
# Specific Heat Models
########################


class Constant_Specific_Heat(BaseModel):
    """
    Constant Specific Heat Model
    """

    consts = ["Cv0"]

    def value(self, *args):
        return self.parent.parameters.Cv0 * np.ones(len(self.parent.state.T))


# class Linear_Specific_Heat(BaseModel):
#    """
#    Linear Specific Heat Model
#    """
#    consts = ['Cv0', 'T0', 'dCdT']

#    def value(self, *args):
#        c0=self.parent.parameters.Cv0
#        t0=self.parent.parameters.T0
#        dcdt=self.parent.parameters.dCdT
#        tnow=self.parent.state.T
#        cnow=c0+(tnow-t0)*dcdt
#        return cnow


class Linear_Specific_Heat(BaseModel):
    """
    Linear Specific Heat Model
    """

    consts = ["c0", "c1"]

    def value(self, *args):
        tnow = self.parent.state.T
        cnow = self.parent.parameters.c0 + self.parent.parameters.c1 * tnow
        return cnow


class Quadratic_Specific_Heat(BaseModel):
    """
    Quadratic Specific Heat Model
    """

    consts = ["c0", "c1", "c2"]

    def value(self, *args):
        tnow = self.parent.state.T
        cnow = (
            self.parent.parameters.c0
            + self.parent.parameters.c1 * tnow
            + self.parent.parameters.c2 * tnow**2
        )
        return cnow


class Cubic_Specific_Heat(BaseModel):
    """
    Cubic Specific Heat Model
    """

    consts = ["c0", "c1", "c2", "c3"]

    def value(self, *args):
        tnow = self.parent.state.T
        cnow = (
            self.parent.parameters.c0
            + self.parent.parameters.c1 * tnow
            + self.parent.parameters.c2 * tnow**2
            + self.parent.parameters.c3 * tnow**3
        )
        return cnow


class Piecewise_Linear_Specific_Heat(BaseModel):
    """
    Piecewise Linear Specific Heat Model
    Cv (T) = c0_0 + c1_0 * T for T<=T_t
    Cv (T) = c0_1 + c1_1 * T for T>T_t
    """

    consts = ["T_t", "c0_0", "c1_0", "c0_1", "c1_1"]

    def value(self, *args):
        tnow = self.parent.state.T
        intercept = np.repeat(self.parent.parameters.c0_0, len(tnow))
        slope = np.repeat(self.parent.parameters.c1_0, len(tnow))
        intercept[np.where(tnow > self.parent.parameters.T_t)] = (
            self.parent.parameters.c0_1
        )
        slope[np.where(tnow > self.parent.parameters.T_t)] = (
            self.parent.parameters.c1_1
        )
        cnow = intercept + slope * tnow
        return cnow


class Piecewise_Quadratic_Specific_Heat(BaseModel):
    """
    Piecewise Quadratic Specific Heat Model
    Cv (T) = c0_0 + c1_0 * T + c2_0 * T**2 for T<=T_t
    Cv (T) = c0_1 + c1_1 * T + c2_1 * T**2 for T>T_t
    """

    consts = ["T_t", "c0_0", "c1_0", "c2_0", "c0_1", "c1_1", "c2_1"]

    def value(self, *args):
        tnow = self.parent.state.T
        pow_0_coeff = np.repeat(self.parent.parameters.c0_0, len(tnow))
        pow_1_coeff = np.repeat(self.parent.parameters.c1_0, len(tnow))
        pow_2_coeff = np.repeat(self.parent.parameters.c2_0, len(tnow))

        pow_0_coeff[np.where(tnow > self.parent.parameters.T_t)] = (
            self.parent.parameters.c0_1
        )
        pow_1_coeff[np.where(tnow > self.parent.parameters.T_t)] = (
            self.parent.parameters.c1_1
        )
        pow_2_coeff[np.where(tnow > self.parent.parameters.T_t)] = (
            self.parent.parameters.c2_1
        )

        cnow = pow_0_coeff + pow_1_coeff * tnow + pow_2_coeff * tnow * tnow
        return cnow


class Piecewise_Cubic_Specific_Heat(BaseModel):
    """
    Cubic Specific Heat Model
    Piecewise Quadratic Specific Heat Model
    Cv (T) = c0_0 + c1_0 * T + c2_0 * T**2  + c3_0 * T**3 for T<=T_t
    Cv (T) = c0_1 + c1_1 * T + c2_1 * T**2  + c3_1 * T**3 for T>T_t
    """

    consts = [
        "T_t",
        "c0_0",
        "c1_0",
        "c2_0",
        "c3_0",
        "c0_1",
        "c1_1",
        "c2_1",
        "c3_1",
    ]

    def value(self, *args):
        tnow = self.parent.state.T
        pow_0_coeff = np.repeat(self.parent.parameters.c0_0, len(tnow))
        pow_1_coeff = np.repeat(self.parent.parameters.c1_0, len(tnow))
        pow_2_coeff = np.repeat(self.parent.parameters.c2_0, len(tnow))
        pow_3_coeff = np.repeat(self.parent.parameters.c3_0, len(tnow))

        pow_0_coeff[np.where(tnow > self.parent.parameters.T_t)] = (
            self.parent.parameters.c0_1
        )
        pow_1_coeff[np.where(tnow > self.parent.parameters.T_t)] = (
            self.parent.parameters.c1_1
        )
        pow_2_coeff[np.where(tnow > self.parent.parameters.T_t)] = (
            self.parent.parameters.c2_1
        )
        pow_3_coeff[np.where(tnow > self.parent.parameters.T_t)] = (
            self.parent.parameters.c3_1
        )

        cnow = (
            pow_0_coeff
            + pow_1_coeff * tnow
            + pow_2_coeff * tnow * tnow
            + pow_3_coeff * tnow * tnow * tnow
        )
        return cnow


########################
# Density Models
########################


class Constant_Density(BaseModel):
    """
    Constant Density Model
    """

    consts = ["rho0"]

    def value(self, *args):
        return self.parent.parameters.rho0 * np.ones(len(self.parent.state.T))


# class Linear_Density(BaseModel):
#    """
#    Linear Density Model
#    """
#    consts = ['rho0', 'T0', 'dRhodT']

#    def value(self, *args):
#        r0=self.parent.parameters.rho0
#        t0=self.parent.parameters.T0
#        drdt=self.parent.parameters.dRhodT
#        tnow=self.parent.state.T
#        rnow=r0+drdt*(tnow-t0)
#        return rnow


class Linear_Density(BaseModel):
    """
    Linear Density Model
    """

    consts = ["r0", "r1"]

    def value(self, *args):
        tnow = self.parent.state.T
        rnow = self.parent.parameters.r0 + self.parent.parameters.r1 * tnow
        return rnow


class Quadratic_Density(BaseModel):
    """
    Quadratic Density Model
    """

    consts = ["r0", "r1", "r2"]

    def value(self, *args):
        tnow = self.parent.state.T
        rnow = (
            self.parent.parameters.r0
            + self.parent.parameters.r1 * tnow
            + self.parent.parameters.r2 * tnow**2
        )
        return rnow


class Cubic_Density(BaseModel):
    """
    Quadratic Density Model
    """

    consts = ["r0", "r1", "r2", "r3"]

    def value(self, *args):
        tnow = self.parent.state.T
        rnow = (
            self.parent.parameters.r0
            + self.parent.parameters.r1 * tnow
            + self.parent.parameters.r2 * tnow**2
            + self.parent.parameters.r3 * tnow**3
        )
        return rnow


########################
# Melt Temperature Models
########################


class Constant_Melt_Temperature(BaseModel):
    """
    Constant Melt Temperature Model
    """

    consts = ["Tmelt0"]

    def value(self, *args):
        return self.parent.parameters.Tmelt0 * np.ones(len(self.parent.state.T))


# class Linear_Melt_Temperature(BaseModel):
#    """
#    Linear Melt Temperature Model
#    """
#    consts=['Tmelt0', 'rho0', 'dTmdRho']#

#    def value(self, *args):
#        tm0=self.parent.parameters.Tmelt0
#        rnow=self.parent.state.rho
#        dtdr=self.parent.parameters.dTmdRho
#        r0=self.parent.parameters.rho0
#        tmeltnow=tm0+dtdr*(rnow-r0)
#        return tmeltnow


class Linear_Melt_Temperature(BaseModel):
    """
    Linear Melt Temperature Model
    """

    consts = ["tm0", "tm1"]

    def value(self, *args):
        rnow = self.parent.state.rho

        tmeltnow = (
            self.parent.parameters.tm0 + self.parent.parameters.tm1 * rnow
        )
        return tmeltnow


class Quadratic_Melt_Temperature(BaseModel):
    """
    Quadratic Melt Temperature Model
    """

    consts = ["tm0", "tm1", "tm2"]

    def value(self, *args):
        rnow = self.parent.state.rho
        tmeltnow = (
            self.parent.parameters.tm0
            + self.parent.parameters.tm1 * rnow
            + self.parent.parameters.tm2 * rnow**2
        )
        return tmeltnow


class Cubic_Melt_Temperature(BaseModel):
    """
    Cubic Melt Temperature Model
    """

    consts = ["tm0", "tm1", "tm2", "tm3"]

    def value(self, *args):
        rnow = self.parent.state.rho
        tmeltnow = (
            self.parent.parameters.tm0
            + self.parent.parameters.tm1 * rnow
            + self.parent.parameters.tm2 * rnow**2
            + self.parent.parameters.tm3 * rnow**3
        )
        # print(tmeltnow)
        return tmeltnow


class BGP_Melt_Temperature(BaseModel):
    consts = ["Tm_0", "rho_m", "gamma_1", "gamma_3", "q3"]

    def value(self, *args):
        mp = self.parent.parameters
        rho = self.parent.state.rho

        melt_temp = (
            mp.Tm_0
            * np.power(rho / mp.rho_m, 1.0 / 3.0)
            * np.exp(
                6
                * mp.gamma_1
                * (np.power(mp.rho_m, -1.0 / 3.0) - np.power(rho, -1.0 / 3.0))
                + 2.0
                * mp.gamma_3
                / mp.q3
                * (np.power(mp.rho_m, -mp.q3) - np.power(rho, -mp.q3))
            )
        )
        return melt_temp


# Shear Modulus Models


class Constant_Shear_Modulus(BaseModel):
    consts = ["G0"]

    def value(self, *args):
        return self.parent.parameters.G0 * np.ones(len(self.parent.state.T))


# class Linear_Shear_Modulus(BaseModel):
#    consts =  ['G0', 'rho0', 'dGdRho' ]#

#    def value(self, *args):
#         g0=self.parent.parameters.G0
#         rho0=self.parent.parameters.rho0
#         dgdr=self.parent.parameters.dGdRho
#         rnow=self.parent.state.rho
#         gnow=g0+dgdr*(rnow-rho0)
#         return gnow


class Linear_Cold_PW_Shear_Modulus(BaseModel):
    consts = ["g0", "g1", "alpha"]

    def value(self, *args):
        mp = self.parent.parameters
        rho = self.parent.state.rho
        temp = self.parent.state.T
        tmelt = self.parent.state.Tmelt
        cold_shear = mp.g0 + mp.g1 * rho
        gnow = cold_shear * (1.0 - mp.alpha * (temp / tmelt))

        gnow[np.where(temp >= tmelt)] = 0.0
        gnow[np.where(gnow < 0)] = 0.0

        return gnow


class Quadratic_Cold_PW_Shear_Modulus(BaseModel):
    consts = ["g0", "g1", "g2", "alpha"]

    def value(self, *args):
        mp = self.parent.parameters
        rho = self.parent.state.rho
        temp = self.parent.state.T
        tmelt = self.parent.state.Tmelt
        cold_shear = mp.g0 + mp.g1 * rho + mp.g2 * rho**2
        gnow = cold_shear * (1.0 - mp.alpha * (temp / tmelt))

        gnow[np.where(temp >= tmelt)] = 0.0
        gnow[np.where(gnow < 0)] = 0.0

        return gnow


class Simple_Shear_Modulus(BaseModel):
    consts = ["G0", "alpha"]

    def value(self, *args):
        mp = self.parent.parameters
        temp = self.parent.state.T
        tmelt = self.parent.state.Tmelt

        return mp.G0 * (1.0 - mp.alpha * (temp / tmelt))


@jit(nopython=True)
def _BGP_PW_Shear_Modulus(
    rho, G0, rho_0, gamma_1, gamma_2, q2, alpha, temp, tmelt
):
    """BPG model provides cold shear, i.e. shear modulus at zero temperature as a function of density.
    PW describes the (linear) temperature dependence of the shear modulus. (Same dependency as
    in Simple_Shear_modulus.)
    With these two models combined, we get the shear modulus as a function of density and temperature;
    see Burakovsky, Greeff, Preston, Phys. Rev. B67 (2003) 094107, DOI:10.1103/PhysRevB.67.094107"""
    cold_shear = (
        G0
        * np.power(rho / rho_0, 4.0 / 3.0)
        * np.exp(
            6.0 * gamma_1 * (1 / np.cbrt(rho_0) - 1 / np.cbrt(rho))
            + 2 * gamma_2 / q2 * (np.power(rho_0, -q2) - np.power(rho, -q2))
        )
    )
    gnow = cold_shear * (1.0 - alpha * (temp / tmelt))

    gnow[np.where(temp >= tmelt)] = 0.0
    gnow[np.where(gnow < 0)] = 0.0

    # if temp >= tmelt: gnow = 0.0
    # if gnow < 0.0:    gnow = 0.0
    return gnow


class BGP_PW_Shear_Modulus(BaseModel):
    """BPG model provides cold shear, i.e. shear modulus at zero temperature as a function of density.
    PW describes the (linear) temperature dependence of the shear modulus. (Same dependency as
    in Simple_Shear_modulus.)
    With these two models combined, we get the shear modulus as a function of density and temperature;
    see Burakovsky, Greeff, Preston, Phys. Rev. B67 (2003) 094107, DOI:10.1103/PhysRevB.67.094107"""

    consts = ["G0", "rho_0", "gamma_1", "gamma_2", "q2", "alpha"]

    def value(self, *args):
        mp = self.parent.parameters
        rho = self.parent.state.rho
        temp = self.parent.state.T
        tmelt = self.parent.state.Tmelt
        gnow = _BGP_PW_Shear_Modulus(
            rho=rho,
            G0=mp.G0,
            rho_0=mp.rho_0,
            gamma_1=mp.gamma_1,
            gamma_2=mp.gamma_2,
            q2=mp.q2,
            alpha=mp.alpha,
            temp=temp,
            tmelt=tmelt,
        )
        return gnow


class Stein_Shear_Modulus(BaseModel):
    # consts = ['G0', 'sgA', 'sgB']
    # assuming constant density and pressure
    # so we only include the temperature dependence
    consts = ["G0", "sgB"]
    eta = 1.0

    def value(self, *args):
        mp = self.parent.parameters
        temp = self.parent.state.T
        tmelt = self.parent.state.Tmelt
        # just putting this here for completeness
        # aterm = a/eta**(1.0/3.0)*pressure
        aterm = 0.0
        bterm = mp.sgB * (temp - 300.0)
        gnow = mp.G0 * (1.0 + aterm - bterm)
        # if temp >= tmelt: gnow = 0.0
        # if gnow < 0.0:    gnow = 0.0
        gnow[np.where(temp >= tmelt)] = 0.0
        gnow[np.where(gnow < 0)] = 0.0
        return gnow


# Yield Stress Models


class Constant_Yield_Stress(BaseModel):
    """
    Constant Yield Stress Model
    """

    consts = ["yield_stress", "chi"]

    def value(self, *args):
        return self.parent.parameters.yield_stress * np.ones(
            len(self.parent.state.T)
        )


# def fast_pow(a, b):
#     """
#     Numpy power is slow, this is faster.  Gets a**b for a and b np arrays.
#     """
#     ## no longer true in numpy 1.24 (=our requirement in toml file) and higher; deprecate / remove?
#     cond = a > 0
#     out = a * 0.0
#     out[cond] = np.exp(b[cond] * np.log(a[cond]))
#     return out
fast_pow = np.power ## do we need this?


def pos(a):
    return np.maximum(0, a)


class JC_Yield_Stress(BaseModel):
    params = ["A", "B", "C", "n", "m"]
    consts = ["Tref", "edot0", "chi"]

    def value(self, edot):
        mp = self.parent.parameters
        eps = self.parent.state.strain
        t = self.parent.state.T
        tmelt = self.parent.state.Tmelt

        # th = np.max([(t - mp.Tref) / (tmelt - mp.Tref), 0.])
        th = pos((t - mp.Tref) / (tmelt - mp.Tref))

        Y = (
            (mp.A + mp.B * np.power(eps, mp.n))
            * (1.0 + mp.C * np.log(edot / mp.edot0))
            * (1.0 - np.power(th, mp.m))
        )
        return Y


@jit(nopython=True)
def _PTW_corefct(
    p,
    kappa,
    s0,
    sInf,
    y0,
    yInf,
    y1,
    y2,
    beta,
    theta,
    lgamma,
    xiDot,
    edot,
    t_hom,
    eps,
    small=1.0e-10,
):
    """jit-compiled subroutine of the PTW_Yield_Stress class"""
    log_xid_ed = np.log(xiDot / edot)
    argErf = kappa * t_hom * (lgamma + log_xid_ed)
    Erfres = erf(argErf)

    saturation1 = s0 - (s0 - sInf) * Erfres
    saturation2 = s0 * np.exp(beta * (-lgamma - log_xid_ed))
    sat_cond = saturation1 > saturation2
    tau_s = np.copy(saturation2)
    tau_s[np.where(sat_cond)] = saturation1[sat_cond]

    ayield = y0 - (y0 - yInf) * Erfres
    byield = y1 * np.exp(-y2 * (lgamma + log_xid_ed))
    cyield = s0 * np.exp(-beta * (lgamma + log_xid_ed))

    y_cond = byield < cyield
    dyield = np.copy(cyield)
    dyield[np.where(y_cond)] = byield[y_cond]

    y_cond2 = ayield > dyield
    tau_y = np.copy(dyield)
    tau_y[np.where(y_cond2)] = ayield[y_cond2]
    scaled_stress = tau_s
    ind = np.where((p > small) * (np.abs(tau_s - tau_y) > small))
    eArg1 = (p * (tau_s - tau_y) / (s0 - tau_y))[ind]
    eArg2 = (
        (eps * p * theta)[ind] / (s0 - tau_y)[ind] / (np.exp(eArg1) - 1.0)
    )  # eArg1 already subsetted by ind
    if np.any((1.0 - (1.0 - np.exp(-eArg1)) * np.exp(-eArg2)) <= 0) or np.any(
        np.isinf(1.0 - (1.0 - np.exp(-eArg1)) * np.exp(-eArg2))
    ):
        print("bad")
    theLog = np.log(1.0 - (1.0 - np.exp(-eArg1)) * np.exp(-eArg2))
    scaled_stress[ind] = tau_s[ind] + (s0[ind] - tau_y[ind]) * theLog / p[ind]
    ind2 = np.where((p <= small) * (tau_s > tau_y))
    scaled_stress[ind2] = +tau_s[ind2] - (tau_s - tau_y)[ind2] * np.exp(
        -eps[ind2] * theta[ind2] / (tau_s - tau_y)[ind2]
    )
    return scaled_stress


@jit(nopython=True)
def PTW_goodparam(s0, sInf, y0, yInf, y1, y2, beta):
    """checks if the given PTW parameter set is valid"""
    return (
        (sInf < s0)
        * (yInf < y0)
        * (y0 < s0)
        * (yInf < sInf)
        * (y1 > s0)
        * (y2 > beta)
    )


class PTW_Yield_Stress(BaseModel):
    """This class implements the PTW flow stress model"""

    params = [
        "theta",
        "p",
        "s0",
        'beta',
        "sInf",
        "kappa",
        "lgamma",
        "y0",
        "yInf",
        "y1",
        "y2",
    ]
    consts = ["rho0", "matomic", "chi"]

    # @profile
    def value(self, edot):
        """
        function used to define PTW flow stress model
        arguments are:
            - edot: scalar, strain rate
            - material: an instance of MaterialModel class
        returns the flow stress at the current material state
        and specified strain rate
        """
        mp = self.parent.parameters
        eps = self.parent.state.strain
        temp = self.parent.state.T
        tmelt = self.parent.state.Tmelt
        shear = self.parent.state.G

        good = PTW_goodparam(
            mp.s0, mp.sInf, mp.y0, mp.yInf, mp.y1, mp.y2, mp.beta
        )
        if np.any(np.logical_not(good)):
            # return np.array([-999.]*len(good))
            raise ConstraintError("PTW bad val")

        # convert to 1/s strain rate since PTW rate is in that unit
        edot = edot * 1.0e6

        t_hom = temp / tmelt
        # this one is commented because it is assumed that
        # the material state computes the temperature dependence of
        # the shear modulus
        # shear = shear * (1.0 - mp.alpha * t_hom)
        # print("ptw shear is "+str(shear))

        afact = (4.0 / 3.0) * pi * mp.rho0 / mp.matomic
        # ainv is 1/a where 4/3 pi a^3 is the atomic volume
        ainv = afact ** (1.0 / 3.0)

        # transverse wave velocity up to units
        xfact = np.sqrt(shear / mp.rho0)
        # PTW characteristic strain rate [ 1/s ]
        xiDot = 0.5 * ainv * xfact * pow(6.022e29, 1.0 / 3.0) * 1.0e4

        # should be flow stress in units of Mbar
        out = (
            _PTW_corefct(
                mp.p,
                mp.kappa,
                mp.s0,
                mp.sInf,
                mp.y0,
                mp.yInf,
                mp.y1,
                mp.y2,
                mp.beta,
                mp.theta,
                mp.lgamma,
                xiDot,
                edot,
                t_hom,
                eps,
                small=1.0e-10,
            )
            * shear
            * 2.0
        )
        out[np.where(np.logical_not(good))] = -999.0
        return out


class Stein_Flow_Stress(BaseModel):
    params = ["y0", "a", "b", "beta", "n", "ymax"]
    consts = ["G0", "epsi", "chi"]

    def value(self, *args):
        mp = self.parent.parameters
        temp = self.parent.state.T
        tmelt = self.parent.state.Tmelt
        shear = self.parent.state.G
        eps = self.parent.state.strain
        fnow = np.power((1.0 + mp.beta * (mp.epsi + eps)), mp.n)

        cond1 = fnow * mp.y0 > mp.ymax
        fnow[cond1] = (mp.ymax / mp.y0)[cond1]
        cond2 = temp > tmelt
        fnow[cond2] = 0.0

        # if fnow*mp.y0 > mp.ymax: fnow = mp.ymax/mp.y0
        # if temp > tmelt: fnow = 0.0
        return mp.y0 * fnow * shear / mp.G0


## Parameters Definition


class ModelParameters:
    params = []
    consts = []
    parent = None

    def update_parameters(self, x):
        if isinstance(x, np.ndarray):
            self.__dict__.update(dict(zip(self.params, x)))
        elif isinstance(x, dict):
            for key in self.params:
                self.__dict__[key] = x[key]
        elif isinstance(x, list):
            assert len(x) == len(self.params), "Incorrect number of parameters!"
            for i, xi in enumerate(self.params):
                self.__dict__[self.params[i]] = xi
        else:
            raise ValueError(f"Type {type(x)} is not supported.")

    def __init__(self, parent):
        self.parent = parent


## State Definition


class MaterialState:
    T = None
    Tmelt = None
    stress = None
    strain = None
    G = None

    def set_state(self, T=300.0, strain=0.0, stress=0.0):
        self.T = T
        self.strain = strain
        self.stress = stress

    def __init__(self, parent, T=300.0, strain=0.0, stress=0.0):
        self.parent = parent
        self.set_state(T, strain, stress)


## Material Model Definition


class MaterialModel:
    def __init__(
        self,
        parameters=ModelParameters,
        initial_state=MaterialState,
        flow_stress_model=Constant_Yield_Stress,
        specific_heat_model=Constant_Specific_Heat,
        shear_modulus_model=Constant_Shear_Modulus,
        melt_model=Constant_Melt_Temperature,
        density_model=Constant_Density,
    ):
        """
        Initialization routine for Material Model.  All of the arguments
        supplied are classes, which are then instantiated within the function.

        The reason for doing this is that then we can pass the MaterialModel
        instance to the physical models so that the model's parent can be
        declared at instantiation.  then the model.value() function can reach
        into the parent class to find whatever it needs.
        """
        self.parameters = parameters(self)
        self.state = initial_state(self)

        self.flow_stress = flow_stress_model(self)
        self.specific_heat = specific_heat_model(self)
        self.shear_modulus = shear_modulus_model(self)
        self.melt_model = melt_model(self)
        self.density = density_model(self)

        params = (
            self.flow_stress.params
            + self.specific_heat.params
            + self.shear_modulus.params
            + self.melt_model.params
            + self.density.params
        )
        consts = set(
            self.flow_stress.consts
            + self.specific_heat.consts
            + self.shear_modulus.consts
            + self.melt_model.consts
            + self.density.consts
        )

        assert len(set(params)) == len(params) , "Some Duplicate Parameters between models"
        assert len(set(params).intersection(set(consts))) == 0 , "Duplicate item in parameters and constants"

        self.parameters.params = params
        self.parameters.consts = consts

    def get_parameter_list(
        self,
    ):
        """
        The list of parameters used in the model.
        This also describes the order of their appearance in the sampling results
        """
        return self.parameters.params

    def get_constants_list(
        self,
    ):
        """
        List of Constants used in the model
        """
        return self.parameters.consts

    def update_state(self, edot, dt):
        chi = self.parameters.chi
        self.state.Cv = self.specific_heat.value()
        self.state.rho = self.density.value()
        # if we are working with microseconds, then this is a reasonable value
        # if we work in seconds, it should be changed to ~1.
        edotcrit = 1.0e-6
        # if edot > edotcrit:
        #  self.state.T += chi * self.state.stress * edot * dt / (self.state.Cv * self.state.rho)
        cond = edot > edotcrit
        # if any(cond):
        self.state.T = (
            self.state.T
            + cond
            * chi
            * self.state.stress
            * edot
            * dt
            / (self.state.Cv * self.state.rho)
        )
        self.state.strain = self.state.strain + edot * dt

        self.state.Tmelt = self.melt_model.value()
        self.state.G = self.shear_modulus.value()
        self.state.stress = self.flow_stress.value(edot)

    def update_parameters(self, x):
        self.parameters.update_parameters(x)

    def initialize(self, parameters, constants):
        """
        Initialize the model at a given set of parameters, constants
        """
        ## if user assumed one or more parameters constant, they would be in the constants var instead;
        ## check for this first:
        if not isinstance(self.parameters.params, list):
            self.parameters.params = list(self.parameters.params)
        if not isinstance(self.parameters.consts, set):
            self.parameters.consts = set(self.parameters.consts)
        user_constants = set(self.parameters.params).difference(parameters)
        for usercnst in user_constants:
            self.parameters.consts |= {usercnst}
        try:
            self.parameters.__dict__.update(
                {
                    key: parameters[key]
                    for key in set(self.parameters.params).difference(
                        user_constants
                    )
                },
            )
            self.parameters.__dict__ |= {
                key: constants[key] for key in user_constants
            }
        except KeyError:
            print(
                "{} missing from list of supplied parameters".format(
                    set(self.parameters.params).difference(
                        set(parameters.keys())
                    )
                )
            )
            raise
        try:
            self.parameters.__dict__.update(
                {key: constants[key] for key in self.parameters.consts},
            )
        except KeyError:
            print(
                "{} missing from list of supplied constants".format(
                    set(self.parameters.consts).difference(
                        set(constants.keys())
                    )
                )
            )
            raise

    def initialize_state(self, T=300.0, stress=0.0, strain=0.0):
        self.state.set_state(T, stress, strain)

    def set_history_variables(self, emax, edot, Nhist):
        self.emax = emax
        self.edot = edot
        self.Nhist = Nhist

    def get_history_variables(self):
        return [self.emax, self.edot, self.Nhist]

    def compute_state_history(self, strain_history):
        strains = strain_history["strains"]
        times = strain_history["times"]
        strain_rate = strain_history["strain_rate"]
        # Nhist = len(strains)
        # nrep = len(self.parameters.kappa)
        nrep, Nhist = strains.shape  # nexp * nhist array

        results = np.empty((Nhist, 6, nrep))

        state = self.state
        self.update_state(strain_rate[:, 0], 0.0)

        # import pdb
        # pdb.set_trace()

        results[0] = np.array([
            times[:, 0],
            state.strain,
            state.stress,
            state.T,
            state.G,
            state.rho,
        ])  # np.repeat(state.rho,nrep)])

        for i in range(1, Nhist):
            self.update_state(
                strain_rate[:, i - 1], times[:, i] - times[:, i - 1]
            )
            # self.update_state(strain_rate.T[i-1], times.T[i] - times.T[i-1])
            # results[i] = [times[i], state.strain, state.stress, state.T, state.G, state.rho]
            results[i] = np.array([
                times[:, i],
                state.strain,
                state.stress,
                state.T,
                state.G,
                state.rho,
            ])  # np.repeat(state.rho, nrep)])

        return results


## function to generate strain history to calculate along
##  Should probably make this a method of MaterialModel class


def generate_strain_history(emax, edot, Nhist):
    tmax = emax / edot
    strains = np.linspace(0.0, emax, Nhist)
    nrep = len(edot)
    times = np.empty((nrep, Nhist))
    # for i in range(nrep):
    #    times[i,:] = np.linspace(0., tmax[i], Nhist)
    times = np.linspace(0.0, tmax, Nhist)
    strain_rate = np.empty((nrep, Nhist - 1))
    # for i in range(nrep):
    #    strain_rate[i, :] = np.diff(strains) / np.diff(times[i, :])

    strain_diffs = np.diff(strains)

    strain_rate = strain_diffs[:, np.newaxis] / np.diff(times, axis=0)
    return dict((
        ["times", times.T],
        ["strains", strains],
        ["strain_rate", strain_rate.T],
    ))


def generate_strain_history_new(emax, edot, nhist):
    tmax = emax / edot
    strains = np.linspace(0, emax, nhist)  # nhist * nexp
    times = np.linspace(0, tmax, nhist)  # nhist * nexp
    rates = np.diff(strains, axis=0) / np.diff(
        times, axis=0
    )  # (nhist - 1) * nexp
    return {"times": times.T, "strains": strains.T, "strain_rate": rates.T}
